Update handler.py

handler.py

@@ -1,11 +1,11 @@
 # -*- coding: utf-8 -*-
-# handler.py — PULSE-7B / LLaVA robust endpoint (fixed: 'inputs' & NoneType.new_ones)
-# - Uses PULSE fork of LLaVA (AIMedLab/PULSE:dev)
-# - Safe image loading + processor normalization
-# - Attention-mask creation + "mask injection" fallback
-# - Fix for duplicate 'inputs' kwarg during fallback
-# - Small forward() patch to drop unknown kwargs
-# - FIXED: NoneType.new_ones error in mask injection
+# handler.py — PULSE-7B / LLaVA robust endpoint (minimal & stable)
+# - PULSE fork (AIMedLab/PULSE:dev) üzerinden LLaVA yükleme
+# - Güvenli image loader + processor normalizasyonu
+# - ANYRES->PAD fallback
+# - Forward patch: cache_position/input_positions sessizce at
+# - KRİTİK FIX: generate çağrısına hem `inputs` hem de `input_ids` ver (NoneType.new_ones biter)
+# - attention_mask gönderme (LLaVA kendi içinde hallediyor)
 
 import os, io, sys, subprocess, base64
 from typing import Any, Dict, List, Optional, Tuple
@@ -134,7 +134,7 @@ except Exception:
         if len(chunks) > 0 and len(chunks[0]) > 0 and chunks[0][0] == tokenizer.bos_token_id:
             offset = 1
             ids.append(chunks[0][0])
-        for x in insert_sep(chunks, [image_token_index]*(offset+1)):
+        for x in insert_sep(chunks, [IMAGE_TOKEN_INDEX]*(offset+1)):
             ids.extend(x[offset:])
         if return_tensors == 'pt':
             return torch.tensor(ids, dtype=torch.long)
@@ -157,7 +157,6 @@ from llava.constants import (
 )
 from llava.conversation import conv_templates
 from llava.utils import disable_torch_init
-
 from transformers import AutoProcessor, AutoImageProcessor, CLIPImageProcessor
 
 DEFAULT_CONV_MODE  = os.getenv("LLAVA_CONV_MODE", "llava_v1")
@@ -174,7 +173,7 @@ class EndpointHandler:
         else:
             model_path = MODEL_ID
 
-        self.model_name = self._get_model_name_from_path(model_path)
+        self.model_name = get_model_name_from_path(model_path)
 
         try:
             import flash_attn  # noqa
@@ -234,7 +233,7 @@ class EndpointHandler:
                 print("[info] image_processor loaded via AutoProcessor(model_path)")
             except Exception as e:
                 print(f"[warn] AutoProcessor başarısız: {e}")
-                vt_id = self._resolve_vision_tower_id(self.model.config, model_path)
+                vt_id = self._resolve_vision_tower_id(self.model.config)
                 print(f"[hotfix] trying to load image_processor from vision_tower: {vt_id}")
                 try:
                     self.image_processor = AutoImageProcessor.from_pretrained(vt_id, trust_remote_code=True)
@@ -264,131 +263,17 @@ class EndpointHandler:
         self.use_im_start_end = getattr(self.model.config, "mm_use_im_start_end", False)
         self.is_multimodal = ('llava' in self.model_name.lower()) or ('pulse' in self.model_name.lower())
 
-    # ---- FIXED: mask injection helper ----
-    def _generate_with_injected_mask(self, input_ids, images, image_sizes, attention_mask, base_kwargs):
-        """
-        Inject attention_mask inside prepare_inputs_for_generation so HF generate uses it,
-        while avoiding duplicate kwargs like 'inputs' or 'attention_mask'.
-        FIXED: Better handling of None values and tensor validation.
-        """
-        orig_prepare = getattr(self.model, "prepare_inputs_for_generation", None)
-        if orig_prepare is None:
-            print("[error] Model has no prepare_inputs_for_generation method")
-            raise RuntimeError("Model doesn't support mask injection fallback")
-
-        def patched_prepare(input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
-            try:
-                # Call original prepare method
-                model_inputs = orig_prepare(input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs)
-                
-                # Validate model_inputs is not None and is a dict
-                if model_inputs is None:
-                    print("[error] prepare_inputs_for_generation returned None")
-                    model_inputs = {}
-                elif not isinstance(model_inputs, dict):
-                    print(f"[error] prepare_inputs_for_generation returned non-dict: {type(model_inputs)}")
-                    model_inputs = {}
-                
-                # Only inject attention_mask if it's not already present and we have a valid mask
-                if model_inputs.get("attention_mask", None) is None and attention_mask is not None:
-                    # Validate attention_mask
-                    if isinstance(attention_mask, torch.Tensor) and attention_mask.numel() > 0:
-                        model_inputs["attention_mask"] = attention_mask
-                        print(f"[debug] Injected attention_mask with shape: {attention_mask.shape}")
-                    else:
-                        print("[warn] Invalid attention_mask, skipping injection")
-                
-                # Ensure input_ids is present
-                if "input_ids" not in model_inputs and input_ids is not None:
-                    model_inputs["input_ids"] = input_ids
-                
-                return model_inputs
-                
-            except Exception as e:
-                print(f"[error] Error in patched_prepare: {e}")
-                # Return minimal valid dict to avoid None errors
-                return {"input_ids": input_ids}
-
-        # Apply the patch
-        self.model.prepare_inputs_for_generation = patched_prepare
-        
-        try:
-            # IMPORTANT: Remove 'attention_mask' and 'inputs' from kwargs to avoid conflicts
-            patched_kwargs = {k: v for k, v in base_kwargs.items() if k not in ("attention_mask", "inputs")}
-            
-            # Add images and image_sizes if they exist
-            if images is not None:
-                patched_kwargs["images"] = images
-            if image_sizes is not None:
-                patched_kwargs["image_sizes"] = image_sizes
-            
-            # Validate input_ids before generation
-            if input_ids is None or not isinstance(input_ids, torch.Tensor) or input_ids.numel() == 0:
-                raise ValueError("Invalid input_ids for generation")
-            
-            print(f"[debug] Starting generation with input_ids shape: {input_ids.shape}")
-            with torch.inference_mode():
-                output = self.model.generate(inputs=input_ids, **patched_kwargs)
-            
-            return output
-            
-        except Exception as e:
-            print(f"[error] Generation failed in mask injection: {e}")
-            raise e
-        finally:
-            # Always restore original method
-            self.model.prepare_inputs_for_generation = orig_prepare
-
-    # ---- ADDED: Simplified fallback without mask injection ----
-    def _generate_without_mask(self, input_ids, images, image_sizes, base_kwargs):
-        """
-        Fallback generation without attention_mask for models that don't support it well.
-        """
-        try:
-            # Remove problematic arguments
-            clean_kwargs = {k: v for k, v in base_kwargs.items() 
-                           if k not in ("attention_mask", "inputs")}
-            
-            # Add multimodal inputs if present
-            if images is not None:
-                clean_kwargs["images"] = images
-            if image_sizes is not None:
-                clean_kwargs["image_sizes"] = image_sizes
-            
-            # Force use_cache=False for stability
-            clean_kwargs["use_cache"] = False
-            
-            # Ensure we have basic required parameters
-            clean_kwargs.setdefault("pad_token_id", self.tokenizer.pad_token_id)
-            clean_kwargs.setdefault("eos_token_id", self.tokenizer.eos_token_id)
-            
-            print(f"[debug] Fallback generation without mask, kwargs: {list(clean_kwargs.keys())}")
-            
-            with torch.inference_mode():
-                output = self.model.generate(inputs=input_ids, **clean_kwargs)
-            
-            return output
-            
-        except Exception as e:
-            print(f"[error] Fallback generation failed: {e}")
-            raise e
-
     # ------------- helpers -------------
-    def _get_model_name_from_path(self, model_path: str) -> str:
-        p = model_path.strip("/").split("/")
-        return (p[-2] + "_" + p[-1]) if p[-1].startswith("checkpoint-") else p[-1]
-
-    def _resolve_vision_tower_id(self, config: Any, model_path: str) -> str:
+    def _resolve_vision_tower_id(self, config: Any) -> str:
         for key in ("mm_vision_tower", "vision_tower", "mm_vision_tower_name", "image_tower", "visual_encoder"):
             v = getattr(config, key, None)
             if isinstance(v, str) and v.strip(): return v.strip()
-        for key in ("mm_vision_tower", "vision_tower"):
-            v = getattr(config, key, None)
-            try:
-                name = getattr(getattr(v, "config", None), "_name_or_path", None)
-                if isinstance(name, str) and name.strip(): return name.strip()
-            except Exception:
-                pass
+        try:
+            v = getattr(config, "vision_tower", None)
+            name = getattr(getattr(v, "config", None), "_name_or_path", None)
+            if isinstance(name, str) and name.strip(): return name.strip()
+        except Exception:
+            pass
         return DEFAULT_VISION_TOWER_ID
 
     def _normalize_image_processor(self) -> bool:
@@ -477,26 +362,7 @@ class EndpointHandler:
         conv.append_message(conv.roles[1], None)
         return conv.get_prompt()
 
-    def _create_robust_attention_mask(self, input_ids: torch.Tensor) -> Optional[torch.Tensor]:
-        try:
-            if input_ids is None or not isinstance(input_ids, torch.Tensor):
-                print("[warn] Invalid input_ids for attention mask creation")
-                return None
-            
-            device = input_ids.device
-            attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=device)
-            
-            if self.tokenizer.pad_token_id is not None:
-                attention_mask = attention_mask.masked_fill(input_ids == self.tokenizer.pad_token_id, 0)
-            
-            print(f"[debug] Created attention_mask: shape={attention_mask.shape}, device={device}")
-            return attention_mask
-            
-        except Exception as e:
-            print(f"[error] Failed to create attention_mask: {e}")
-            return None
-
-    # ------------- IMPROVED: inference with better error handling -------------
+    # ------------- inference -------------
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         inputs = data.get("inputs") or {}
         params = data.get("parameters") or {}
@@ -543,7 +409,7 @@ class EndpointHandler:
                 import traceback; traceback.print_exc()
                 images = None; image_sizes = None
 
-        # 3) tokenize + robust mask
+        # 3) tokenize
         try:
             mdev = next(self.model.parameters()).device
             input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt') \
@@ -551,12 +417,15 @@ class EndpointHandler:
             print(f"[debug] input_ids shape: {input_ids.shape} | has images: {images is not None}")
         except Exception as e:
             print(f"[error] Tokenization failed: {e}")
-            input_ids = self.tokenizer(query_text, return_tensors="pt").input_ids.to(next(self.model.parameters()).device)
-            images = None; image_sizes = None
-
-        attention_mask = self._create_robust_attention_mask(input_ids)
+            try:
+                input_ids = self.tokenizer(query_text, return_tensors="pt").input_ids.to(next(self.model.parameters()).device)
+                images = None; image_sizes = None
+                print("[warn] Fallback to basic tokenization without image tokens")
+            except Exception as e2:
+                print(f"[error] Even basic tokenization failed: {e2}")
+                return [{"generated_text": f"Error: Tokenization failed: {str(e)}"}]
 
-         # 4) gen params
+        # 4) gen params (attention_mask YOK)
         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))
@@ -568,76 +437,50 @@ class EndpointHandler:
         if max_new_tokens < 1:
             return [{"generated_text": "Error: Input too long, exceeds max token length."}]
 
-        # --- Strategy 1: Normal path with attention_mask ---
         gen_kwargs = {
+            # KRİTİK: Hem `inputs` hem de `input_ids` veriyoruz
             "inputs": input_ids,
-            "attention_mask": attention_mask,
-            "use_cache": bool(params.get("use_cache", True)),
+            "input_ids": input_ids,
             "max_new_tokens": max_new_tokens,
             "temperature": temperature,
             "top_p": top_p,
             "repetition_penalty": repetition_penalty,
             "do_sample": do_sample,
+            # attention_mask verme!
+            "use_cache": bool(params.get("use_cache", True)),
             "pad_token_id": self.tokenizer.pad_token_id,
             "eos_token_id": getattr(self.tokenizer, "eos_token_id", None),
             "bos_token_id": getattr(self.tokenizer, "bos_token_id", None),
         }
-        if images is not None:
+        if images is not None and image_sizes is not None:
             gen_kwargs["images"] = images
             gen_kwargs["image_sizes"] = image_sizes
 
+        # 5) generate
         try:
-            print("[debug] Trying generation: normal path (with attention_mask)")
             with torch.inference_mode():
                 output = self.model.generate(**gen_kwargs)
-        except ValueError as e:
-            msg = str(e)
-            if "model_kwargs" in msg and "attention_mask" in msg and "not used" in msg:
-                print("[hotfix] model rejected attention_mask; retrying via mask injection (no kwargs mask) + use_cache=False")
-                gen_kwargs_no_mask = {k: v for k, v in gen_kwargs.items() if k not in ("attention_mask", "inputs")}
-                gen_kwargs_no_mask["use_cache"] = False
-                output = self._generate_with_injected_mask(
-                    input_ids=input_ids,
-                    images=images,
-                    image_sizes=image_sizes,
-                    attention_mask=attention_mask,
-                    base_kwargs=gen_kwargs_no_mask
-                )
-            else:
-                print(f"Generation error: {e}")
-                import traceback; traceback.print_exc()
-                return [{"generated_text": f"Error during generation: {msg}"}]
         except Exception as e:
-            emsg = str(e)
-            print(f"[warn] Normal path failed: {emsg}")
-            print("[hotfix] retry via NO-MASK fallback")
+            # Son çare: cache kapalı tekrar dene
+            print(f"[warn] First generate failed: {e} | retry with use_cache=False")
+            gen_kwargs["use_cache"] = False
             try:
-                output = self._generate_without_mask(
-                    input_ids=input_ids,
-                    images=images,
-                    image_sizes=image_sizes,
-                    base_kwargs=gen_kwargs
-                )
+                with torch.inference_mode():
+                    output = self.model.generate(**gen_kwargs)
             except Exception as e2:
-                print(f"[error] No-mask fallback failed: {e2}")
+                print(f"[error] Generation failed: {e2}")
                 import traceback; traceback.print_exc()
-                return [{"generated_text": f"Error during generation: {emsg} | fallback: {e2}"}]
+                return [{"generated_text": f"Error during generation: {str(e2)}"}]
 
-        # 5) Decode response
+        # 6) decode
         try:
             sequences = output.sequences if hasattr(output, "sequences") else output
             input_len = input_ids.shape[1]
-            if sequences.shape[-1] > input_len:
-                response_ids = sequences[:, input_len:]
-            else:
-                response_ids = sequences
+            response_ids = sequences[:, input_len:] if sequences.shape[-1] > input_len else sequences
             text = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True)[0].strip()
-            return [{
-                "generated_text": text,
-                "input_tokens": int(input_len),
-                "output_tokens": int(response_ids.shape[-1]),
-                "strategy_used": "normal_or_injected_mask_or_nomask"
-            }]
+            if not text:
+                text = "Error: Empty response generated"
+            return [{"generated_text": text}]
         except Exception as e:
-            print(f"[error] Decoding failed: {e}")
-            return [{"generated_text": f"Error during decoding: {str(e)}"}]
+            print(f"[error] Response decoding failed: {e}")
+            return [{"generated_text": f"Error: Response decoding failed: {str(e)}"}]