letting windsurf handle it

handler.py

@@ -1,10 +1,20 @@
-from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from typing import Dict, List, Any, Optional, Union
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor
 import torch
+import base64
+from io import BytesIO
+from PIL import Image
+import requests
 
 
 class EndpointHandler:
     def __init__(self, path=""):
+        # If path is empty, use the GSAI-ML/LLaDA-8B-Instruct model
+        if not path:
+            path = "GSAI-ML/LLaDA-8B-Instruct"
+
+        print(f"Loading model from {path}...")
+
         # Load model with half precision to save memory
         self.model = AutoModelForCausalLM.from_pretrained(
             path, torch_dtype=torch.float16, device_map="auto"
@@ -13,6 +23,9 @@ class EndpointHandler:
         # Load tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(path)
 
+        # Load processor for handling images
+        self.processor = AutoProcessor.from_pretrained(path)
+
         # Ensure pad token is properly set
         if self.tokenizer.pad_token_id is None:
             if (
@@ -28,15 +41,52 @@ class EndpointHandler:
 
         print(f"Model loaded successfully. Pad token ID: {self.tokenizer.pad_token_id}")
 
+    def _load_image(self, image_data: Union[str, bytes]) -> Image.Image:
+        """Load image from URL or base64 encoded string"""
+        if isinstance(image_data, str):
+            if image_data.startswith("http"):
+                # Load from URL
+                response = requests.get(image_data, stream=True)
+                response.raise_for_status()
+                return Image.open(BytesIO(response.content))
+            elif image_data.startswith("data:image"):
+                # Handle base64 encoded image
+                base64_data = image_data.split(",")[1]
+                image_bytes = base64.b64decode(base64_data)
+                return Image.open(BytesIO(image_bytes))
+            else:
+                # Assume it's a base64 string without the prefix
+                try:
+                    image_bytes = base64.b64decode(image_data)
+                    return Image.open(BytesIO(image_bytes))
+                except Exception as e:
+                    raise ValueError(f"Invalid image data format: {e}")
+        elif isinstance(image_data, bytes):
+            return Image.open(BytesIO(image_data))
+        else:
+            raise ValueError(f"Unsupported image data type: {type(image_data)}")
+
+    def _format_prompt(self, text: str, system_prompt: Optional[str] = None) -> str:
+        """Format the prompt according to LLaDA's expected format"""
+        # Default system prompt for LLaDA if none provided
+        if system_prompt is None:
+            system_prompt = (
+                "You are a helpful AI assistant that can understand images and text."
+            )
+
+        # Format the prompt following LLaDA's expected structure
+        formatted_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        return formatted_prompt
+
     def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:
         """Handle inference requests"""
         # Extract inputs and parameters from request data
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
 
-        # Convert single string input to list for consistent handling
-        if isinstance(inputs, str):
-            inputs = [inputs]
+        # Extract image data if present
+        image_data = parameters.get("image", None)
+        system_prompt = parameters.get("system_prompt", None)
 
         # Extract generation parameters with sensible defaults
         max_new_tokens = parameters.get("max_new_tokens", 256)
@@ -44,24 +94,75 @@ class EndpointHandler:
         top_p = parameters.get("top_p", 0.95)
         do_sample = parameters.get("do_sample", True)
 
-        # Tokenize inputs
-        input_tokens = self.tokenizer(inputs, return_tensors="pt", padding=True).to(
-            self.model.device
-        )
+        # Convert single string input to list for consistent handling
+        if isinstance(inputs, str):
+            inputs = [inputs]
 
-        # Generate text
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **input_tokens,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=do_sample,
-                pad_token_id=self.tokenizer.pad_token_id,
-            )
+        # Process each input
+        generated_texts = []
+        for input_text in inputs:
+            # Format the prompt according to LLaDA's expected format
+            formatted_prompt = self._format_prompt(input_text, system_prompt)
+
+            if image_data:
+                try:
+                    # Process image if present
+                    image = self._load_image(image_data)
+                    inputs_processor = self.processor(
+                        text=formatted_prompt, images=image, return_tensors="pt"
+                    )
+
+                    # Move inputs to the same device as the model
+                    for k, v in inputs_processor.items():
+                        if isinstance(v, torch.Tensor):
+                            inputs_processor[k] = v.to(self.model.device)
+
+                    # Generate text with image context
+                    with torch.no_grad():
+                        outputs = self.model.generate(
+                            **inputs_processor,
+                            max_new_tokens=max_new_tokens,
+                            temperature=temperature,
+                            top_p=top_p,
+                            do_sample=do_sample,
+                            pad_token_id=self.tokenizer.pad_token_id,
+                        )
+
+                    # Decode generated text
+                    generated_text = self.tokenizer.decode(
+                        outputs[0], skip_special_tokens=True
+                    )
+                    generated_texts.append(generated_text)
+
+                except Exception as e:
+                    # If image processing fails, fall back to text-only
+                    print(
+                        f"Error processing image: {e}. Falling back to text-only processing."
+                    )
+                    image_data = None
+
+            if not image_data:
+                # Text-only processing
+                input_tokens = self.tokenizer(formatted_prompt, return_tensors="pt").to(
+                    self.model.device
+                )
+
+                # Generate text
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **input_tokens,
+                        max_new_tokens=max_new_tokens,
+                        temperature=temperature,
+                        top_p=top_p,
+                        do_sample=do_sample,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                    )
 
-        # Decode generated text
-        generated_texts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+                # Decode generated text
+                generated_text = self.tokenizer.decode(
+                    outputs[0], skip_special_tokens=True
+                )
+                generated_texts.append(generated_text)
 
         # Return results in expected format
         return {"generated_text": generated_texts}

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.0.0
+transformers>=4.30.0
+pillow>=9.0.0
+requests>=2.25.0
+accelerate>=0.20.0