Update handler.py

handler.py

@@ -1,61 +1,46 @@
 import torch
 from transformers import AutoProcessor, LlavaForConditionalGeneration
 from PIL import Image
-import requests
-from io import BytesIO
 import base64
+from io import BytesIO
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # This loading logic is robust and correct...
-        base_model_id = "llava-hf/llava-1.5-7b-hf"
-        
-        print("Loading processor...")
-        self.processor = AutoProcessor.from_pretrained(base_model_id)
+        # The 'path' is now a self-contained directory with the complete, merged model.
+        # No internet access is needed here.
         
-        print("Loading base model...")
+        print("Loading model and processor from local path...")
+        self.processor = AutoProcessor.from_pretrained(path, trust_remote_code=True)
         self.model = LlavaForConditionalGeneration.from_pretrained(
-            base_model_id,
+            path,
             load_in_4bit=True,
             torch_dtype=torch.float16,
-            device_map="auto"
+            device_map="auto",
+            trust_remote_code=True
         )
-        
-        print(f"Loading LoRA adapters from repository path: {path}...")
-        # This assumes your repo at `path` contains the LoRA adapter files
-        self.model.load_adapter(path)
-        print("✅ Model and adapters loaded successfully.")
+        print("✅ Model loaded successfully.")
 
     def __call__(self, data: dict) -> dict:
-        # --- Simplified and Corrected Inference Logic ---
-        payload = data.pop("inputs", data)
-            
-        prompt_text = payload.pop("prompt", "What can you do?")
-        image_b64 = payload.pop("image_b64", None)
-        max_new_tokens = payload.pop("max_new_tokens", 200)
-
-        image = None
-        if image_b64:
-            try:
-                image_bytes = base64.b64decode(image_b64)
-                image = Image.open(BytesIO(image_bytes)).convert("RGB")
-            except Exception as e:
-                return {"error": f"Failed to decode base64 image: {e}"}
-
-        # This is the key change: a simple, clear separation of logic
-        if image is not None:
-            # --- Case 1: Multimodal (Image + Text) Request ---
-            prompt = f"USER: <image>\n{prompt_text} ASSISTANT:"
-            inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(self.model.device)
-            output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
-        else:
-            # --- Case 2: Text-Only Request ---
-            prompt = f"USER: {prompt_text} ASSISTANT:"
-            inputs = self.processor(text=prompt, return_tensors="pt").to(self.model.device)
+        prompt_text = data.pop("prompt", "Describe the image in detail.")
+        image_b64 = data.pop("image_b64", None)
+        max_new_tokens = data.pop("max_new_tokens", 200)
+
+        if not image_b64:
+            return {"error": "No image provided. Please use the 'image_b64' key."}
+
+        try:
+            image_bytes = base64.b64decode(image_b64)
+            image = Image.open(BytesIO(image_bytes))
+        except Exception as e:
+            return {"error": f"Failed to decode or open base64 image: {e}"}
+
+        prompt = f"USER: <image>\n{prompt_text} ASSISTANT:"
+        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to("cuda")
+
+        with torch.no_grad():
             output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
 
-        # Decode the output and extract the assistant's response
         full_response = self.processor.decode(output[0], skip_special_tokens=True)
         assistant_response = full_response.split("ASSISTANT:")[-1].strip()
-        
+
         return {"generated_text": assistant_response}