Back to my working implementation

labelizer/__init__.py

@@ -1,28 +1,34 @@
 import torch
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import Florence2ForConditionalGeneration, Florence2Processor
 
 MODEL_ID = "ducviet00/Florence-2-large-hf"
 
 # Global variables for lazy loading
 _model = None
 _processor = None
+_device = None
+_torch_dtype = None
 
 
 def _load_model():
     """Load model and processor lazily"""
-    global _model, _processor
+    global _model, _processor, _device, _torch_dtype
 
     if _model is None:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-        print(f"Loading model {MODEL_ID} on {device}...")
-        _model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
-        _processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+        _device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        _torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+        print(f"Loading model {MODEL_ID} on {_device} with dtype {_torch_dtype}...")
+        _model = Florence2ForConditionalGeneration.from_pretrained(
+            MODEL_ID, torch_dtype=_torch_dtype, trust_remote_code=True
+        ).to(_device)  # type: ignore
+        _processor = Florence2Processor.from_pretrained(
+            MODEL_ID, trust_remote_code=True
+        )
         print("Model loaded successfully!")
 
-    return _model, _processor
+    return _model, _processor, _device, _torch_dtype
 
 
 def get_task_response(task_prompt: str, image: Image.Image, text_input=None):
@@ -35,7 +41,7 @@ def get_task_response(task_prompt: str, image: Image.Image, text_input=None):
 
     """
     # Lazy load model only when needed
-    model, processor = _load_model()
+    model, processor, device, torch_dtype = _load_model()
     if text_input is None:
         prompt = task_prompt
     else:
@@ -47,13 +53,11 @@ def get_task_response(task_prompt: str, image: Image.Image, text_input=None):
 
     if processor is None:
         raise ValueError("processor is None")
-
-    # Process inputs using the correct API
-    inputs = processor(text=prompt, images=image, return_tensors="pt")
-
-    # Move inputs to device if model is on CUDA
-    device = next(model.parameters()).device
-    inputs = {k: v.to(device) for k, v in inputs.items()}
+    inputs = processor(
+        text=prompt,
+        images=image,
+        return_tensors="pt",  # type: ignore
+    ).to(device, torch_dtype)
 
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],