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""
f"
🚀 Prompt Enhancement
"
f"
Original: {original_prompt}
"
f"
Enhanced: {enhanced_instruction}
"
f"
"
f"
⚠️ Enhancement Not Applied
"
f"
Using original prompt. Error: {str(e)}
"
f"
"
f"
📝 Original Prompt
"
f"
{original_prompt}
"
f"
"
f"
⚠️ Error: {str(e)}
"
f""
)
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("""
⚡️ Qwen-Image-Edit Lightning
8-step image editing with local prompt enhancement | Powered by NVIDIA H200
""")
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(
""
"
Prompt details will appear here after generation