Update handler.py

handler.py

@@ -9,6 +9,10 @@ import requests
 # ===== Kullanılacak HF model id =====
 MODEL_ID = os.getenv("HF_MODEL_ID", "PULSE-ECG/PULSE-7B")
 
+# Flash Attention için environment
+os.environ.setdefault("FLASH_ATTENTION", "1")
+os.environ.setdefault("ATTN_IMPLEMENTATION", "flash_attention_2")
+
 # ===== LLaVA kaynak kodunu runtime'da getir (pip yok) =====
 LLAVA_GIT_URL = os.getenv("LLAVA_GIT_URL", "https://github.com/haotian-liu/LLaVA.git")
 LLAVA_GIT_REF = os.getenv("LLAVA_GIT_REF", "v1.2.2.post1")  # kanıtlı, stabil
@@ -28,35 +32,26 @@ _ensure_llava()
 
 # ---- LLaVA parçaları (demo akışı) ----
 from llava.model.builder import load_pretrained_model
-from llava.mm_utils import tokenizer_image_token, process_images
+from llava.mm_utils import tokenizer_image_token, process_images, get_model_name_from_path
 from llava.constants import (
     IMAGE_TOKEN_INDEX,
     DEFAULT_IMAGE_TOKEN,
     DEFAULT_IM_START_TOKEN,
     DEFAULT_IM_END_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    IMAGE_PLACEHOLDER,
+    IGNORE_INDEX,
 )
 from llava.conversation import conv_templates
 from llava.utils import disable_torch_init
 
-# Eksik fonksiyonu kendimiz tanımlıyoruz
-def get_model_name_from_path(model_path):
-
-    model_path = model_path.strip("/")
-
-    model_paths = model_path.split("/")
-
-    if model_paths[-1].startswith('checkpoint-'):
-
-        return model_paths[-2] + "_" + model_paths[-1]
-
-    else:
-
-        return model_paths[-1]
+# Eksik fonksiyonu kaldır - artık mm_utils'ten import ediyoruz
+# def get_model_name_from_path() artık gerekli değil
 
 # Varsayılanlar
 DEFAULT_CONV_MODE  = os.getenv("LLAVA_CONV_MODE", "llava_llama_2")
 MAX_NEW_TOKENS_DEF = int(os.getenv("MAX_NEW_TOKENS", "256"))
-os.environ.setdefault("ATTN_IMPLEMENTATION", os.getenv("ATTN_IMPLEMENTATION", "sdpa"))
+# ATTN_IMPLEMENTATION artık otomatik seçiliyor, bu satırı kaldırıyoruz
 
 class EndpointHandler:
     """
@@ -82,22 +77,28 @@ class EndpointHandler:
 
         self.model_name = get_model_name_from_path(model_path)
 
+        # Attention implementation otomatik seç
+        try:
+            import flash_attn
+            attn_impl = "flash_attention_2"
+        except ImportError:
+            attn_impl = "sdpa"
+        
         # PULSE, LLaVA tabanlı olduğundan LLaVA loader ile yüklenir
         self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
             model_path=model_path,
             model_base=None,
             model_name=self.model_name,
             torch_dtype="auto",
-            attn_implementation=os.getenv("ATTN_IMPLEMENTATION", "sdpa"),
+            attn_implementation=attn_impl,
             device_map="auto",
         )
         self.model.eval()
 
         # Görsel token işaretleri (LLaVA config)
         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
+        # Constants'tan direkt kullan
+        # self.image_token, self.im_start, self.im_end artık gerekli değil
 
     # ---- yardımcılar ----
     def _load_image(self, img_field: str) -> Optional[Image.Image]:
@@ -122,10 +123,13 @@ class EndpointHandler:
         if conv_mode not in conv_templates:
             conv_mode = DEFAULT_CONV_MODE
         conv = conv_templates[conv_mode].copy()
+        
+        # Image token'ları doğru yerleştir
         if self.use_im_start_end:
-            content = f"{self.im_start}{self.image_token}{self.im_end}\n{user_text}"
+            content = f"{DEFAULT_IM_START_TOKEN}{DEFAULT_IMAGE_TOKEN}{DEFAULT_IM_END_TOKEN}\n{user_text}"
         else:
-            content = f"{self.image_token}\n{user_text}"
+            content = f"{DEFAULT_IMAGE_TOKEN}\n{user_text}"
+            
         conv.append_message(conv.roles[0], content)
         conv.append_message(conv.roles[1], None)
         return conv.get_prompt()
@@ -148,13 +152,25 @@ class EndpointHandler:
         if image_f:
             pil = self._load_image(image_f)
             if pil is not None:
-                image_tensors = process_images([pil], self.image_processor, self.model.config)
-                image_tensors = image_tensors.to(self.model.device, dtype=self.model.dtype, non_blocking=True)
+                try:
+                    # LLaVA'nın gelişmiş process_images fonksiyonunu kullan
+                    # Bu fonksiyon anyres, pad gibi farklı aspect ratio modlarını destekler
+                    image_tensors = process_images([pil], self.image_processor, self.model.config)
+                    if image_tensors is not None and len(image_tensors) > 0:
+                        image_tensors = image_tensors.to(self.model.device, dtype=torch.float16, non_blocking=True)
+                except Exception as e:
+                    print(f"[warn] image processing failed: {e}")
+                    image_tensors = None
 
         # 3) tokenize (image token'ı gömülü)
         input_ids = tokenizer_image_token(
             prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
-        ).to(self.model.device, non_blocking=True)
+        ).unsqueeze(0).to(self.model.device, non_blocking=True)  # unsqueeze ekledik
+        
+        # Input uzunluk kontrolü
+        if input_ids.shape[-1] > self.context_len - 100:
+            # Prompt'u kısalt
+            input_ids = input_ids[:, -(self.context_len - 200):]
 
         # 4) güvenli max_new_tokens
         requested = int(params.get("max_new_tokens", MAX_NEW_TOKENS_DEF))
@@ -162,16 +178,29 @@ class EndpointHandler:
         max_new_tokens = max(1, min(requested, avail))
 
         gen_kwargs = {
+            "input_ids": input_ids,
+            "images": image_tensors,
             "max_new_tokens": max_new_tokens,
             "temperature": float(params.get("temperature", 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)),
+            "pad_token_id": self.tokenizer.eos_token_id,
         }
 
-        with torch.inference_mode():
-            output_ids = self.model.generate(input_ids, images=image_tensors, **gen_kwargs)
-
-        text = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
+        try:
+            with torch.inference_mode():
+                output_ids = self.model.generate(**gen_kwargs)
+            
+            # Output'u input'tan ayır
+            if output_ids.shape[-1] > input_ids.shape[-1]:
+                response_ids = output_ids[:, input_ids.shape[-1]:]
+                text = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True)[0].strip()
+            else:
+                text = "Error: No response generated"
+                
+        except Exception as e:
+            print(f"Generation error: {e}")
+            text = f"Error during generation: {str(e)}"
         return [{"generated_text": text}]