import gradio as gr import numpy as np import spaces import torch import random import os import json # from diffusers import QwenImageEditInpaintPipeline # NOTE: `optimization.py` (AOTI compile path) imports torchao at module top and is # unused here (the optimize_pipeline_ call below is commented out). torchao is left # out of requirements entirely: the pinned diffusers branch crashes on import with # newer torchao (NameError in its torchao quantizer), and peft rejects the older # 0.11.0. With torchao absent, both diffusers and peft skip their torchao paths. from diffusers.utils import load_image from diffusers import FlowMatchEulerDiscreteScheduler from qwenimage.pipeline_qwenimage_edit_inpaint import QwenImageEditInpaintPipeline from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel import math from huggingface_hub import InferenceClient from PIL import Image, ImageOps # Set environment variable for parallel loading # os.environ["HF_ENABLE_PARALLEL_LOADING"] = "YES" # --- Prompt Enhancement using Hugging Face InferenceClient --- def polish_prompt_hf(original_prompt, system_prompt): """ Rewrites the prompt using a Hugging Face InferenceClient. """ # Ensure HF_TOKEN is set api_key = os.environ.get("HF_TOKEN") if not api_key: print("Warning: HF_TOKEN not set. Falling back to original prompt.") return original_prompt try: # Initialize the client client = InferenceClient( provider="cerebras", api_key=api_key, ) # Format the messages for the chat completions API messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": original_prompt} ] # Call the API completion = client.chat.completions.create( model="Qwen/Qwen3-235B-A22B-Instruct-2507", messages=messages, ) # Parse the response result = completion.choices[0].message.content # Try to extract JSON if present if '{"Rewritten"' in result: try: # Clean up the response result = result.replace('```json', '').replace('```', '') result_json = json.loads(result) polished_prompt = result_json.get('Rewritten', result) except: polished_prompt = result else: polished_prompt = result polished_prompt = polished_prompt.strip().replace("\n", " ") return polished_prompt except Exception as e: print(f"Error during API call to Hugging Face: {e}") # Fallback to original prompt if enhancement fails return original_prompt def polish_prompt(prompt, img): """ Main function to polish prompts for image editing using HF inference. """ SYSTEM_PROMPT = ''' # Edit Instruction Rewriter You 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. Please strictly follow the rewriting rules below: ## 1. General Principles - Keep the rewritten prompt **concise**. Avoid overly long sentences and reduce unnecessary descriptive language. - If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary. - Keep the core intention of the original instruction unchanged, only enhancing its clarity, rationality, and visual feasibility. - All added objects or modifications must align with the logic and style of the edited input image's overall scene. - If the instruction refers to multiple input images (e.g. "Picture 1", "Picture 2", a "reference image", or "the object from the second image"), preserve those references exactly so the relationship between the images stays clear. ## 2. Task Type Handling Rules ### 1. Add, Delete, Replace Tasks - If the instruction is clear (already includes task type, target entity, position, quantity, attributes), preserve the original intent and only refine the grammar. - If the description is vague, supplement with minimal but sufficient details (category, color, size, orientation, position, etc.). For example: > Original: "Add an animal" > Rewritten: "Add a light-gray cat in the bottom-right corner, sitting and facing the camera" - Remove meaningless instructions: e.g., "Add 0 objects" should be ignored or flagged as invalid. - For replacement tasks, specify "Replace Y with X" and briefly describe the key visual features of X. ### 2. Text Editing Tasks - 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. - **For text replacement tasks, always use the fixed template:** - Replace "xx" to "yy". - Replace the xx bounding box to "yy". - 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: > Original: "Add a line of text" (poster) > Rewritten: "Add text "LIMITED EDITION" at the top center with slight shadow" - Specify text position, color, and layout in a concise way. ### 3. Human Editing Tasks - Maintain the person's core visual consistency (ethnicity, gender, age, hairstyle, expression, outfit, etc.). - If modifying appearance (e.g., clothes, hairstyle), ensure the new element is consistent with the original style. - **For expression changes, they must be natural and subtle, never exaggerated.** - If deletion is not specifically emphasized, the most important subject in the original image (e.g., a person, an animal) should be preserved. - For background change tasks, emphasize maintaining subject consistency at first. - Example: > Original: "Change the person's hat" > Rewritten: "Replace the man's hat with a dark brown beret; keep smile, short hair, and gray jacket unchanged" ### 4. Style Transformation or Enhancement Tasks - If a style is specified, describe it concisely with key visual traits. For example: > Original: "Disco style" > Rewritten: "1970s disco: flashing lights, disco ball, mirrored walls, colorful tones" - 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. - **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" - If there are other changes, place the style description at the end. ## 3. Rationality and Logic Checks - Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" should be logically corrected. - Add missing key information: if position is unspecified, choose a reasonable area based on composition (near subject, empty space, center/edges). # Output Format Return only the rewritten instruction text directly, without JSON formatting or any other wrapper. ''' # The HF client already sends SYSTEM_PROMPT as the system message, so we pass # only the user's instruction as the user message. The old code concatenated the # entire SYSTEM_PROMPT into the user text; when the API errored (as the cerebras # provider currently does for this model), the fallback returned that giant blob # as the "prompt" -- which is the garbled "Rewritten Prompt" you saw in the logs. return polish_prompt_hf(prompt, SYSTEM_PROMPT) MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 2048 # --- Helper functions for reuse feature --- def clear_result(): """Clears the result image.""" return gr.update(value=None) def use_output_as_input(output_image): """Sets the generated output as the new input image.""" if output_image is not None: return gr.update(value=output_image[1]) return gr.update() # Initialize Qwen Image Edit pipeline # Scheduler configuration for Lightning scheduler_config = { "base_image_seq_len": 256, "base_shift": math.log(3), "invert_sigmas": False, "max_image_seq_len": 8192, "max_shift": math.log(3), "num_train_timesteps": 1000, "shift": 1.0, "shift_terminal": None, "stochastic_sampling": False, "time_shift_type": "exponential", "use_beta_sigmas": False, "use_dynamic_shifting": True, "use_exponential_sigmas": False, "use_karras_sigmas": False, } # Initialize scheduler with Lightning config scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config) pipe = QwenImageEditInpaintPipeline.from_pretrained("Qwen/Qwen-Image-Edit", scheduler=scheduler, torch_dtype=torch.bfloat16).to("cuda") pipe.load_lora_weights( "lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors" ) pipe.fuse_lora() # pipe.transformer.__class__ = QwenImageTransformer2DModel # NOTE: The FA3 attention processor (QwenDoubleStreamAttnProcessorFA3) was removed. # The vllm-flash-attn3 kernel fetched via `kernels` has no compiled kernel image for # the current ZeroGPU device, causing a runtime # "CUDA error: no kernel image is available for execution on the device" at inference. # Without an explicit processor the transformer falls back to native SDPA # (dispatch_attention_fn backend=None), which runs on any device. # dummy_mask = load_image("https://github.com/Trgtuan10/Image_storage/blob/main/mask_cat.png?raw=true") # # --- Ahead-of-time compilation --- # optimize_pipeline_(pipe, image=Image.new("RGB", (1328, 1328)), prompt="prompt", mask_image=dummy_mask) def composite_reference(base_image, mask_layer, reference_image): """Paste the reference image into the masked region of the base image. This inpaint pipeline's text encoder conditions on a SINGLE image, so we cannot feed it a second image directly (doing so triggers the "Image features and image tokens do not match" error). Instead we paste the reference into the area you masked and use the result as the init image for inpainting. With a moderate `strength` (~0.6-0.85) the inpaint pass then blends the pasted content into the scene (lighting, edges, perspective). The reference is scaled to fit *inside* the mask's bounding box (aspect ratio preserved) and centered. """ base = base_image.convert("RGB") # Derive a binary mask from the painted layer. The ImageEditor brush paints # opaque pixels, so non-zero alpha (or non-zero luminance) marks the region. if mask_layer.mode == "RGBA": binary = np.array(mask_layer.split()[-1]) > 0 else: binary = np.array(mask_layer.convert("L")) > 0 ys, xs = np.where(binary) if xs.size == 0: # Nothing was painted -> leave the base image unchanged. return base x0, x1 = int(xs.min()), int(xs.max()) y0, y1 = int(ys.min()), int(ys.max()) box_w, box_h = (x1 - x0 + 1), (y1 - y0 + 1) ref_fitted = ImageOps.contain(reference_image.convert("RGB"), (box_w, box_h)) paste_x = x0 + (box_w - ref_fitted.width) // 2 paste_y = y0 + (box_h - ref_fitted.height) // 2 composited = base.copy() composited.paste(ref_fitted, (paste_x, paste_y)) return composited @spaces.GPU(duration=120) def infer(edit_images, reference_image, prompt, negative_prompt="", seed=42, randomize_seed=False, strength=0.85, num_inference_steps=8, true_cfg_scale=1.0, rewrite_prompt=True, progress=gr.Progress(track_tqdm=True)): # `edit_images` is the base image you draw the mask on. The mask marks the # region that will be regenerated. image = edit_images["background"].convert("RGB") mask = edit_images["layers"][0] # If a reference image is supplied, paste it into the masked region so it # becomes the starting point for inpainting. We keep this to a SINGLE image # because the pipeline's text encoder only supports one. if reference_image is not None: image = composite_reference(image, mask, reference_image) if randomize_seed: seed = random.randint(0, MAX_SEED) if rewrite_prompt: prompt = polish_prompt(prompt, image) print(f"Rewritten Prompt: {prompt}") # Generate image using Qwen pipeline result_image = pipe( prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask, strength=strength, num_inference_steps=num_inference_steps, true_cfg_scale=true_cfg_scale, generator=torch.Generator(device="cuda").manual_seed(seed) ).images[0] return [image,result_image], seed examples = [ "a vase of flowers on the table, matching the scene's lighting", "the person standing here, natural shadows and perspective", "this product on the shelf, blended seamlessly", ] css = """ #col-container { margin: 0 auto; max-width: 1024px; } #logo-title { text-align: center; } #logo-title img { width: 400px; } #edit_text{margin-top: -62px !important} """ with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.HTML("""
Qwen-Image Edit Logo

Inpaint

""") gr.Markdown(""" Reference-guided inpainting with Qwen Image Edit. [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. **How to use:** draw a mask on the **base image** to mark the region to replace, upload a **reference image** holding the thing you want placed there, then describe it in the prompt (e.g. *"a coffee mug on the table, blended naturally with matching lighting"*). The reference is scaled to fit inside your masked area and pasted in, then the model blends it. Use **Strength ~0.6–0.85** so the reference survives the inpaint pass — near 1.0 the model mostly ignores it. The reference image is optional; leave it empty for plain text-driven inpainting (then 1.0 is fine). This demo uses the [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) LoRA for accelerated 8-step inference. 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. """) with gr.Row(): with gr.Column(): edit_image = gr.ImageEditor( label='Base image — draw a mask over the region to inpaint', type='pil', sources=["upload", "webcam"], image_mode='RGB', layers=False, brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"), height=600 ) reference_image = gr.Image( label='Reference image — the content to fit into the masked area (optional)', type='pil', sources=["upload", "webcam"], image_mode='RGB', height=300 ) prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Describe what to put in the mask (e.g., 'place the object from the reference image here')", container=False, ) negative_prompt = gr.Text( label="Negative Prompt", show_label=True, max_lines=1, placeholder="Enter what you don't want (optional)", container=False, value="", visible=False ) run_button = gr.Button("Run") with gr.Column(): result = gr.ImageSlider(label="Result", show_label=False, interactive=False) use_as_input_button = gr.Button("🔄 Use as Input Image", visible=False, variant="secondary") with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): strength = gr.Slider( label="Strength", minimum=0.0, maximum=1.0, step=0.05, value=0.85, info="How much the masked region changes. With a reference image, lower (~0.6-0.85) keeps it faithful to the reference; near 1.0 mostly ignores the paste. Without a reference, use 1.0 for a full replacement." ) true_cfg_scale = gr.Slider( label="True CFG Scale", minimum=1.0, maximum=10.0, step=0.5, value=1.0, info="Classifier-free guidance scale" ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=8, ) rewrite_prompt = gr.Checkbox( label="Enhance prompt (using HF Inference)", value=True ) # Event handlers for reuse functionality use_as_input_button.click( fn=use_output_as_input, inputs=[result], outputs=[edit_image], show_api=False ) # Main generation pipeline with result clearing and button visibility gr.on( triggers=[run_button.click, prompt.submit], fn=clear_result, inputs=None, outputs=result, show_api=False ).then( fn = infer, inputs = [edit_image, reference_image, prompt, negative_prompt, seed, randomize_seed, strength, num_inference_steps, true_cfg_scale, rewrite_prompt], outputs = [result, seed] ).then( fn=lambda: gr.update(visible=True), inputs=None, outputs=use_as_input_button, show_api=False ) demo.launch()