Update app.py

app.py

@@ -179,8 +179,9 @@ def handle_file(file, prompt, extra_prompt, max_new_tokens, progress=gr.Progress
 
 # --- 3. Inference Function ---
 @spaces.GPU
-def run_inference(img: Image.Image, prompt: str = "", max_new_tokens: int = 512) -> str:
+def run_inference(img, prompt="", max_new_tokens=512):
     model.to("cuda")
+
     if img.mode != "RGB":
         img = img.convert("RGB")
     prompt_text = prompt.strip()
@@ -199,9 +200,9 @@ def run_inference(img: Image.Image, prompt: str = "", max_new_tokens: int = 512)
 
     inputs = processor(
         text=[text_prompt], images=[img], return_tensors="pt", padding=True
-    ).to(device)
+    ).to("cuda")  # Sửa ở đây
 
-    with torch.inference_mode(), torch.cuda.amp.autocast(enabled=(device.type == 'cuda')):
+    with torch.inference_mode(), torch.cuda.amp.autocast():
         gen = model.generate(
             **inputs,
             max_new_tokens=max_new_tokens,
@@ -215,7 +216,6 @@ def run_inference(img: Image.Image, prompt: str = "", max_new_tokens: int = 512)
     )[0].strip()
 
     return result
-
 # --- 9. Prompt templates & JSON export ---
 prompt_templates = {
     "Electrolux": """Extract all structured information from the delivery order document image.