Update llm.py

llm.py

@@ -1,6 +1,7 @@
 import torch
 import asyncio
 from functools import partial
+import time
 from transformers import AutoProcessor, AutoModelForImageTextToText #, BitsAndBytesConfig
 
 # Quantization config
@@ -13,7 +14,7 @@ processor = AutoProcessor.from_pretrained("datalab-to/chandra-ocr-2")
 model = AutoModelForImageTextToText.from_pretrained(
     "datalab-to/chandra-ocr-2",
     # quantization_config=quant_config,
-    device_map="cuda"
+    device_map="auto"
 )
 
 print("CUDA available:", torch.cuda.is_available())
@@ -22,11 +23,12 @@ print("Model device:", model.device)
 if torch.cuda.is_available():
     print("GPU name:", torch.cuda.get_device_name(0))
     print("Memory allocated:", torch.cuda.memory_allocated() / 1e9, "GB")
+
     
-# ✅ SYNC function (runs in thread)
 def execute_llm(model, processor, image, prompt: str):
     print("execute llm")
-
+    print(prompt)
+    a=time.time()
     # ✅ Use passed prompt (FIXED)
     if not prompt:
         prompt = """
@@ -43,7 +45,7 @@ def execute_llm(model, processor, image, prompt: str):
         }
     ]
 
-    # ✅ Inference mode (faster + less memory)
+    
     with torch.inference_mode():
 
         inputs = processor.apply_chat_template(
@@ -68,16 +70,14 @@ def execute_llm(model, processor, image, prompt: str):
             outputs[0][inputs["input_ids"].shape[-1]:],
             skip_special_tokens=True
         )
-
+    print('total time taken',time.time()-a)
     print(result)
     return result
 
 
-# ✅ ASYNC wrapper (non-blocking FastAPI)
 async def call_llm(image, prompt: str = ""):
     print("call llm")
 
-    loop = asyncio.get_event_loop()
 
     result=execute_llm(model,processor,image,prompt)