app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
import os
import json

from PIL import Image, ImageDraw
import torch 
import math

from optimization import optimize_pipeline_
from qwenimage.pipeline_qwen_image_edit import QwenImageEditPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

from huggingface_hub import InferenceClient
import math

# --- 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.  

## 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.
'''
    
    # Note: We're not actually using the image in the HF version, 
    # but keeping the interface consistent
    full_prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
    
    return polish_prompt_hf(full_prompt, SYSTEM_PROMPT)

# --- Outpainting Functions ---
def can_expand(source_width, source_height, target_width, target_height, alignment):
    """Checks if the image can be expanded based on the alignment."""
    if alignment in ("Left", "Right") and source_width >= target_width:
        return False
    if alignment in ("Top", "Bottom") and source_height >= target_height:
        return False
    return True

def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
    """Prepares the image with white margins and creates a mask for outpainting."""
    target_size = (width, height)
    
    # Calculate the scaling factor to fit the image within the target size
    scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)
    new_width = int(image.width * scale_factor)
    new_height = int(image.height * scale_factor)
    
    # Resize the source image to fit within target size
    source = image.resize((new_width, new_height), Image.LANCZOS)
    
    # Apply resize option using percentages
    if resize_option == "Full":
        resize_percentage = 100
    elif resize_option == "50%":
        resize_percentage = 50
    elif resize_option == "33%":
        resize_percentage = 33
    elif resize_option == "25%":
        resize_percentage = 25
    else:  # Custom
        resize_percentage = custom_resize_percentage
    
    # Calculate new dimensions based on percentage
    resize_factor = resize_percentage / 100
    new_width = int(source.width * resize_factor)
    new_height = int(source.height * resize_factor)
    
    # Ensure minimum size of 64 pixels
    new_width = max(new_width, 64)
    new_height = max(new_height, 64)
    
    # Resize the image
    source = source.resize((new_width, new_height), Image.LANCZOS)
    
    # Calculate the overlap in pixels based on the percentage
    overlap_x = int(new_width * (overlap_percentage / 100))
    overlap_y = int(new_height * (overlap_percentage / 100))
    
    # Ensure minimum overlap of 1 pixel
    overlap_x = max(overlap_x, 1)
    overlap_y = max(overlap_y, 1)
    
    # Calculate margins based on alignment
    if alignment == "Middle":
        margin_x = (target_size[0] - new_width) // 2
        margin_y = (target_size[1] - new_height) // 2
    elif alignment == "Left":
        margin_x = 0
        margin_y = (target_size[1] - new_height) // 2
    elif alignment == "Right":
        margin_x = target_size[0] - new_width
        margin_y = (target_size[1] - new_height) // 2
    elif alignment == "Top":
        margin_x = (target_size[0] - new_width) // 2
        margin_y = 0
    elif alignment == "Bottom":
        margin_x = (target_size[0] - new_width) // 2
        margin_y = target_size[1] - new_height
    
    # Adjust margins to eliminate gaps
    margin_x = max(0, min(margin_x, target_size[0] - new_width))
    margin_y = max(0, min(margin_y, target_size[1] - new_height))
    
    # Create a new background image with white margins and paste the resized source image
    background = Image.new('RGB', target_size, (255, 255, 255))
    background.paste(source, (margin_x, margin_y))
    
    # Create the mask
    mask = Image.new('L', target_size, 255)
    mask_draw = ImageDraw.Draw(mask)
    
    # Calculate overlap areas
    white_gaps_patch = 2
    left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch
    right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch
    top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch
    bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch
    
    if alignment == "Left":
        left_overlap = margin_x + overlap_x if overlap_left else margin_x
    elif alignment == "Right":
        right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width
    elif alignment == "Top":
        top_overlap = margin_y + overlap_y if overlap_top else margin_y
    elif alignment == "Bottom":
        bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
    
    # Draw the mask
    mask_draw.rectangle([
        (left_overlap, top_overlap),
        (right_overlap, bottom_overlap)
    ], fill=0)
    
    return background, mask

def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
    """Creates a preview showing the mask overlay."""
    background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
    
    # Create a preview image showing the mask
    preview = background.copy().convert('RGBA')
    
    # Create a semi-transparent red overlay
    red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64))  # Reduced alpha to 64 (25% opacity)
    
    # Convert black pixels in the mask to semi-transparent red
    red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))
    red_mask.paste(red_overlay, (0, 0), mask)
    
    # Overlay the red mask on the background
    preview = Image.alpha_composite(preview, red_mask)
    
    return preview

# --- Model Loading ---
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)
pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

# --- Ahead-of-time compilation ---
optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt="prompt")

# --- UI Constants and Helpers ---
MAX_SEED = np.iinfo(np.int32).max

def preload_presets(target_ratio, ui_width, ui_height):
    """Updates the width and height sliders based on the selected aspect ratio."""
    if target_ratio == "9:16":
        changed_width = 720
        changed_height = 1280
        return changed_width, changed_height, gr.update()
    elif target_ratio == "16:9":
        changed_width = 1280
        changed_height = 720
        return changed_width, changed_height, gr.update()
    elif target_ratio == "1:1":
        changed_width = 1024
        changed_height = 1024
        return changed_width, changed_height, gr.update()
    elif target_ratio == "Custom":
        return ui_width, ui_height, gr.update(open=True)

def select_the_right_preset(user_width, user_height):
    if user_width == 720 and user_height == 1280:
        return "9:16"
    elif user_width == 1280 and user_height == 720:
        return "16:9"
    elif user_width == 1024 and user_height == 1024:
        return "1:1"
    else:
        return "Custom"

def toggle_custom_resize_slider(resize_option):
    return gr.update(visible=(resize_option == "Custom"))

# --- Main Inference Function (with outpainting preprocessing) ---
@spaces.GPU(duration=120)
def infer(
    image,
    prompt,
    width,
    height,
    overlap_percentage,
    resize_option,
    custom_resize_percentage,
    alignment,
    overlap_left,
    overlap_right,
    overlap_top,
    overlap_bottom,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=4.0,
    num_inference_steps=50,
    rewrite_prompt=True,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generates an outpainted image using the Qwen-Image-Edit pipeline.
    """
    # Hardcode the negative prompt as requested
    negative_prompt = " "
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Set up the generator for reproducibility
    generator = torch.Generator(device=device).manual_seed(seed)
    
    print(f"Original Prompt: '{prompt}'")
    print(f"Negative Prompt: '{negative_prompt}'")
    print(f"Seed: {seed}, Steps: {num_inference_steps}")

    if rewrite_prompt:
        prompt = polish_prompt(prompt, image)
        print(f"Rewritten Prompt: {prompt}")

    # Prepare the image with white margins for outpainting
    outpaint_image, mask = prepare_image_and_mask(
        image, width, height, overlap_percentage,
        resize_option, custom_resize_percentage, alignment,
        overlap_left, overlap_right, overlap_top, overlap_bottom
    )
    
    # Check if expansion is possible
    if not can_expand(image.width, image.height, width, height, alignment):
        alignment = "Middle"
        outpaint_image, mask = prepare_image_and_mask(
            image, width, height, overlap_percentage,
            resize_option, custom_resize_percentage, "Middle",
            overlap_left, overlap_right, overlap_top, overlap_bottom
        )
    
    print(f"Outpaint dimensions: {outpaint_image.size}")

    # Generate the image with outpainting preprocessing
    result_image = pipe(
        outpaint_image,  # Use the preprocessed image with white margins
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
    ).images[0]

    return result_image, seed

# --- Examples and UI Layout ---
examples = []

css = """
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
#edit_text{margin-top: -62px !important}
.preview-container {
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 10px;
    margin-top: 10px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML("""
        <div id="logo-title">
            <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;">
        </div>
        """)
        gr.Markdown("""
        ## Qwen-Image Edit with Outpainting
        
        Extend your images beyond their original boundaries with intelligent outpainting. The model will generate new content that seamlessly blends with your original image.
        
        [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. 
        Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit) to run locally.
        """)
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="pil")
                
                prompt = gr.Text(
                    label="Prompt",
                    placeholder="Describe what should appear in the extended areas",
                    container=False,
                )
                
                with gr.Row():
                    target_ratio = gr.Radio(
                        label="Target Ratio",
                        choices=["9:16", "16:9", "1:1", "Custom"],
                        value="16:9",
                        scale=2
                    )
                    alignment_dropdown = gr.Dropdown(
                        choices=["Middle", "Left", "Right", "Top", "Bottom"],
                        value="Middle",
                        label="Alignment"
                    )
                
                run_button = gr.Button("Outpaint!", variant="primary")
                
                with gr.Accordion("Outpainting Settings", open=False) as settings_panel:
                    with gr.Row():
                        width_slider = gr.Slider(
                            label="Target Width",
                            minimum=512,
                            maximum=2048,
                            step=8,
                            value=1280,
                        )
                        height_slider = gr.Slider(
                            label="Target Height",
                            minimum=512,
                            maximum=2048,
                            step=8,
                            value=720,
                        )
                    
                    with gr.Group():
                        overlap_percentage = gr.Slider(
                            label="Mask overlap (%)",
                            minimum=1,
                            maximum=50,
                            value=10,
                            step=1,
                            info="Controls the blending area between original and new content"
                        )
                        
                        with gr.Row():
                            overlap_top = gr.Checkbox(label="Overlap Top", value=True)
                            overlap_right = gr.Checkbox(label="Overlap Right", value=True)
                        with gr.Row():
                            overlap_left = gr.Checkbox(label="Overlap Left", value=True)
                            overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)
                    
                    with gr.Row():
                        resize_option = gr.Radio(
                            label="Resize input image",
                            choices=["Full", "50%", "33%", "25%", "Custom"],
                            value="Full",
                            info="How much of the target canvas the original image should occupy"
                        )
                        custom_resize_percentage = gr.Slider(
                            label="Custom resize (%)",
                            minimum=1,
                            maximum=100,
                            step=1,
                            value=50,
                            visible=False
                        )
                    
                    preview_button = gr.Button("Preview alignment and mask", variant="secondary")

                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="True guidance scale",
                            minimum=1.0,
                            maximum=10.0,
                            step=0.1,
                            value=1.0
                        )

                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=50,
                            step=1,
                            value=50,
                        )

                    rewrite_prompt = gr.Checkbox(
                        label="Enhance prompt (using HF Inference)", 
                        value=True
                    )

            with gr.Column():
                result = gr.Image(label="Result", type="pil")
                
                with gr.Column(visible=False) as preview_container:
                    preview_image = gr.Image(label="Preview (red area will be generated)", type="pil")

    # Event handlers
    target_ratio.change(
        fn=preload_presets,
        inputs=[target_ratio, width_slider, height_slider],
        outputs=[width_slider, height_slider, settings_panel],
        queue=False,
    )
    
    width_slider.change(
        fn=select_the_right_preset,
        inputs=[width_slider, height_slider],
        outputs=[target_ratio],
        queue=False,
    )
    
    height_slider.change(
        fn=select_the_right_preset,
        inputs=[width_slider, height_slider],
        outputs=[target_ratio],
        queue=False,
    )
    
    resize_option.change(
        fn=toggle_custom_resize_slider,
        inputs=[resize_option],
        outputs=[custom_resize_percentage],
        queue=False,
    )
    
    preview_button.click(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=[preview_container],
        queue=False,
    ).then(
        fn=preview_image_and_mask,
        inputs=[
            input_image, width_slider, height_slider, overlap_percentage,
            resize_option, custom_resize_percentage, alignment_dropdown,
            overlap_left, overlap_right, overlap_top, overlap_bottom
        ],
        outputs=preview_image,
        queue=False,
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            input_image,
            prompt,
            width_slider,
            height_slider,
            overlap_percentage,
            resize_option,
            custom_resize_percentage,
            alignment_dropdown,
            overlap_left,
            overlap_right,
            overlap_top,
            overlap_bottom,
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            rewrite_prompt,
        ],
        outputs=[result, seed],
    )

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