Update handler.py

handler.py

@@ -8,29 +8,21 @@ import ssl
 import urllib3
 import os
 
-# Disable SSL warnings and errors for custom models
+# Disable SSL warnings and errors for dev/debugging
 urllib3.disable_warnings()
 ssl._create_default_https_context = ssl._create_unverified_context
 
 class EndpointHandler:
     def __init__(self, model_dir=None):
-        self.model = None
-        self.tokenizer = None
         self.load_model()
 
     def load_model(self):
         model_name = "openbmb/MiniCPM-V-2_6"
-        
-        # Ensure that the Hugging Face token is set in environment variables or passed as a parameter
-        hf_token = os.getenv("HF_AUTH_TOKEN")  # or use your token directly here
+        hf_token = os.getenv("HF_AUTH_TOKEN")  # Set this as a secret in Hugging Face Space/Endpoint
 
-        # Load the tokenizer and model with the token for authentication
         self.tokenizer = AutoTokenizer.from_pretrained(
-            model_name, 
-            trust_remote_code=True, 
-            use_auth_token=hf_token
+            model_name, trust_remote_code=True, use_auth_token=hf_token
         )
-
         self.model = AutoModel.from_pretrained(
             model_name,
             trust_remote_code=True,
@@ -41,11 +33,11 @@ class EndpointHandler:
 
     def predict(self, request):
         """
-        Expected request format:
+        Expected input:
         {
-            "image": "<url or base64 string>",
-            "question": "What is shown in the image?",
-            "stream": false (optional)
+            "image": "<image URL or base64>",
+            "question": "What is this?",
+            "stream": false
         }
         """
         image_input = request.get("image")
@@ -53,52 +45,32 @@ class EndpointHandler:
         stream = request.get("stream", False)
 
         if not image_input:
-            return {"error": "Image input is required."}
+            return {"error": "Missing image."}
 
         try:
+            # Load image from URL or base64
             if image_input.startswith("http"):
-                # Load image from URL
-                response = requests.get(image_input, verify=False)
-                image = Image.open(BytesIO(response.content)).convert("RGB")
+                resp = requests.get(image_input, verify=False)
+                image = Image.open(BytesIO(resp.content)).convert("RGB")
             else:
-                # Load image from base64 string
                 image = Image.open(BytesIO(base64.b64decode(image_input))).convert("RGB")
         except Exception as e:
-            return {"error": f"Failed to load image: {e}"}
+            return {"error": f"Invalid image format or URL: {e}"}
 
-        # Prepare message for the model
-        msgs = [{"role": "user", "content": [image, question]}]
+        # Prepare message with <image> placeholder
+        msgs = [{"role": "user", "content": f"<image>\n{question}"}]
 
         try:
             if stream:
-                # If streaming is enabled, collect the output incrementally
                 generated_text = ""
-                for new_text in self.model.chat(
-                    image=None,
-                    msgs=msgs,
-                    tokenizer=self.tokenizer,
-                    sampling=True,
-                    stream=True
+                for chunk in self.model.chat(
+                    image=image, msgs=msgs, tokenizer=self.tokenizer,
+                    sampling=True, stream=True
                 ):
-                    generated_text += new_text
+                    generated_text += chunk
                 return {"output": generated_text}
             else:
-                # If streaming is not enabled, get the complete output
-                output = self.model.chat(
-                    image=None,
-                    msgs=msgs,
-                    tokenizer=self.tokenizer
-                )
+                output = self.model.chat(image=image, msgs=msgs, tokenizer=self.tokenizer)
                 return {"output": output}
         except Exception as e:
             return {"error": f"Inference failed: {e}"}
-
-
-# Test block (optional, remove in production)
-if __name__ == "__main__":
-    handler = EndpointHandler()
-    result = handler.predict({
-        "image": "https://upload.wikimedia.org/wikipedia/commons/9/9e/Ours_brun_parcanimalierpyrenees_1.jpg",
-        "question": "What animal is this?"
-    })
-    print(result)