Update app.py

app.py

@@ -16,6 +16,7 @@ import math
 
 import os
 import base64
+from io import BytesIO
 import json
 
 SYSTEM_PROMPT = '''
@@ -98,7 +99,7 @@ Please strictly follow the rewriting rules below:
 }
 '''
 
-def polish_prompt_hf(prompt, img_list):
+def polish_prompt_hf(original_prompt, img_list):
     """
     Rewrites the prompt using a Hugging Face InferenceClient.
     Supports multiple images via img_list.
@@ -107,8 +108,8 @@ def polish_prompt_hf(prompt, img_list):
     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:"
+        return original_prompt
+    prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {original_prompt}\n\nRewritten Prompt:"
     system_prompt = "you are a helpful assistant, you should provide useful answers to users."
     try:
         # Initialize the client
@@ -197,7 +198,7 @@ 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 prompt 
+        return original_prompt