Create handler.py

handler.py

@@ -0,0 +1,144 @@
+from typing import Dict, List, Any
+from transformers import AutoModelForVision2Seq, AutoTokenizer, AutoProcessor
+from PIL import Image
+import torch
+import io
+import base64
+from peft import PeftModel
+
+class EndpointHandler():
+    def __init__(self, model_dir: str):
+        self.path = model_dir
+        # Load base model and tokenizer
+        base_model_id = "Qwen/Qwen2-VL-2B-Instruct"
+        
+        # Load tokenizer/processor
+        self.processor = AutoProcessor.from_pretrained(
+            self.path,
+            trust_remote_code=True
+        )
+        
+        # Load base model
+        self.model = AutoModelForVision2Seq.from_pretrained(
+            base_model_id,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        
+        # Load LoRA adapter
+        self.model = PeftModel.from_pretrained(
+            self.model,
+            self.path,
+            device_map="auto"
+        )
+        
+        # Merge and unload for faster inference
+        self.model = self.model.merge_and_unload()
+        
+        # Set to eval mode
+        self.model.eval()
+        
+        # Store the instruction template
+        self.instruction = """
+A conversation between a Healthcare Provider and an AI Medical Image Analysis Assistant. The provider shares a medical image, and the Assistant generates a clear description/report. The assistant first analyzes the image systematically, then provides a concise report. The analysis process and report are enclosed within <thinking> </thinking><answer> </answer>.
+Always respond in this format:
+<thinking>
+1. Initial Assessment:
+   - What type of image is this? (X-ray, CT, MRI, etc.)
+   - Which body part/region is shown?
+   - Is the image quality adequate?
+2. Key Findings:
+   - What are the normal structures visible?
+   - Are there any abnormalities?
+   - What are the important measurements?
+3. Clinical Significance:
+   - What are the main clinical findings?
+   - Are there any critical findings?
+</thinking>
+<answer>
+Brief Structured Report:
+1. EXAM TYPE: [imaging type and body region]
+2. FINDINGS: [key observations and abnormalities]
+3. IMPRESSION: [summary and clinical significance]
+</answer>
+"""
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
+        """
+        data args:
+            inputs (:obj: `str` | `PIL.Image` | `np.array`)
+            parameters (:obj: `Dict[str, Any]`, *optional*)
+        Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # Extract inputs and parameters
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+        
+        # Handle different input formats
+        if isinstance(inputs, str):
+            # Base64 encoded image
+            image_bytes = base64.b64decode(inputs)
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        elif isinstance(inputs, dict):
+            # Dictionary with image key
+            image_data = inputs.get("image", inputs.get("inputs", ""))
+            if isinstance(image_data, str):
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+            else:
+                image = image_data
+        else:
+            # Direct image
+            image = inputs
+            
+        # Ensure image is RGB
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        
+        # Prepare messages in Qwen format
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": self.instruction}
+                ]
+            }
+        ]
+        
+        # Process inputs
+        text = self.processor.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        
+        # Prepare inputs for model
+        inputs = self.processor(
+            text=[text],
+            images=[image],
+            padding=True,
+            return_tensors="pt"
+        ).to(self.model.device)
+        
+        # Generate response
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                **inputs,
+                max_new_tokens=parameters.get("max_new_tokens", 512),
+                temperature=parameters.get("temperature", 0.7),
+                top_p=parameters.get("top_p", 0.9),
+                do_sample=True,
+                pad_token_id=self.processor.tokenizer.pad_token_id,
+                eos_token_id=self.processor.tokenizer.eos_token_id,
+            )
+        
+        # Decode output - only the generated part
+        output_text = self.processor.batch_decode(
+            output_ids[:, inputs.input_ids.shape[1]:],
+            skip_special_tokens=True
+        )[0]
+        
+        return [{"generated_text": output_text}]