Update handler.py

handler.py

@@ -8,6 +8,9 @@ from PIL import Image
 import requests
 import math
 import ast
+from io import BytesIO
+from urllib.parse import urlparse
+
 
 # ===== Kullanılacak HF model id =====
 MODEL_ID = os.getenv("HF_MODEL_ID", "PULSE-ECG/PULSE-7B")
@@ -260,22 +263,107 @@ class EndpointHandler:
         self.use_im_start_end = getattr(self.model.config, "mm_use_im_start_end", False)
 
     # ---- yardımcılar ----
-    def _load_image(self, img_field: str) -> Optional[Image.Image]:
-        """URL / base64 / path -> PIL.Image"""
-        if not img_field:
+    def _load_image(self, image_input: str) -> Optional[Image.Image]:
+        """
+        URL / base64 / yerel path -> PIL.Image
+        - URL için: içerik tipi/uzantı kontrolü, UA header, redirect, boyut sınırı
+        - base64 için: data URL prefix temizliği + padding düzeltme
+        - path için: doğrudan aç
+        """
+        if not image_input:
             return None
+    
+        # 25MB üstünü reddet (isteğe göre ayarlanabilir)
+        MAX_IMAGE_BYTES = int(os.getenv("MAX_IMAGE_BYTES", "26214400"))
+        ALLOWED_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp"}
+    
+        def _is_valid_image_format(url: str) -> bool:
+            try:
+                ext = os.path.splitext(urlparse(url).path.lower())[1]
+                # Uzantı yoksa da kabul et; Content-Type ile kontrol ederiz
+                return (ext in ALLOWED_EXTS) or (ext == "")
+            except Exception:
+                return True
+    
         try:
-            if img_field.startswith("data:image"):
-                _, b64 = img_field.split(",", 1)
-                return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
-            if img_field.startswith(("http://", "https://")):
-                r = requests.get(img_field, timeout=20)
-                r.raise_for_status()
-                return Image.open(io.BytesIO(r.content)).convert("RGB")
-            return Image.open(img_field).convert("RGB")
+            # ---- URL input ----
+            if isinstance(image_input, str) and image_input.startswith(("http://", "https://")):
+                if not _is_valid_image_format(image_input):
+                    print("[warn] Invalid image extension in URL")
+                    return None
+    
+                headers = {"User-Agent": "Mozilla/5.0"}
+                resp = requests.get(image_input, timeout=20, headers=headers, allow_redirects=True, stream=True)
+                resp.raise_for_status()
+    
+                # Content-Type kontrolü (varsa)
+                ctype = resp.headers.get("Content-Type", "").lower()
+                if ctype and not ctype.startswith("image/"):
+                    print(f"[warn] Non-image content-type: {ctype}")
+                    return None
+    
+                # Boyut kontrolü (Content-Length varsa)
+                clen = resp.headers.get("Content-Length")
+                if clen is not None:
+                    try:
+                        if int(clen) > MAX_IMAGE_BYTES:
+                            print(f"[warn] Image too large: {clen} bytes")
+                            return None
+                    except Exception:
+                        pass
+    
+                # Stream’den kontrollü oku
+                data = resp.content
+                if len(data) > MAX_IMAGE_BYTES:
+                    print(f"[warn] Image too large (actual): {len(data)} bytes")
+                    return None
+    
+                img = Image.open(BytesIO(data)).convert("RGB")
+                print(f"[info] URL image loaded: size={img.size}")
+                return img
+    
+            # ---- Base64 input (data URL dahil) ----
+            if isinstance(image_input, str):
+                b64 = image_input.strip()
+    
+                # data URL prefix varsa ayıkla
+                if b64.startswith("data:image"):
+                    # ör: data:image/png;base64,AAAA...
+                    if "base64," in b64:
+                        b64 = b64.split("base64,", 1)[1]
+                    else:
+                        # bazen ;base64 sonrası newline vb olabilir
+                        b64 = b64.split(",", 1)[-1]
+    
+                # boşluk/newline temizliği + padding düzeltme
+                b64 = b64.replace("\n", "").replace("\r", "").replace(" ", "")
+                missing = (4 - len(b64) % 4) % 4
+                b64 += "=" * missing
+    
+                try:
+                    data = base64.b64decode(b64, validate=False)
+                    if len(data) > MAX_IMAGE_BYTES:
+                        print(f"[warn] Base64 image too large: {len(data)} bytes")
+                        return None
+                    img = Image.open(BytesIO(data)).convert("RGB")
+                    print(f"[info] Base64 image loaded: size={img.size}")
+                    return img
+                except Exception as e:
+                    print(f"[warn] Base64 decode/open failed: {e}")
+                    # Devam edip path olarak deneyeceğiz
+    
+            # ---- Yerel path ----
+            if isinstance(image_input, str) and os.path.exists(image_input):
+                img = Image.open(image_input).convert("RGB")
+                print(f"[info] Local image loaded: size={img.size}")
+                return img
+    
         except Exception as e:
             print(f"[warn] image load failed: {e}")
             return None
+    
+        return None
+
 
     def _build_prompt(self, user_text: str, conv_mode: str) -> str:
         """Model worker tarzında prompt oluştur"""