Update app.py

app.py

@@ -4,15 +4,19 @@ import random
 import torch
 import spaces
 from PIL import Image
-
 from diffusers import QwenImageEditPipeline
 from diffusers.utils import is_xformers_available
-
 import os
 import base64
 import json
 from huggingface_hub import InferenceClient
-
+import logging
+#############################
+os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
+os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+#############################
 def get_caption_language(prompt):
     """Detects if the prompt contains Chinese characters."""
     ranges = [
@@ -22,7 +26,6 @@ def get_caption_language(prompt):
         if any(start <= char <= end for start, end in ranges):
             return 'zh'
     return 'en'
-
 def polish_prompt(original_prompt, system_prompt, hf_token):
     """
     Rewrites the prompt using a Hugging Face InferenceClient.
@@ -31,7 +34,6 @@ def polish_prompt(original_prompt, system_prompt, hf_token):
     if not hf_token or not hf_token.strip():
         gr.Warning("HF Token is required for prompt rewriting but was not provided!")
         return original_prompt
-
     client = InferenceClient(
         provider="cerebras",
         api_key=hf_token,
@@ -53,7 +55,6 @@ def polish_prompt(original_prompt, system_prompt, hf_token):
         print(f"Error during Hugging Face API call: {e}")
         gr.Warning("Failed to rewrite prompt. Using original.")
         return original_prompt
-
 SYSTEM_PROMPT_EDIT = '''
 # Edit Instruction Rewriter
 You are a professional edit instruction rewriter. Your task is to generate a precise, concise, and visually achievable instruction based on the user's intent and the input image.
@@ -85,7 +86,6 @@ Please provide the rewritten instruction in a clean `json` format as:
   "Rewritten": "..."
 }
 '''
-
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
@@ -94,12 +94,10 @@ pipe.load_lora_weights(
     "lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
 )
 pipe.fuse_lora()
-
 if is_xformers_available():
     pipe.enable_xformers_memory_efficient_attention()
 else:
     print("xformers not available or failed to load.")
-
 @spaces.GPU(duration=60)
 def infer(
     image,
@@ -116,21 +114,49 @@ def infer(
     """
     Requires user-provided HF token for prompt rewriting.
     """
+    original_prompt = prompt  # Save original prompt for display
     negative_prompt = " "
+    prompt_info = ""  # Initialize info text
+    
+    # Handle prompt rewriting with status messages
+    if rewrite_prompt:
+        if not hf_token.strip():
+            gr.Warning("HF Token is required for prompt rewriting but was not provided!")
+            prompt_info = f"""## ⚠️ Prompt Rewriting Skipped (No HF Token)
+**Original Prompt:**  
+{original_prompt}"""
+            rewritten_prompt = original_prompt
+        else:
+            try:
+                rewritten_prompt = polish_prompt(original_prompt, SYSTEM_PROMPT_EDIT, hf_token)
+                prompt_info = f"""## ✅ Prompt Rewrite Successful
+**Original Prompt:**  
+{original_prompt}
+
+**Enhanced Prompt:**  
+{rewritten_prompt}"""
+            except Exception as e:
+                gr.Warning(f"Prompt rewriting failed: {str(e)}")
+                rewritten_prompt = original_prompt
+                prompt_info = f"""## ❌ Prompt Rewrite Failed
+**Original Prompt:**  
+{original_prompt}
+**Error:**  
+{str(e)}"""
+    else:
+        rewritten_prompt = original_prompt
+        prompt_info = f"""## Original Prompt (No Rewrite)
+**User Input:**  
+{original_prompt}"""
+    
+    # Generate images
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    if rewrite_prompt:
-        lang = get_caption_language(prompt)
-        system_prompt = SYSTEM_PROMPT_EDIT
-        polished_prompt = polish_prompt(prompt, system_prompt, hf_token)
-        print(f"Rewritten Prompt: {polished_prompt}")
-        prompt = polished_prompt
-    
     edited_images = pipe(
         image,
-        prompt=prompt,
+        prompt=rewritten_prompt,
         negative_prompt=negative_prompt,
         num_inference_steps=num_inference_steps,
         generator=generator,
@@ -138,7 +164,7 @@ def infer(
         num_images_per_prompt=num_images_per_prompt,
     ).images
     
-    return edited_images, seed
+    return edited_images, seed, prompt_info
 
 MAX_SEED = np.iinfo(np.int32).max
 examples = [
@@ -154,7 +180,7 @@ with gr.Blocks() as demo:
     gr.Markdown("✨ **8-step lightning inferencing with lightx2v's LoRA.**")
     gr.Markdown("⚠️ **Prompt rewriting requires your own [Hugging Face token](https://huggingface.co/settings/tokens)**")
     gr.Markdown("🚧 **Work in progress, further improvements coming soon.**")
-
+    
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="pil")
@@ -168,7 +194,6 @@ with gr.Blocks() as demo:
                     value=0
                 )
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                
                 with gr.Row():
                     true_guidance_scale = gr.Slider(
                         label="True Guidance Scale",
@@ -191,12 +216,12 @@ with gr.Blocks() as demo:
                         step=1,
                         value=1
                     )
-
-
             run_button = gr.Button("Edit", variant="primary")
-            
+        
         with gr.Column():
             result = gr.Gallery(label="Output Images", show_label=False, columns=1)
+            # New prompt display component
+            prompt_info = gr.Markdown("## Prompt Details", visible=False)
             
             with gr.Group():
                 rewrite_toggle = gr.Checkbox(label="Use Prompt Rewriter (Requires HF Token)", value=False, interactive=True)
@@ -207,7 +232,6 @@ with gr.Blocks() as demo:
                     visible=False,
                     info="Required for prompt rewriting - get yours from [Hugging Face settings](https://huggingface.co/settings/tokens). API tokens are kept safe locally, but as good practice please make sure to double check the source code. Invalid or missing keys will revert to the original prompt entered."
                 )
-                
                 def toggle_token_visibility(checked):
                     return gr.update(visible=checked)
                 
@@ -217,8 +241,6 @@ with gr.Blocks() as demo:
                     outputs=[hf_token_input]
                 )
 
-        
-
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -233,7 +255,24 @@ with gr.Blocks() as demo:
             hf_token_input,
             num_images_per_prompt
         ],
-        outputs=[result, seed],
+        outputs=[result, seed, prompt_info]
+    )
+    
+    # Show prompt info box after processing
+    def set_prompt_visible():
+        return gr.update(visible=True)
+    
+    run_button.click(
+        fn=set_prompt_visible,
+        inputs=None,
+        outputs=[prompt_info],
+        queue=False
+    )
+    prompt.submit(
+        fn=set_prompt_visible,
+        inputs=None,
+        outputs=[prompt_info],
+        queue=False
     )
 
 if __name__ == "__main__":