multimodalart/Qwen-Image-Edit-Fast
494
1import os2os.system('pip install --upgrade spaces')3 4import gradio as gr5import numpy as np6import random7import torch8import spaces9import os10import json11 12from PIL import Image13from diffusers import QwenImageEditPipeline, FlowMatchEulerDiscreteScheduler14 15from huggingface_hub import InferenceClient16import math17 18from optimization import optimize_pipeline_19from qwenimage.pipeline_qwen_image_edit import QwenImageEditPipeline as QwenImageEditPipelineCustom20from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel21from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA322 23# --- Prompt Enhancement using Hugging Face InferenceClient ---24def polish_prompt_hf(original_prompt, system_prompt):25 """26 Rewrites the prompt using a Hugging Face InferenceClient.27 """28 # Ensure HF_TOKEN is set29 api_key = os.environ.get("HF_TOKEN")30 if not api_key:31 print("Warning: HF_TOKEN not set. Falling back to original prompt.")32 return original_prompt33 34 try:35 # Initialize the client36 client = InferenceClient(37 provider="cerebras",38 api_key=api_key,39 )40 41 # Format the messages for the chat completions API42 messages = [43 {"role": "system", "content": system_prompt},44 {"role": "user", "content": original_prompt}45 ]46 47 # Call the API48 completion = client.chat.completions.create(49 model="Qwen/Qwen3-235B-A22B-Instruct-2507",50 messages=messages,51 )52 53 # Parse the response54 result = completion.choices[0].message.content55 56 # Try to extract JSON if present57 if '{"Rewritten"' in result:58 try:59 # Clean up the response60 result = result.replace('```json', '').replace('```', '')61 result_json = json.loads(result)62 polished_prompt = result_json.get('Rewritten', result)63 except:64 polished_prompt = result65 else:66 polished_prompt = result67 68 polished_prompt = polished_prompt.strip().replace("\n", " ")69 return polished_prompt70 71 except Exception as e:72 print(f"Error during API call to Hugging Face: {e}")73 # Fallback to original prompt if enhancement fails74 return original_prompt75 76 77def polish_prompt(prompt, img):78 """79 Main function to polish prompts for image editing using HF inference.80 """81 SYSTEM_PROMPT = '''82# Edit Instruction Rewriter83You are a professional edit instruction rewriter. Your task is to generate a precise, concise, and visually achievable professional-level edit instruction based on the user-provided instruction and the image to be edited. 84 85Please strictly follow the rewriting rules below:86 87## 1. General Principles88- Keep the rewritten prompt **concise**. Avoid overly long sentences and reduce unnecessary descriptive language. 89- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary. 90- Keep the core intention of the original instruction unchanged, only enhancing its clarity, rationality, and visual feasibility. 91- All added objects or modifications must align with the logic and style of the edited input image's overall scene. 92 93## 2. Task Type Handling Rules94### 1. Add, Delete, Replace Tasks95- If the instruction is clear (already includes task type, target entity, position, quantity, attributes), preserve the original intent and only refine the grammar. 96- If the description is vague, supplement with minimal but sufficient details (category, color, size, orientation, position, etc.). For example: 97 > Original: "Add an animal" 98 > Rewritten: "Add a light-gray cat in the bottom-right corner, sitting and facing the camera" 99- Remove meaningless instructions: e.g., "Add 0 objects" should be ignored or flagged as invalid. 100- For replacement tasks, specify "Replace Y with X" and briefly describe the key visual features of X. 101 102### 2. Text Editing Tasks103- All text content must be enclosed in English double quotes " ". Do not translate or alter the original language of the text, and do not change the capitalization. 104- **For text replacement tasks, always use the fixed template:**105 - Replace "xx" to "yy". 106 - Replace the xx bounding box to "yy". 107- If the user does not specify text content, infer and add concise text based on the instruction and the input image's context. For example: 108 > Original: "Add a line of text" (poster) 109 > Rewritten: "Add text "LIMITED EDITION" at the top center with slight shadow" 110- Specify text position, color, and layout in a concise way. 111 112### 3. Human Editing Tasks113- Maintain the person's core visual consistency (ethnicity, gender, age, hairstyle, expression, outfit, etc.). 114- If modifying appearance (e.g., clothes, hairstyle), ensure the new element is consistent with the original style. 115- **For expression changes, they must be natural and subtle, never exaggerated.** 116- If deletion is not specifically emphasized, the most important subject in the original image (e.g., a person, an animal) should be preserved.117 - For background change tasks, emphasize maintaining subject consistency at first. 118- Example: 119 > Original: "Change the person's hat" 120 > Rewritten: "Replace the man's hat with a dark brown beret; keep smile, short hair, and gray jacket unchanged" 121 122### 4. Style Transformation or Enhancement Tasks123- If a style is specified, describe it concisely with key visual traits. For example: 124 > Original: "Disco style" 125 > Rewritten: "1970s disco: flashing lights, disco ball, mirrored walls, colorful tones" 126- If the instruction says "use reference style" or "keep current style," analyze the input image, extract main features (color, composition, texture, lighting, art style), and integrate them concisely. 127- **For coloring tasks, including restoring old photos, always use the fixed template:** "Restore old photograph, remove scratches, reduce noise, enhance details, high resolution, realistic, natural skin tones, clear facial features, no distortion, vintage photo restoration" 128- If there are other changes, place the style description at the end.129 130## 3. Rationality and Logic Checks131- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" should be logically corrected. 132- Add missing key information: if position is unspecified, choose a reasonable area based on composition (near subject, empty space, center/edges). 133 134# Output Format135Return only the rewritten instruction text directly, without JSON formatting or any other wrapper.136'''137 138 # Note: We're not actually using the image in the HF version, 139 # but keeping the interface consistent140 full_prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"141 142 return polish_prompt_hf(full_prompt, SYSTEM_PROMPT)143 144 145# --- Model Loading ---146dtype = torch.bfloat16147device = "cuda" if torch.cuda.is_available() else "cpu"148 149# Scheduler configuration for Lightning150scheduler_config = {151 "base_image_seq_len": 256,152 "base_shift": math.log(3),153 "invert_sigmas": False,154 "max_image_seq_len": 8192,155 "max_shift": math.log(3),156 "num_train_timesteps": 1000,157 "shift": 1.0,158 "shift_terminal": None,159 "stochastic_sampling": False,160 "time_shift_type": "exponential",161 "use_beta_sigmas": False,162 "use_dynamic_shifting": True,163 "use_exponential_sigmas": False,164 "use_karras_sigmas": False,165}166 167# Initialize scheduler with Lightning config168scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)169 170# Load the edit pipeline with Lightning scheduler171pipe = QwenImageEditPipelineCustom.from_pretrained(172 "Qwen/Qwen-Image-Edit", 173 scheduler=scheduler,174 torch_dtype=dtype175).to(device)176 177# Load Lightning LoRA weights for acceleration178try:179 pipe.load_lora_weights(180 "lightx2v/Qwen-Image-Lightning", 181 weight_name="Qwen-Image-Lightning-8steps-V2.0.safetensors"182 )183 pipe.fuse_lora()184 print("Successfully loaded Lightning LoRA weights")185except Exception as e:186 print(f"Warning: Could not load Lightning LoRA weights: {e}")187 print("Continuing with base model...")188 189#spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/Qwen-Image", variant="fa3")190 191# Apply the same optimizations from the first version192#pipe.transformer.__class__ = QwenImageTransformer2DModel193#pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())194 195# --- Ahead-of-time compilation ---196#optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt="prompt")197 198# --- UI Constants and Helpers ---199MAX_SEED = np.iinfo(np.int32).max200 201# --- Main Inference Function ---202@spaces.GPU(duration=60)203def infer(204 image,205 prompt,206 seed=42,207 randomize_seed=False,208 true_guidance_scale=1.0,209 num_inference_steps=8, # Default to 8 steps for fast inference210 rewrite_prompt=True,211 progress=gr.Progress(track_tqdm=True),212):213 """214 Generates an edited image using the Qwen-Image-Edit pipeline with Lightning acceleration.215 """216 # Hardcode the negative prompt as in the original217 negative_prompt = " "218 219 if randomize_seed:220 seed = random.randint(0, MAX_SEED)221 222 # Set up the generator for reproducibility223 generator = torch.Generator(device=device).manual_seed(seed)224 225 print(f"Original prompt: '{prompt}'")226 print(f"Negative Prompt: '{negative_prompt}'")227 print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")228 229 if rewrite_prompt:230 prompt = polish_prompt(prompt, image)231 print(f"Rewritten Prompt: {prompt}")232 233 # Generate the edited image - always generate just 1 image234 try:235 images = pipe(236 image,237 prompt=prompt,238 negative_prompt=negative_prompt,239 num_inference_steps=num_inference_steps,240 generator=generator,241 true_cfg_scale=true_guidance_scale,242 num_images_per_prompt=1 # Always generate only 1 image243 ).images244 245 # Return the first (and only) image246 return images[0], seed247 248 except Exception as e:249 print(f"Error during inference: {e}")250 raise e251 252# --- Examples and UI Layout ---253examples = [254 # You can add example pairs of [image_path, prompt] here255 # ["path/to/image1.jpg", "Replace the background with a beach scene"],256 # ["path/to/image2.jpg", "Add a red hat to the person"],257]258 259css = """260#col-container {261 margin: 0 auto;262 max-width: 1024px;263}264#logo-title {265 text-align: center;266}267#logo-title img {268 width: 400px;269}270#edit_text{margin-top: -62px !important}271"""272 273with gr.Blocks(css=css) as demo:274 with gr.Column(elem_id="col-container"):275 gr.HTML("""276 <div id="logo-title">277 <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_edit_logo.png" alt="Qwen-Image Edit Logo" width="400" style="display: block; margin: 0 auto;">278 <h2 style="font-style: italic;color: #5b47d1;margin-top: -27px !important;margin-left: 96px">Fast, 8-steps with Lightning LoRA</h2>279 </div>280 """)281 gr.Markdown("""282 [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. 283 This demo uses the [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) LoRA for accelerated inference.284 Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit) to run locally with ComfyUI or diffusers.285 """)286 287 with gr.Row():288 with gr.Column():289 input_image = gr.Image(290 label="Input Image", 291 show_label=True, 292 type="pil"293 )294 # Changed from Gallery to Image295 result = gr.Image(296 label="Result", 297 show_label=True, 298 type="pil"299 )300 301 with gr.Row():302 prompt = gr.Text(303 label="Edit Instruction",304 show_label=False,305 placeholder="Describe the edit instruction (e.g., 'Replace the background with a sunset', 'Add a red hat', 'Remove the person')",306 container=False,307 )308 run_button = gr.Button("Edit!", variant="primary")309 310 with gr.Accordion("Advanced Settings", open=False):311 seed = gr.Slider(312 label="Seed",313 minimum=0,314 maximum=MAX_SEED,315 step=1,316 value=0,317 )318 319 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)320 321 with gr.Row():322 true_guidance_scale = gr.Slider(323 label="True guidance scale",324 minimum=1.0,325 maximum=10.0,326 step=0.1,327 value=1.0328 )329 330 num_inference_steps = gr.Slider(331 label="Number of inference steps",332 minimum=4,333 maximum=28,334 step=1,335 value=8336 )337 338 # Removed num_images_per_prompt slider entirely339 rewrite_prompt = gr.Checkbox(340 label="Enhance prompt (using HF Inference)", 341 value=True342 )343 344 # gr.Examples(examples=examples, inputs=[input_image, prompt], outputs=[result, seed], fn=infer, cache_examples=False)345 346 gr.on(347 triggers=[run_button.click, prompt.submit],348 fn=infer,349 inputs=[350 input_image,351 prompt,352 seed,353 randomize_seed,354 true_guidance_scale,355 num_inference_steps,356 rewrite_prompt,357 # Removed num_images_per_prompt from inputs358 ],359 outputs=[result, seed],360 )361 362if __name__ == "__main__":363 demo.launch()