Update handler.py

handler.py

@@ -1,76 +1,28 @@
 import torch
 from PIL import Image
 from transformers import AutoModel, AutoTokenizer
-import requests
-from io import BytesIO
 import base64
-import ssl
-import urllib3
-import os
-
-# Check if CUDA is available
-print("CUDA Available:", torch.cuda.is_available())
-
-if torch.cuda.is_available():
-    device_name = torch.cuda.get_device_name(torch.cuda.current_device())
-    print(f"CUDA Device: {device_name}")
-    print(f"Memory Allocated: {torch.cuda.memory_allocated()} bytes")
-    print(f"Memory Cached: {torch.cuda.memory_reserved()} bytes")
-    print(f"PyTorch Version: {torch.__version__}")
-    print(f"CUDA Version (PyTorch uses): {torch.version.cuda}")
-else:
-    print("CUDA is not available!")
-
-urllib3.disable_warnings()
-ssl._create_default_https_context = ssl._create_unverified_context
+from io import BytesIO
+import requests
 
 class EndpointHandler:
     def __init__(self, model_dir=None):
         self.load_model()
 
     def load_model(self):
-        model_name = "openbmb/MiniCPM-V-2_6"
-        hf_token = os.getenv("HF_AUTH_TOKEN")
-
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            model_name, trust_remote_code=True, use_auth_token=hf_token
-        )
-        self.model = AutoModel.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            attn_implementation="sdpa",
-            torch_dtype=torch.float16,
-            use_auth_token=hf_token
-        ).eval().cuda()
+        model_name = "openbmb/MiniCPM-V-2_6-int4"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16)
+        self.model.eval()
 
-    def load_image(self, image_input):
-        if image_input.startswith("http"):
-            try:
-                resp = requests.get(image_input, verify=False)
-                image = Image.open(BytesIO(resp.content)).convert("RGB")
-                return image
-            except Exception as e:
-                raise ValueError(f"Failed to fetch image from URL: {e}")
-        
-        elif image_input.startswith("data:image"):
-            try:
-                image = Image.open(BytesIO(base64.b64decode(image_input.split(",")[1]))).convert("RGB")
-                return image
-            except Exception as e:
-                raise ValueError(f"Invalid base64 image format: {e}")
-        
-        else:
-            try:
-                image = Image.open(image_input).convert("RGB")
-                return image
-            except Exception as e:
-                raise ValueError(f"Failed to open image from file path: {e}")
+    def load_image(self, image_bytes):
+        try:
+            image = Image.open(BytesIO(image_bytes)).convert("RGB")
+            return image
+        except Exception as e:
+            raise ValueError(f"Failed to open image from bytes: {e}")
 
     def predict(self, request):
-        # Unwrap Hugging Face format
-        if "inputs" in request:
-            request = request["inputs"]
-
         image_input = request.get("image")
         question = request.get("question", "What is in the image?")
         stream = request.get("stream", False)
@@ -80,28 +32,30 @@ class EndpointHandler:
 
         try:
             image = self.load_image(image_input)
-            msgs = [{"role": "user", "content": f"<image>\n{question}"}]
-
-            try:
-                if stream:
-                    generated_text = ""
-                    for chunk in self.model.chat(
-                        image=None, msgs=msgs, tokenizer=self.tokenizer,
-                        sampling=True, stream=True
-                    ):
-                        generated_text += chunk
-                    return {"output": generated_text}
-                else:
-                    output = self.model.chat(image=None, msgs=msgs, tokenizer=self.tokenizer)
-                    return {"output": output}
-            except Exception as e:
-                return {"error": f"Inference failed: {e}"}
-
-        except ValueError as e:
-            return {"error": f"Image processing error: {e}"}
+            msgs = [{"role": "user", "content": [image, question]}]
+
+            if stream:
+                generated_text = ""
+                res = self.model.chat(
+                    image=None,
+                    msgs=msgs,
+                    tokenizer=self.tokenizer,
+                    sampling=True,
+                    temperature=0.7,
+                    stream=True
+                )
+                for new_text in res:
+                    generated_text += new_text
+                return {"output": generated_text}
+            else:
+                output = self.model.chat(image=None, msgs=msgs, tokenizer=self.tokenizer)
+                return {"output": output}
+
+        except Exception as e:
+            return {"error": str(e)}
 
     def __call__(self, data):
         return self.predict(data)
 
-# Hugging Face looks for a callable handler
-handler = EndpointHandler()
+# Example usage
+handler = EndpointHandler()