Update app.py

app.py

@@ -11,7 +11,7 @@ import subprocess
 subprocess.run('pip install --upgrade transformers', shell=True)
 subprocess.run('pip install accelerate', shell=True)
 
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, AutoModelForVisionText2Text
 
 # Model and processor initialization with trust_remote_code=True
 processor = AutoProcessor.from_pretrained(
@@ -19,7 +19,7 @@ processor = AutoProcessor.from_pretrained(
     trust_remote_code=True
 )
 
-model = AutoModelForCausalLM.from_pretrained(
+model = AutoModelForVisionText2Text.from_pretrained(
     "Qwen/QVQ-72B-Preview",
     trust_remote_code=True,
     device_map="auto"
@@ -35,25 +35,28 @@ footer = """
 # Vision model function
 @spaces.GPU()
 def process_image(image, text_input=None):
-    # Convert image to PIL format
-    image = Image.fromarray(image).convert("RGB")
-    
-    # Prepare inputs
-    if text_input:
-        inputs = processor(text=text_input, images=image, return_tensors="pt")
-    else:
-        inputs = processor(images=image, return_tensors="pt")
-    
-    # Move inputs to the same device as the model
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    
-    # Generate output
-    outputs = model.generate(**inputs, max_new_tokens=1000)
-    
-    # Decode response
-    response = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-    
-    return response
+    try:
+        # Convert image to PIL format
+        image = Image.fromarray(image).convert("RGB")
+        
+        # Prepare inputs
+        if text_input:
+            inputs = processor(text=text_input, images=image, return_tensors="pt")
+        else:
+            inputs = processor(images=image, return_tensors="pt")
+        
+        # Move inputs to the same device as the model
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        # Generate output
+        outputs = model.generate(**inputs, max_new_tokens=1000)
+        
+        # Decode response
+        response = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        
+        return response
+    except Exception as e:
+        return f"Error processing image: {str(e)}"
 
 # CSS styling
 css = """