Update app.py

app.py

@@ -4,7 +4,10 @@ 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
@@ -41,7 +44,7 @@ def polish_prompt(original_prompt, system_prompt, hf_token):
         completion = client.chat.completions.create(
             model="Qwen/Qwen3-235B-A22B-Instruct-2507",
             messages=messages,
-            max_tokens=2000,
+            max_tokens=512,
         )
         polished_prompt = completion.choices[0].message.content
         polished_prompt = polished_prompt.strip().replace("\n", " ")
@@ -92,6 +95,10 @@ pipe.load_lora_weights(
 )
 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(
@@ -143,66 +150,71 @@ examples = [
 ]
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Qwen-Image-Edit with Prompt Enhancement")
+    gr.Markdown("# Qwen-Image-Edit [FAST] with HF Inference Prompt Enhancement")
+    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)**")
-    
-    with gr.Column():
-        input_image = gr.Image(label="Input Image", type="pil")
-        prompt = gr.Text(label="Edit Instruction", placeholder="e.g. Add a dog to the right side.")
-        run_button = gr.Button("Edit", variant="primary")
-        result = gr.Gallery(label="Output Images", show_label=False)
-    
-    with gr.Accordion("Advanced Settings", open=False):
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
-            step=1,
-            value=0
-        )
-        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-        
-        with gr.Row():
-            true_guidance_scale = gr.Slider(
-                label="True Guidance Scale",
-                minimum=1.0,
-                maximum=5.0,
-                step=0.1,
-                value=4.0
-            )
-            num_inference_steps = gr.Slider(
-                label="Inference Steps (Fast 8-step mode)",
-                minimum=4,
-                maximum=8,
-                step=1,
-                value=8
-            )
-            num_images_per_prompt = gr.Slider(
-                label="Images per Prompt",
-                minimum=1,
-                maximum=4,
-                step=1,
-                value=1
-            )
-        
-        with gr.Group():
-            rewrite_toggle = gr.Checkbox(label="Use Prompt Rewriter (Requires HF Token)", value=False, interactive=True)
-            hf_token_input = gr.Textbox(
-                label="Your Hugging Face Token",
-                type="password",
-                placeholder="hf_xxxxxxxxxxxxxxxx",
-                visible=False,
-                info="Required for prompt rewriting - get yours from [Hugging Face settings](https://huggingface.co/settings/tokens)"
-            )
-            
-            def toggle_token_visibility(checked):
-                return gr.update(visible=checked)
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input Image", type="pil")
+            prompt = gr.Text(label="Edit Instruction", placeholder="e.g. Add a dog to the right side.")
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0
+                )
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                
+                with gr.Row():
+                    true_guidance_scale = gr.Slider(
+                        label="True Guidance Scale",
+                        minimum=1.0,
+                        maximum=5.0,
+                        step=0.1,
+                        value=4.0
+                    )
+                    num_inference_steps = gr.Slider(
+                        label="Inference Steps (Fast 8-step mode)",
+                        minimum=4,
+                        maximum=16,
+                        step=1,
+                        value=8
+                    )
+                    num_images_per_prompt = gr.Slider(
+                        label="Images per Prompt",
+                        minimum=1,
+                        maximum=4,
+                        step=1,
+                        value=1
+                    )
+                with gr.Group():
+                    rewrite_toggle = gr.Checkbox(label="Use Prompt Rewriter (Requires HF Token)", value=False, interactive=True)
+                    hf_token_input = gr.Textbox(
+                        label="Your Hugging Face Token",
+                        type="password",
+                        placeholder="hf_xxxxxxxxxxxxxxxx",
+                        visible=False,
+                        info="Required for prompt rewriting - get yours from [Hugging Face settings](https://huggingface.co/settings/tokens)"
+                    )
+                    
+                    def toggle_token_visibility(checked):
+                        return gr.update(visible=checked)
+                    
+                    rewrite_toggle.change(
+                        toggle_token_visibility,
+                        inputs=[rewrite_toggle],
+                        outputs=[hf_token_input]
+                    )
+
+            run_button = gr.Button("Edit", variant="primary")
             
-            rewrite_toggle.change(
-                toggle_token_visibility,
-                inputs=[rewrite_toggle],
-                outputs=[hf_token_input]
-            )
+        with gr.Column():
+            result = gr.Gallery(label="Output Images", show_label=False)
+
+        
 
     gr.on(
         triggers=[run_button.click, prompt.submit],