Update app.py

app.py

@@ -1,21 +1,24 @@
+# app.py
 import gradio as gr
 from transformers import AutoModelForVision2Seq, AutoProcessor
 import torch
 from PIL import Image
 import os
 
-# Load Qwen-VL model and processor
+# Load Qwen-VL model and processor (trust custom code)
 model_id = "Qwen/Qwen-VL-Chat"
-processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForVision2Seq.from_pretrained(model_id, trust_remote_code=True)
+model = model.to("cpu")
 
 # Inference function
 def ocr_with_qwen(image):
+    # Fallback to test.png if no image uploaded
     if image is None:
         image = Image.open("test.png")
 
     prompt = "<|im_start|>system\nYou are a helpful assistant. Extract all text from the image and output only the text.<|im_end|>\n<|im_start|>user\n"
-    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
+    inputs = processor(images=image, text=prompt, return_tensors="pt").to("cpu")
     outputs = model.generate(**inputs, max_new_tokens=512)
     result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
     return result.strip()