Create app.py

app.py

@@ -0,0 +1,535 @@
+import gradio as gr
+import numpy as np
+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 = [
+        ('\u4e00', '\u9fff'),  # CJK Unified Ideographs
+    ]
+    for char in 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.
+    Requires user-provided HF token for API access.
+    """
+    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,
+    )
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": original_prompt}
+    ]
+    try:
+        completion = client.chat.completions.create(
+            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
+            messages=messages,
+            max_tokens=512,
+        )
+        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 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.
+## 1. General Principles
+- Keep the rewritten instruction **concise** and clear.
+- Avoid contradictions, vagueness, or unachievable instructions.
+- Maintain the core logic of the original instruction; only enhance clarity and feasibility.
+- Ensure new added elements or modifications align with the image's original context and art style.
+## 2. Task Types
+### Add, Delete, Replace:
+- When the input is detailed, only refine grammar and clarity.
+- For vague instructions, infer minimal but sufficient details.
+- For replacement, use the format: `"Replace X with Y"`.
+### Text Editing (e.g., text replacement):
+- Enclose text content in quotes, e.g., `Replace "abc" with "xyz"`.
+- Preserving the original structure and language—**do not translate** or alter style.
+### Human Editing (e.g., change a person's face/hair):
+- Preserve core visual identity (gender, ethnic features).
+- Describe expressions in subtle and natural terms.
+- Maintain key clothing or styling details unless explicitly replaced.
+### Style Transformation:
+- If a style is specified, e.g., `Disco style`, rewrite it to encapsulate the essential visual traits.
+- Use a fixed template for **coloring/restoration**:  
+  `"Restore old photograph, remove scratches, reduce noise, enhance details, high resolution, realistic, natural skin tones, clear facial features, no distortion, vintage photo restoration"`  
+  if applicable.
+## 4. Output Format
+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)
+
+# Load LoRA weights for acceleration
+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,
+    prompt,
+    seed=42,
+    randomize_seed=False,
+    true_guidance_scale=1.0,
+    num_inference_steps=8,
+    rewrite_prompt=False,
+    hf_token="",
+    num_images_per_prompt=1,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    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"""<div class="prompt-info-box warning">
+<h3>⚠️ Prompt Rewriting Skipped</h3>
+<p><strong>Original:</strong> {original_prompt}</p>
+<p class="note">HF Token required for enhancement</p>
+</div>"""
+            rewritten_prompt = original_prompt
+        else:
+            try:
+                rewritten_prompt = polish_prompt(original_prompt, SYSTEM_PROMPT_EDIT, hf_token)
+                prompt_info = f"""<div class="prompt-info-box success">
+<h3>✨ Enhanced Successfully</h3>
+<p><strong>Original:</strong> {original_prompt}</p>
+<p><strong>Enhanced:</strong> {rewritten_prompt}</p>
+</div>"""
+            except Exception as e:
+                gr.Warning(f"Prompt rewriting failed: {str(e)}")
+                rewritten_prompt = original_prompt
+                prompt_info = f"""<div class="prompt-info-box error">
+<h3>❌ Enhancement Failed</h3>
+<p><strong>Original:</strong> {original_prompt}</p>
+<p class="note">Error: {str(e)}</p>
+</div>"""
+    else:
+        rewritten_prompt = original_prompt
+        prompt_info = f"""<div class="prompt-info-box default">
+<h3>📝 Original Prompt</h3>
+<p>{original_prompt}</p>
+</div>"""
+    
+    # Generate images
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    edited_images = pipe(
+        image,
+        prompt=rewritten_prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        true_cfg_scale=true_guidance_scale,
+        num_images_per_prompt=num_images_per_prompt,
+    ).images
+    
+    return edited_images, seed, prompt_info
+
+MAX_SEED = np.iinfo(np.int32).max
+
+examples = [
+    "Replace the cat with a friendly golden retriever. Make it look happier, and add more background details.",
+    "Add text 'Qwen - AI for image editing' in Chinese at the bottom center with a small shadow.",
+    "Change the style to 1970s vintage, add old photo effect, restore any scratches on the wall or window.",
+    "Remove the blue sky and replace it with a dark night cityscape.",
+    """Replace "Qwen" with "通义" in the Image. Ensure Chinese font is used for "通义" and position it to the top left with a light heading-style font."""
+]
+
+# Custom CSS for enhanced visual design
+custom_css = """
+/* Gradient background */
+.gradio-container {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #fecfef 75%, #fecfef 100%);
+    min-height: 100vh;
+}
+
+/* Main container styling */
+.container {
+    max-width: 1400px !important;
+    margin: 0 auto !important;
+    padding: 2rem !important;
+}
+
+/* Card-like sections */
+.gr-box {
+    background: rgba(255, 255, 255, 0.95) !important;
+    backdrop-filter: blur(10px) !important;
+    border-radius: 20px !important;
+    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1) !important;
+    border: 1px solid rgba(255, 255, 255, 0.5) !important;
+    padding: 1.5rem !important;
+    margin-bottom: 1.5rem !important;
+}
+
+/* Header styling */
+h1 {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    font-size: 3rem !important;
+    font-weight: 800 !important;
+    text-align: center;
+    margin-bottom: 0.5rem !important;
+    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
+}
+
+h2 {
+    color: #4a5568 !important;
+    font-size: 1.5rem !important;
+    font-weight: 600 !important;
+    margin-bottom: 1rem !important;
+}
+
+/* Button styling */
+.gr-button-primary {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    border: none !important;
+    color: white !important;
+    font-weight: 600 !important;
+    font-size: 1.1rem !important;
+    padding: 0.8rem 2rem !important;
+    border-radius: 12px !important;
+    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-button-primary:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.5) !important;
+}
+
+/* Input fields styling */
+.gr-input, .gr-text-input, .gr-slider, .gr-dropdown {
+    border-radius: 10px !important;
+    border: 2px solid #e2e8f0 !important;
+    background: white !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-input:focus, .gr-text-input:focus {
+    border-color: #667eea !important;
+    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+}
+
+/* Accordion styling */
+.gr-accordion {
+    background: rgba(255, 255, 255, 0.8) !important;
+    border-radius: 12px !important;
+    border: 1px solid rgba(102, 126, 234, 0.2) !important;
+    overflow: hidden !important;
+}
+
+/* Gallery styling */
+.gr-gallery {
+    border-radius: 12px !important;
+    overflow: hidden !important;
+}
+
+/* Prompt info boxes */
+.prompt-info-box {
+    padding: 1.5rem;
+    border-radius: 12px;
+    margin: 1rem 0;
+    animation: fadeIn 0.5s ease;
+}
+
+.prompt-info-box h3 {
+    margin: 0 0 0.75rem 0;
+    font-size: 1.2rem;
+    font-weight: 600;
+}
+
+.prompt-info-box p {
+    margin: 0.5rem 0;
+    line-height: 1.6;
+}
+
+.prompt-info-box.success {
+    background: linear-gradient(135deg, #d4f4dd 0%, #e3f9e5 100%);
+    border-left: 4px solid #48bb78;
+}
+
+.prompt-info-box.warning {
+    background: linear-gradient(135deg, #fef5e7 0%, #fff9ec 100%);
+    border-left: 4px solid #f6ad55;
+}
+
+.prompt-info-box.error {
+    background: linear-gradient(135deg, #fed7d7 0%, #fee5e5 100%);
+    border-left: 4px solid #fc8181;
+}
+
+.prompt-info-box.default {
+    background: linear-gradient(135deg, #e6f3ff 0%, #f0f7ff 100%);
+    border-left: 4px solid #667eea;
+}
+
+.prompt-info-box .note {
+    font-size: 0.9rem;
+    color: #718096;
+    font-style: italic;
+}
+
+/* Checkbox styling */
+.gr-checkbox {
+    background: white !important;
+    border-radius: 8px !important;
+    padding: 0.5rem !important;
+}
+
+/* Token input field */
+input[type="password"] {
+    font-family: monospace !important;
+    letter-spacing: 0.05em !important;
+}
+
+/* Info badges */
+.gr-markdown p {
+    color: #4a5568;
+    line-height: 1.6;
+}
+
+.gr-markdown a {
+    color: #667eea !important;
+    text-decoration: none !important;
+    font-weight: 500 !important;
+    transition: color 0.3s ease !important;
+}
+
+.gr-markdown a:hover {
+    color: #764ba2 !important;
+    text-decoration: underline !important;
+}
+
+/* Animation */
+@keyframes fadeIn {
+    from {
+        opacity: 0;
+        transform: translateY(10px);
+    }
+    to {
+        opacity: 1;
+        transform: translateY(0);
+    }
+}
+
+/* Slider styling */
+.gr-slider input[type="range"] {
+    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%) !important;
+}
+
+/* Group styling */
+.gr-group {
+    background: rgba(249, 250, 251, 0.8) !important;
+    border-radius: 12px !important;
+    padding: 1rem !important;
+    margin-top: 1rem !important;
+}
+
+/* Loading spinner customization */
+.gr-loading {
+    color: #667eea !important;
+}
+
+/* Example buttons */
+.gr-examples button {
+    background: white !important;
+    border: 2px solid #e2e8f0 !important;
+    border-radius: 8px !important;
+    padding: 0.5rem 1rem !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-examples button:hover {
+    border-color: #667eea !important;
+    background: rgba(102, 126, 234, 0.05) !important;
+}
+"""
+
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎨 Qwen-Image-Edit Lightning")
+    gr.Markdown("✨ **Ultra-fast 8-step image editing with AI-powered prompt enhancement**")
+    gr.Markdown("🔐 **Secure prompt rewriting with your [Hugging Face token](https://huggingface.co/settings/tokens)**")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Group():
+                input_image = gr.Image(
+                    label="📸 Input Image", 
+                    type="pil",
+                    elem_classes="gr-box"
+                )
+                prompt = gr.Text(
+                    label="✏️ Edit Instruction", 
+                    placeholder="e.g. Add a dog to the right side, change the sky to sunset...",
+                    lines=3,
+                    elem_classes="gr-box"
+                )
+                
+                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="Guidance Scale",
+                            minimum=1.0,
+                            maximum=5.0,
+                            step=0.1,
+                            value=4.0
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Inference Steps",
+                            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
+                    )
+                
+                run_button = gr.Button("🚀 Generate Edit", variant="primary", size="lg")
+        
+        with gr.Column(scale=1):
+            result = gr.Gallery(
+                label="🖼️ Output Images", 
+                show_label=True, 
+                columns=2,
+                rows=2,
+                elem_classes="gr-box"
+            )
+            
+            # Prompt display component
+            prompt_info = gr.HTML(visible=False)
+            
+            with gr.Group():
+                rewrite_toggle = gr.Checkbox(
+                    label="🤖 Enable AI Prompt Enhancement", 
+                    value=False, 
+                    interactive=True
+                )
+                hf_token_input = gr.Textbox(
+                    label="🔑 Hugging Face API Token",
+                    type="password",
+                    placeholder="hf_xxxxxxxxxxxxxxxx",
+                    visible=False,
+                    info="Your token is secure and only used for API calls. Get yours from HuggingFace settings.",
+                    elem_classes="gr-box"
+                )
+                
+                def toggle_token_visibility(checked):
+                    return gr.update(visible=checked)
+                
+                rewrite_toggle.change(
+                    toggle_token_visibility,
+                    inputs=[rewrite_toggle],
+                    outputs=[hf_token_input]
+                )
+    
+    # Examples section
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        label="💡 Example Prompts"
+    )
+    
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[
+            input_image,
+            prompt,
+            seed,
+            randomize_seed,
+            true_guidance_scale,
+            num_inference_steps,
+            rewrite_toggle,
+            hf_token_input,
+            num_images_per_prompt
+        ],
+        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__":
+    demo.launch()