Update handler.py

handler.py

@@ -1,30 +1,69 @@
-# handler.py (örnek iskelet)
-import base64, io, os
-from typing import Any, Dict, List
+# /repository/handler.py
+import os, io, base64
+from typing import Any, Dict, List, Optional
 import torch
 from PIL import Image
-from transformers import AutoTokenizer, AutoProcessor, AutoModelForVision2Seq  # model tipinize göre
 
-HF_MODEL_ID = os.getenv("HF_MODEL_ID", "PULSE-ECG/PULSE-7B")  # ağırlıkların olduğu repo id
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DT = torch.bfloat16 if torch.cuda.is_available() else torch.float32  # bfloat16 GPU varsa
+# ---- LLaVA demodaki modüller ----
+from llava.model.builder import load_pretrained_model
+from llava.mm_utils import tokenizer_image_token, process_images
+from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from llava.conversation import conv_templates, SeparatorStyle
+from llava.utils import disable_torch_init
+from llava.model.builder import get_model_name_from_path
+
+# Ortam değişkenleri (modeli nereden alacağımız)
+# 1) Yerel klasörden yüklemek istersen HF_MODEL_LOCAL_DIR kullan
+# 2) HF Hub repo id ile yüklemek istersen HF_MODEL_ID kullan
+HF_MODEL_LOCAL_DIR = os.getenv("HF_MODEL_LOCAL_DIR", "").strip()
+HF_MODEL_ID        = os.getenv("HF_MODEL_ID", "").strip()  # ör: "your-org/your-llava-model"
+
+DEFAULT_CONV_MODE  = os.getenv("LLAVA_CONV_MODE", "llava_v2")  # demo: llava_v2
+MAX_NEW_TOKENS_DEF = int(os.getenv("MAX_NEW_TOKENS", "4096"))
+
+# Flash-Attention yoksa SDPA güvenli yoldur
+os.environ.setdefault("ATTN_IMPLEMENTATION", "sdpa")
 
 class EndpointHandler:
     def __init__(self, path: str = "") -> None:
-        # path: /repository (bu repo klasörü)
-        # NOT: Ağırlıkları bu repodan değil, HF Hub’dan alıyoruz
-        self.tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, use_fast=True, trust_remote_code=True)
-        self.processor = AutoProcessor.from_pretrained(HF_MODEL_ID, trust_remote_code=True)
-        self.model = AutoModelForVision2Seq.from_pretrained(
-            HF_MODEL_ID,
-            torch_dtype=DT,
-            device_map="auto",              # GPU varsa otomatik yerleşim
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-            # attn_implementation="sdpa",   # flash-attn yoksa güvenlisi SDPA
+        """
+        path: /repository  (endpoint bu klasörü model_dir olarak geçer)
+        """
+        disable_torch_init()
+
+        # Model yolunu belirle
+        if HF_MODEL_LOCAL_DIR:
+            model_path = HF_MODEL_LOCAL_DIR
+        elif HF_MODEL_ID:
+            model_path = HF_MODEL_ID
+        else:
+            # Eğer ağırlık/konfig bu repository içindeyse path= "/repository"
+            model_path = path
+
+        # model adı (LLaVA utils)
+        self.model_name = get_model_name_from_path(model_path)
+
+        # LLaVA yükleme — demo ile aynı giriş:
+        # Dönüş: tokenizer, model, image_processor, context_len
+        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
+            model_path=model_path,
+            model_base=None,               # LoRA vb. yoksa None
+            model_name=self.model_name,
+            torch_dtype="auto",            # ortam GPU'ya göre seçsin
+            attn_implementation=os.getenv("ATTN_IMPLEMENTATION", "sdpa"),
+            device_map="auto"
         )
 
-    def _load_image(self, img_field: str) -> Image.Image:
+        # Görüntü başlangıç/bitiş tokenları (model sürümüne göre aktif)
+        self.use_im_start_end = getattr(self.model.config, "mm_use_im_start_end", False)
+        self.image_token = DEFAULT_IMAGE_TOKEN
+        self.im_start = DEFAULT_IM_START_TOKEN
+        self.im_end = DEFAULT_IM_END_TOKEN
+
+    # ---- Yardımcılar ----
+    def _load_image(self, img_field: str) -> Optional[Image.Image]:
+        if not img_field:
+            return None
         if img_field.startswith("data:image"):
             head, b64 = img_field.split(",", 1)
             return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
@@ -34,26 +73,77 @@ class EndpointHandler:
             r.raise_for_status()
             return Image.open(io.BytesIO(r.content)).convert("RGB")
         else:
+            # container içinden dosya okunacaksa
             return Image.open(img_field).convert("RGB")
 
+    def _build_prompt(self, user_text: str, conv_mode: str) -> str:
+        # Demo: conv_templates ile diyalog kur
+        conv = conv_templates[conv_mode].copy()
+        if self.use_im_start_end:
+            # <im_start> <image> <im_end> + kullanıcı metni
+            content = f"{self.im_start}{self.image_token}{self.im_end}\n{user_text}"
+        else:
+            content = f"{self.image_token}\n{user_text}"
+        conv.append_message(conv.roles[0], content)
+        conv.append_message(conv.roles[1], None)  # assistant boş, model dolduracak
+        return conv.get_prompt()
+
+    # ---- Inference giriş noktası ----
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Beklenen giriş (demo ile uyumlu):
+        {
+          "inputs": { "query": "...", "image": "<url|dataurl|path>" },
+          "parameters": { "max_new_tokens": 256, "temperature": 0.0, "top_p": 1.0, ... },
+          "conv_mode": "llava_v2"  # opsiyonel; yoksa varsayılanı kullanırız
+        }
+        """
         inputs = data.get("inputs") or {}
         params = data.get("parameters") or {}
-        query = inputs.get("query", "")
-        img_f = inputs.get("image", "")
-        image = self._load_image(img_f) if img_f else None
+        conv_mode = data.get("conv_mode") or DEFAULT_CONV_MODE
+
+        query_text = inputs.get("query", "")
+        image_f    = inputs.get("image", "")
+        pil_img    = self._load_image(image_f)
+
+        # 1) Prompt (conversation şablonu)
+        prompt = self._build_prompt(query_text, conv_mode)
+
+        # 2) Görsel tensörü (demo: process_images)
+        image_tensors = None
+        if pil_img is not None:
+            image_tensors = process_images([pil_img], self.image_processor, self.model.config)
 
-        # Model türüne göre preprocessing (örnek akış)
-        model_inputs = self.processor(images=image, text=query, return_tensors="pt").to(self.model.device)
+        # 3) Tokenize (görüntü tokenını metne göm)
+        input_ids = tokenizer_image_token(
+            prompt,
+            self.tokenizer,
+            IMAGE_TOKEN_INDEX,
+            return_tensors="pt"
+        )
+
+        # 4) Cihaza taşı
+        input_ids = input_ids.to(self.model.device, non_blocking=True)
+        if image_tensors is not None:
+            image_tensors = image_tensors.to(self.model.device, dtype=self.model.dtype, non_blocking=True)
+
+        # 5) Generate (demo parametreleri)
         gen_kwargs = {
-            "max_new_tokens": int(params.get("max_new_tokens", 4096)),
+            "max_new_tokens": int(params.get("max_new_tokens", MAX_NEW_TOKENS_DEF)),
             "temperature": float(params.get("temperature", 0.0)),
-            "do_sample": bool(params.get("do_sample", params.get("temperature", 0.0) > 0)),
             "top_p": float(params.get("top_p", 1.0)),
             "repetition_penalty": float(params.get("repetition_penalty", 1.0)),
+            "do_sample": bool(params.get("do_sample", float(params.get("temperature", 0.0)) > 0)),
+            "use_cache": bool(params.get("use_cache", True)),
         }
 
         with torch.no_grad():
-            out_ids = self.model.generate(**model_inputs, **gen_kwargs)
-        text = self.tokenizer.decode(out_ids[0], skip_special_tokens=True)
-        return [{"generated_text": text}]
+            output_ids = self.model.generate(
+                input_ids,
+                images=image_tensors,
+                **gen_kwargs
+            )
+
+        # 6) Decode (assistant’ın cevabı)
+        outputs = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
+        return [{"generated_text": outputs}]