Update app.py

app.py

@@ -1,18 +1,21 @@
-import re
-import base64
-import io
 import torch
 import gradio as gr
 from PIL import Image
-from transformers import MllamaForConditionalGeneration, AutoProcessor
+from transformers import AutoProcessor, AutoModel
 
 # Load the model and processor
-model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-model = MllamaForConditionalGeneration.from_pretrained(
+model_id = "OpenGVLab/InternVL2_5-78B"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Initialize the model and processor
+model = AutoModel.from_pretrained(
     model_id,
     torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
+    low_cpu_mem_usage=True,
+    use_flash_attn=True,
+    trust_remote_code=True
+).eval().to(device)
+
 processor = AutoProcessor.from_pretrained(model_id)
 
 def generate_model_response(image_file, user_query):
@@ -24,49 +27,38 @@ def generate_model_response(image_file, user_query):
     - user_query: The user's question about the image.
 
     Returns:
-    - str: The generated response from the model, formatted as HTML.
+    - str: The generated response from the model.
     """
     try:
         # Load and prepare the image
         raw_image = Image.open(image_file).convert("RGB")
         
-        # Prepare input for the model using the processor
-        conversation = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "url": "<|image|>"},  # Placeholder for image
-                    {"type": "text", "text": user_query}
-                ]
-            }
-        ]
-
-        # Apply chat template to prepare inputs for the model
-        inputs = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+        # Prepare inputs for the model using the processor
+        inputs = processor(
+            text=user_query,
+            images=raw_image,
+            return_tensors="pt"
+        ).to(device)
         
-        # Process the image and text inputs together
-        inputs = processor(inputs, raw_image, return_tensors="pt").to(model.device)
-
         # Generate response from the model
-        outputs = model.generate(**inputs)
-        
-        # Decode and format the response
-        generated_text = processor.decode(outputs[0], skip_special_tokens=True)
+        outputs = model.generate(**inputs, max_new_tokens=50)
         
-        return generated_text
+        # Decode and return the response
+        response_text = processor.decode(outputs[0], skip_special_tokens=True)
+        return response_text
     
     except Exception as e:
         print(f"Error in generating response: {e}")
-        return f"<p>An error occurred: {str(e)}</p>"
+        return f"An error occurred: {str(e)}"
 
 # Gradio Interface
 iface = gr.Interface(
     fn=generate_model_response,
     inputs=[
         gr.Image(type="file", label="Upload Image"),
-        gr.Textbox(label="Enter your question", placeholder="How many calories are in this food?")
+        gr.Textbox(label="Enter your question", placeholder="What do you want to know about this image?")
     ],
-    outputs=gr.HTML(label="Response from Model"),
+    outputs="text",
 )
 
 iface.launch(share=True)