Update handler.py

handler.py

@@ -34,30 +34,32 @@ class EndpointHandler:
                 return {"error": f"Failed to decode or open base64 image: {e}"}
 
         if image is not None:
-            # --- Case 1: Multimodal (Image + Text) - This part is correct ---
+            # --- Case 1: Multimodal (Image + Text) ---
             print("Processing multimodal request...")
             prompt = f"USER: <image>\n{prompt_text} ASSISTANT:"
             inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(self.model.device)
         else:
-            # --- Case 2: Text-Only - This is the corrected logic ---
+            # --- Case 2: Text-Only - CORRECTED LOGIC ---
             print("Processing text-only request...")
             prompt = f"USER: {prompt_text} ASSISTANT:"
 
             # First, process the text to get input_ids
             inputs = self.processor(text=prompt, return_tensors="pt")
             
-            # Second, create a dummy image tensor with the correct shape, type, and device.
-            # This placeholder satisfies the model's input requirements.
+            # --- THE FIX: Get image dimensions from the processor's .config ---
             image_processor = self.processor.image_processor
+            config = image_processor.config
+            
+            # Create a dummy image tensor using the correct config values
             dummy_pixel_values = torch.zeros(
                 (
                     1, 
-                    image_processor.num_channels, 
-                    image_processor.size['height'], 
-                    image_processor.size['width']
+                    config.num_channels, 
+                    config.crop_size['height'], 
+                    config.crop_size['width']
                 ),
-                dtype=self.model.dtype,      # Use the model's dtype (e.g., float16)
-                device=self.model.device    # Put the tensor on the same device as the model
+                dtype=self.model.dtype,
+                device=self.model.device
             )
             
             # Add the dummy tensor to the inputs dictionary