Create app.py

app.py

@@ -0,0 +1,130 @@
+import streamlit as st
+import os
+import torch
+from byaldi import RAGMultiModalModel
+from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+# Check for CUDA availability
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Caching the model loading
+@st.cache_resource
+def load_rag_model():
+    return RAGMultiModalModel.from_pretrained("vidore/colpali")
+
+@st.cache_resource
+def load_qwen_model():
+    return Qwen2VLForConditionalGeneration.from_pretrained(
+        "Qwen/Qwen2-VL-2B-Instruct",
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16
+    ).to(device).eval()
+
+@st.cache_resource
+def load_processor():
+    return AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", trust_remote_code=True)
+
+# Load models
+RAG = load_rag_model()
+model = load_qwen_model()
+processor = load_processor()
+
+st.title("Multimodal RAG App")
+
+st.warning("⚠️ Disclaimer: This app is currently running on CPU, which may result in slow processing times. For optimal performance, download and run the app locally on a machine with GPU support.")
+
+# Add download link
+st.markdown("[📥 Download the app code](https://huggingface.co/spaces/clayton07/colpali-qwen2-ocr/blob/main/app.py)")
+
+# Initialize session state for tracking if index is created
+if 'index_created' not in st.session_state:
+    st.session_state.index_created = False
+
+# File uploader
+image_source = st.radio("Choose image source:", ("Upload an image", "Use example image"))
+
+if image_source == "Upload an image":
+    uploaded_file = st.file_uploader("Choose an image file", type=["png", "jpg", "jpeg"])
+else:
+    # Use a pre-defined example image
+    example_image_path = "hindi-qp.jpg"
+    uploaded_file = example_image_path
+
+if uploaded_file is not None:
+    # If using the example image, no need to save it
+    if image_source == "Upload an image":
+        with open("temp_image.png", "wb") as f:
+            f.write(uploaded_file.getvalue())
+        image_path = "temp_image.png"
+    else:
+        image_path = uploaded_file
+
+    if not st.session_state.index_created:
+        # Initialize the index for the first image
+        RAG.index(
+            input_path=image_path,
+            index_name="temp_index",
+            store_collection_with_index=False,
+            overwrite=True
+        )
+        st.session_state.index_created = True
+    else:
+        # Add to the existing index for subsequent images
+        RAG.add_to_index(
+            input_item=image_path,
+            store_collection_with_index=False
+        )
+
+    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+
+    # Text query input
+    text_query = st.text_input("Enter your query about the image:")
+
+    if text_query:
+        # Perform RAG search
+        results = RAG.search(text_query, k=2)
+
+        # Process with Qwen2VL model
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "image": image_path,
+                    },
+                    {"type": "text", "text": text_query},
+                ],
+            }
+        ]
+        text = processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        image_inputs, video_inputs = process_vision_info(messages)
+        inputs = processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+        )
+        inputs = inputs.to(device)
+        generated_ids = model.generate(**inputs, max_new_tokens=100)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+
+        # Display results
+        st.subheader("Results:")
+        st.write(output_text[0])
+
+    # Clean up temporary file
+    if image_source == "Upload an image":
+        os.remove("temp_image.png")
+else:
+    st.write("Please upload an image to get started.")