app_local.py

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 re
import gc
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

#############################
os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')

# Model configuration
REWRITER_MODEL = "Qwen/Qwen1.5-1.8B-Chat"
rewriter_tokenizer = None
rewriter_model = None
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Quantization configuration
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True
)

def load_rewriter():
    """Lazily load the prompt enhancement model"""
    global rewriter_tokenizer, rewriter_model
    if rewriter_tokenizer is None or rewriter_model is None:
        print("🔄 Loading enhancement model...")
        rewriter_tokenizer = AutoTokenizer.from_pretrained(REWRITER_MODEL)
        rewriter_model = AutoModelForCausalLM.from_pretrained(
            REWRITER_MODEL,
            torch_dtype=dtype,
            device_map="auto",
            quantization_config=bnb_config
        )
        print("✅ Enhancement model loaded")

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": "..."
}
'''

def polish_prompt(original_prompt: str) -> str:
    """Enhanced prompt rewriting using Qwen1.5-1.8B"""
    load_rewriter()
    
    # Format as Qwen chat with system prompt
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT_EDIT},
        {"role": "user", "content": original_prompt}
    ]
    
    # Generate enhanced prompt
    text = rewriter_tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    model_inputs = rewriter_tokenizer(text, return_tensors="pt").to(device)
    
    with torch.no_grad():
        generated_ids = rewriter_model.generate(
            **model_inputs,
            max_new_tokens=120,
            do_sample=True,
            temperature=0.7,
            top_p=0.95,
            no_repeat_ngram_size=2
        )
    
    # Extract and clean response
    enhanced = rewriter_tokenizer.decode(
        generated_ids[0][model_inputs.input_ids.shape[1]:],
        skip_special_tokens=True
    )
    
    # Clean possible artifacts
    enhanced = enhanced.strip()
    if enhanced.lower().startswith(("rewritten instruction:", "enhanced:", "output:")):
        enhanced = re.split(r':', enhanced, 1)[-1].strip()
    
    # Remove any quotes around the prompt if present
    if enhanced.startswith('"') and enhanced.endswith('"'):
        enhanced = enhanced[1:-1]
    
    return enhanced

# Load main image editing pipeline
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")

def unload_rewriter():
    """Clear enhancement model from memory"""
    global rewriter_tokenizer, rewriter_model
    if rewriter_model:
        del rewriter_tokenizer, rewriter_model
        rewriter_tokenizer = None
        rewriter_model = None
    torch.cuda.empty_cache()
    gc.collect()

@spaces.GPU(duration=60)
def infer(
    image,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=4.0,
    num_inference_steps=8,
    rewrite_prompt=False,
    num_images_per_prompt=1,
):
    """Image editing endpoint with optimized prompt handling"""
    original_prompt = prompt
    prompt_info = ""
    
    # Handle prompt rewriting
    if rewrite_prompt:
        try:
            enhanced_instruction = polish_prompt(original_prompt)
            prompt_info = (
                f"<div style='margin:10px; padding:10px; border-radius:8px; border-left:4px solid #4CAF50; background: #f5f9fe'>"
                f"<h4 style='margin-top: 0;'>🚀 Prompt Enhancement</h4>"
                f"<p><strong>Original:</strong> {original_prompt}</p>"
                f"<p><strong>Enhanced:</strong> {enhanced_instruction}</p>"
                f"</div>"
            )
            prompt = enhanced_instruction
        except Exception as e:
            gr.Warning(f"Prompt enhancement failed: {str(e)}")
            prompt_info = (
                f"<div style='margin:10px; padding:10px; border-radius:8px; border-left:4px solid #FF5252; background: #fef5f5'>"
                f"<h4 style='margin-top: 0;'>⚠️ Enhancement Not Applied</h4>"
                f"<p>Using original prompt. Error: {str(e)}</p>"
                f"</div>"
            )
    else:
        prompt_info = (
            f"<div style='margin:10px; padding:10px; border-radius:8px; background: #f8f9fa'>"
            f"<h4 style='margin-top: 0;'>📝 Original Prompt</h4>"
            f"<p>{original_prompt}</p>"
            f"</div>"
        )
    
    # Free VRAM after enhancement
    unload_rewriter()
    
    # Set seed for reproducibility
    seed_val = seed
    if randomize_seed:
        seed_val = random.randint(0, 2**32 - 1)
    generator = torch.Generator(device=device).manual_seed(seed_val)
    
    try:
        # Generate images
        edited_images = pipe(
            image=image,
            prompt=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
    except Exception as e:
        gr.Error(f"Image generation failed: {str(e)}")
        prompt_info = (
            f"<div style='margin:10px; padding:10px; border-radius:8px; border-left:4px solid #dd2c00; background: #fef5f5'>"
            f"<h4 style='margin-top: 0;'><strong>⚠️ Error:</strong> {str(e)}</h4>"
            f"</div>"
        )
        return [], seed_val, prompt_info
    
    return edited_images, seed_val, 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 and position it at top left."""
]

with gr.Blocks(title="Qwen Image Editor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    <div style="text-align: center;">
        <h1>⚡️ Qwen-Image-Edit Lightning</h1>
        <p>8-step image editing with local prompt enhancement | Powered by NVIDIA H200</p>
    </div>
    """)

    with gr.Row():
        # Input Column
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil")
            prompt = gr.Textbox(label="Edit Instruction", placeholder="e.g. Add a dog to the right side", lines=2)
            
            with gr.Accordion("Advanced Settings", open=False):
                gr.Markdown("### Generation Parameters")
                with gr.Row():
                    seed = gr.Slider(label="Seed", min=0, max=MAX_SEED, step=1, value=42)
                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                with gr.Row():
                    true_guidance_scale = gr.Slider(
                        label="Guidance Scale", min=1.0, max=5.0, step=0.1, value=4.0
                    )
                    num_inference_steps = gr.Slider(
                        label="Inference Steps", min=4, max=16, step=1, value=8
                    )
                    num_images_per_prompt = gr.Slider(
                        label="Output Images", min=1, max=4, step=1, value=1
                    )
            
            rewrite_toggle = gr.Checkbox(
                label="Enable AI Prompt Enhancement", 
                value=True,
                info="Uses local Qwen1.5-1.8B model to improve your instructions"
            )
            
            run_button = gr.Button("Generate Edits", variant="primary")

        # Output Column
        with gr.Column():
            result = gr.Gallery(
                label="Output Images", 
                columns=lambda x: 2 if x > 1 else 1,
                object_fit="contain",
                height="auto"
            )
            prompt_info = gr.HTML(
                "<div style='margin-top:20px; padding:15px; border-radius:8px; background:#f8f9fa'>"
                "<p>Prompt details will appear here after generation</p></div>"
            )
    
    gr.Examples(
        examples=examples,
        inputs=[prompt],
        label="Try These Examples",
        cache_examples=True
    )

    # Main processing handler
    inputs = [
        input_image,
        prompt,
        seed,
        randomize_seed,
        true_guidance_scale,
        num_inference_steps,
        rewrite_toggle,
        num_images_per_prompt
    ]
    
    outputs = [result, seed, prompt_info]
    
    run_button.click(
        fn=infer,
        inputs=inputs,
        outputs=outputs
    )
    
    prompt.submit(
        fn=infer,
        inputs=inputs,
        outputs=outputs
    )

if __name__ == "__main__":
    demo.launch()