Update handler.py

handler.py

@@ -6,13 +6,16 @@ from io import BytesIO
 import base64
 import ssl
 import urllib3
+
+# Disable SSL warnings and errors for custom models
 urllib3.disable_warnings()
 ssl._create_default_https_context = ssl._create_unverified_context
 
 class EndpointHandler:
-    def __init__(self):
+    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"
@@ -43,17 +46,21 @@ class EndpointHandler:
 
         try:
             if image_input.startswith("http"):
+                # Load image from URL
                 response = requests.get(image_input, verify=False)
                 image = Image.open(BytesIO(response.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}"}
 
+        # Prepare message for the model
         msgs = [{"role": "user", "content": [image, question]}]
 
         try:
             if stream:
+                # If streaming is enabled, collect the output incrementally
                 generated_text = ""
                 for new_text in self.model.chat(
                     image=None,
@@ -65,6 +72,7 @@ class EndpointHandler:
                     generated_text += new_text
                 return {"output": generated_text}
             else:
+                # If streaming is not enabled, get the complete output
                 output = self.model.chat(
                     image=None,
                     msgs=msgs,
@@ -77,8 +85,7 @@ class EndpointHandler:
 
 # Test block (optional, remove in production)
 if __name__ == "__main__":
-    handler = ModelHandler()
-    handler.load_model()
+    handler = EndpointHandler()
     result = handler.predict({
         "image": "https://upload.wikimedia.org/wikipedia/commons/9/9e/Ours_brun_parcanimalierpyrenees_1.jpg",
         "question": "What animal is this?"