Update app.py

app.py

@@ -94,19 +94,33 @@ def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_t
 @spaces.GPU
 def process_vision_query(image, text_input):
     prompt = f"<|user|>\n<|image_1|>\n{text_input}<|end|>\n<|assistant|>\n"
-    image = Image.fromarray(image).convert("RGB")
-    inputs = vision_processor(prompt, image, return_tensors="pt").to(device)
     
-    with torch.no_grad():
-        generate_ids = vision_model.generate(
-            **inputs, 
-            max_new_tokens=1000, 
-            eos_token_id=vision_processor.tokenizer.eos_token_id
-        )
+    # Check if image is already a PIL Image
+    if isinstance(image, Image.Image):
+        pil_image = image
+    elif isinstance(image, np.ndarray):
+        pil_image = Image.fromarray(image).convert("RGB")
+    else:
+        raise ValueError("Unsupported image type. Expected PIL Image or numpy array.")
+    
+    inputs = vision_processor(prompt, pil_image, return_tensors="pt").to(device)
     
-    generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
-    response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    return response
+    try:
+        with torch.no_grad():
+            generate_ids = vision_model.generate(
+                **inputs, 
+                max_new_tokens=1000, 
+                eos_token_id=vision_processor.tokenizer.eos_token_id
+            )
+        
+        generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
+        response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        return response
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            return "Error: GPU out of memory. Try processing a smaller image or freeing up GPU resources."
+        else:
+            raise e
 
 @spaces.GPU
 def generate_speech(prompt, description):