Update app.py

app.py

@@ -10,6 +10,7 @@ from diffusers import QwenImageEditPipeline
 import os
 import base64
 import json
+from huggingface_hub import InferenceClient
 
 SYSTEM_PROMPT = '''
 # Edit Instruction Rewriter
@@ -71,12 +72,64 @@ Please strictly follow the rewriting rules below:
 }
 '''
 
-def polish_prompt(prompt, img):
-    prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
+# def polish_prompt(original_prompt, system_prompt):
+#     """
+#     Rewrites the prompt using a Hugging Face InferenceClient.
+#     """
+#     # Ensure HF_TOKEN is set
+#     api_key = os.environ.get("HF_TOKEN")
+#     if not api_key:
+#         raise EnvironmentError("HF_TOKEN is not set. Please set it in your environment.")
+
+#     # Initialize the client
+#     client = InferenceClient(
+#         provider="cerebras",
+#         api_key=api_key,
+#     )
+
+#     # Format the messages for the chat completions API
+#     messages = [
+#         {"role": "system", "content": system_prompt},
+#         {"role": "user", "content": original_prompt}
+#     ]
+
+#     try:
+#         # Call the API
+#         completion = client.chat.completions.create(
+#             model="Qwen/Qwen3-235B-A22B-Instruct-2507",
+#             messages=messages,
+#         )
+#         polished_prompt = completion.choices[0].message.content
+#         polished_prompt = polished_prompt.strip().replace("\n", " ")
+#         return polished_prompt
+#     except Exception as e:
+#         print(f"Error during API call to Hugging Face: {e}")
+#         # Fallback to original prompt if enhancement fails
+#         return original_prompt
+
+
+def polish_prompt(prompt, system_prompt):
+    """
+    Rewrites the prompt using a Hugging Face InferenceClient.
+    """
+    # Ensure HF_TOKEN is set
+    api_key = os.environ.get("HF_TOKEN")
+    if not api_key:
+        raise EnvironmentError("HF_TOKEN is not set. Please set it in your environment.")
+    # Initialize the client
+    client = InferenceClient(
+        provider="cerebras",
+        api_key=api_key,
+    )
+    
+    prompt = f"{system_prompt}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
     success=False
     while not success:
         try:
-            result = api(prompt, [img])
+            completion = client.chat.completions.create(
+                model="Qwen/Qwen3-235B-A22B-Instruct-2507",
+                messages=messages,
+            )
             # print(f"Result: {result}")
             # print(f"Polished Prompt: {polished_prompt}")
             if isinstance(result, str):
@@ -91,8 +144,9 @@ def polish_prompt(prompt, img):
             polished_prompt = polished_prompt.replace("\n", " ")
             success = True
         except Exception as e:
-            print(f"[Warning] Error during API call: {e}")
-    return polished_prompt
+            print(f"Error during API call to Hugging Face: {e}")
+            # Fallback to original prompt if enhancement fails
+            return prompt
 
 
 def encode_image(pil_image):