Update app.py

app.py

@@ -1,28 +1,36 @@
+import os
 import torch
-from PIL import Image
 import gradio as gr
+from PIL import Image
 from transformers import AutoProcessor, AutoModelForVision2Seq
 
-processor = AutoProcessor.from_pretrained("scb10x/typhoon-ocr-3b", use_auth_token=False)
+hf_token = os.getenv("HUGGINGFACE_TOKEN")
+
+token_args = {}
+if hf_token:
+    token_args = {"use_auth_token": hf_token}
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+processor = AutoProcessor.from_pretrained("scb10x/typhoon-ocr-3b", **token_args)
 model = AutoModelForVision2Seq.from_pretrained(
     "scb10x/typhoon-ocr-3b",
     torch_dtype=torch.float16,
     device_map="auto",
-    use_auth_token=False
+    **token_args
 )
 
 def ocr_infer(image):
+    if image is None:
+        return "❌ Error: No image provided"
     try:
-        if image is None:
-            return "❌ Error: No image provided"
-        
         image = image.convert("RGB")
         inputs = processor(images=image, return_tensors="pt")
-        
+
         if inputs is None or "pixel_values" not in inputs:
             return "❌ Error: Invalid processor output"
-        
-        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+        inputs = {k: v.to(device) for k, v in inputs.items()}
         generated_ids = model.generate(**inputs, max_new_tokens=256)
         result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         return result
@@ -30,11 +38,5 @@ def ocr_infer(image):
     except Exception as e:
         return f"❌ Error during inference: {e}"
 
-iface = gr.Interface(
-    fn=ocr_infer,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Typhoon OCR 3B"
-)
-
+iface = gr.Interface(fn=ocr_infer, inputs=gr.Image(type="pil"), outputs="text", title="Typhoon OCR 3B")
 iface.launch()