Update app.py

app.py

@@ -97,63 +97,79 @@ Please strictly follow the rewriting rules below:
    "Rewritten": "..."
 }
 '''
-# --- Prompt Enhancement using Hugging Face InferenceClient ---
+
 def polish_prompt_hf(prompt, img_list):
     """
     Rewrites the prompt using a Hugging Face InferenceClient.
+    Supports multiple images via img_list.
     """
     # Ensure HF_TOKEN is set
     api_key = os.environ.get("HF_TOKEN")
     if not api_key:
         print("Warning: HF_TOKEN not set. Falling back to original prompt.")
         return prompt
-
+    prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
+    system_prompt = "you are a helpful assistant, you should provide useful answers to users."
     try:
         # Initialize the client
-        prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
         client = InferenceClient(
             provider="nebius",
             api_key=api_key,
         )
 
+        # Convert list of images to base64 data URLs
+        image_urls = []
+        if img_list is not None:
+            # Ensure img_list is actually a list
+            if not isinstance(img_list, list):
+                img_list = [img_list]
+            
+            for img in img_list:
                 image_url = None
-        if img is not None:
-            # If img is a PIL Image
-            if hasattr(img, 'save'):  # Check if it's a PIL Image
-                buffered = BytesIO()
-                img.save(buffered, format="PNG")
-                img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
-                image_url = f"data:image/png;base64,{img_base64}"
-            # If img is already a file path (string)
-            elif isinstance(img, str):
-                with open(img, "rb") as image_file:
-                    img_base64 = base64.b64encode(image_file.read()).decode('utf-8')
-                image_url = f"data:image/png;base64,{img_base64}"
-            else:
-                print(f"Warning: Unexpected image type: {type(img)}")
-                return original_prompt
+                # If img is a PIL Image
+                if hasattr(img, 'save'):  # Check if it's a PIL Image
+                    buffered = BytesIO()
+                    img.save(buffered, format="PNG")
+                    img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+                    image_url = f"data:image/png;base64,{img_base64}"
+                # If img is already a file path (string)
+                elif isinstance(img, str):
+                    with open(img, "rb") as image_file:
+                        img_base64 = base64.b64encode(image_file.read()).decode('utf-8')
+                    image_url = f"data:image/png;base64,{img_base64}"
+                else:
+                    print(f"Warning: Unexpected image type: {type(img)}, skipping...")
+                    continue
+                
+                if image_url:
+                    image_urls.append(image_url)
+
+        # Build the content array with text first, then all images
+        content = [
+            {
+                "type": "text",
+                "text": prompt
+            }
+        ]
+        
+        # Add all images to the content
+        for image_url in image_urls:
+            content.append({
+                "type": "image_url",
+                "image_url": {
+                    "url": image_url
+                }
+            })
 
         # Format the messages for the chat completions API
         messages = [
             {"role": "system", "content": system_prompt},
             {
                 "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": original_prompt
-                    },
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": image_url
-                        }
-                    }
-                ]
+                "content": content
             }
         ]
 
-
         # Call the API
         completion = client.chat.completions.create(
             model="Qwen/Qwen2.5-VL-72B-Instruct",
@@ -181,7 +197,9 @@ def polish_prompt_hf(prompt, img_list):
     except Exception as e:
         print(f"Error during API call to Hugging Face: {e}")
         # Fallback to original prompt if enhancement fails
-        return original_prompt
+        return prompt 
+
+
 
 def encode_image(pil_image):
     import io