fixing parameters

main.py

@@ -118,21 +118,28 @@ def predict(data: PredictRequest):
     Generates a description for an image using the Qwen-2-VL model.
 
     Args:
-        data (any): The encoded image and the prompt to be used.
-        prompt (str): The text prompt to guide the model's response.
+        data (PredictRequest): The request containing encoded images and a prompt.
 
     Returns:
-        str: The generated description of the image.
+        dict: The generated description of the image(s).
     """
 
-    # Create the input message structure
+    # Ensure image_base64 is a list (even if a single image is provided)
+    image_list = (
+        data.image_base64
+        if isinstance(data.image_base64, list)
+        else [data.image_base64]
+    )
+
+    # Create the input message structure with multiple images
     messages = [
         {
             "role": "user",
             "content": [
-                {"type": "image", "image": f"data:image;base64,{data.image_base64}"},
-                {"type": "text", "text": data.prompt},
-            ],
+                {"type": "image", "image": f"data:image;base64,{image}"}
+                for image in image_list
+            ]
+            + [{"type": "text", "text": data.prompt}],
         }
     ]