Update models/gpt4o.py

models/gpt4o.py

@@ -4,16 +4,13 @@ import os
 import json
 import base64
 from PIL import Image
-from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
-import numpy as np
-
 import openai
+import torch
 
 model = "gpt-4o"
 
-# Load Pix2Struct model + processor (vision-language OCR)
-processor = Pix2StructProcessor.from_pretrained("google/pix2struct-textcaps-base")
-pix2struct_model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-textcaps-base")
+pix2struct_model = None
+processor = None
 
 
 def load_prompt(prompt_file="prompts/prompt.txt"):
@@ -52,23 +49,29 @@ def encode_image_base64(image: Image.Image):
 
 
 def extract_all_text_pix2struct(image: Image.Image):
-    inputs = processor(images=image, return_tensors="pt")
+    global pix2struct_model, processor
+
+    # Lazy-load the Pix2Struct model
+    if processor is None or pix2struct_model is None:
+        from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration
+        processor = Pix2StructProcessor.from_pretrained("google/pix2struct-textcaps-base")
+        pix2struct_model = Pix2StructForConditionalGeneration.from_pretrained(
+            "google/pix2struct-textcaps-base"
+        ).to("cuda" if torch.cuda.is_available() else "cpu")
+
+    inputs = processor(images=image, return_tensors="pt").to(pix2struct_model.device)
     predictions = pix2struct_model.generate(**inputs, max_new_tokens=512)
     output_text = processor.decode(predictions[0], skip_special_tokens=True)
     return output_text.strip()
 
 
-# Optional: assign best-matching label from full extracted text using proximity (simplified version)
 def assign_event_gateway_names_from_ocr(image: Image.Image, json_data, ocr_text):
     if not ocr_text:
         return json_data
 
-    # You could use NLP matching or regex in complex cases
-    words = ocr_text.split()
-
     def guess_name_fallback(obj):
         if not obj.get("name") or obj["name"].strip() == "":
-            obj["name"] = "(label unknown)"  # fallback if matching logic isn't yet implemented
+            obj["name"] = "(label unknown)"
 
     for evt in json_data.get("events", []):
         guess_name_fallback(evt)
@@ -83,9 +86,7 @@ def run_model(image: Image.Image, api_key: str = None):
     prompt_text = load_prompt()
     encoded_image = encode_image_base64(image)
 
-    # Pull from environment if not passed explicitly
     api_key = api_key or os.getenv("OPENAI_API_KEY")
-
     if not api_key:
         return {"json": None, "raw": "⚠️ API key is missing. Please set it as a secret in your Space or upload it as a file."}
 
@@ -107,8 +108,8 @@ def run_model(image: Image.Image, api_key: str = None):
     output_text = response.choices[0].message.content.strip()
     parsed_json = try_extract_json(output_text)
 
-    # Vision-language OCR assist step (Pix2Struct)
+    # Use Pix2Struct OCR enrichment
     full_ocr_text = extract_all_text_pix2struct(image)
     parsed_json = assign_event_gateway_names_from_ocr(image, parsed_json, full_ocr_text)
 
-    return {"json": parsed_json, "raw": output_text}
+    return {"json": parsed_json, "raw": output_text}