app.py

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

from PIL import Image
from diffusers import QwenImageEditPipeline, FlowMatchEulerDiscreteScheduler
from huggingface_hub import InferenceClient
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

import base64
from io import BytesIO

# --- Prompt Enhancement using Hugging Face InferenceClient ---
def polish_prompt_hf(original_prompt, system_prompt, img):
    """
    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="nebius",
            api_key=api_key,
        )

                # Convert PIL Image to base64 data URL
        image_url = None
        if img is not None:
            # If img is a PIL Image
            if hasattr(img, 'save'):  # Check if it's a PIL Image
                buffered = BytesIO()
                img.save(buffered, format="PNG")
                img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
                image_url = f"data:image/png;base64,{img_base64}"
            # If img is already a file path (string)
            elif isinstance(img, str):
                with open(img, "rb") as image_file:
                    img_base64 = base64.b64encode(image_file.read()).decode('utf-8')
                image_url = f"data:image/png;base64,{img_base64}"
            else:
                print(f"Warning: Unexpected image type: {type(img)}")
                return original_prompt

        # Format the messages for the chat completions API
        messages = [
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": original_prompt
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_url
                        }
                    }
                ]
            }
        ]


        # Call the API
        completion = client.chat.completions.create(
            model="Qwen/Qwen2.5-VL-72B-Instruct",
            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 = '''
# Lighting Edit Instruction Rewriter

You are a professional lighting edit instruction rewriter. Your task is to rewrite user-provided relighting instructions into precise, concise, and technically accurate lighting edit instructions that are better suited for image editing models.

Please strictly follow the rewriting rules below:

## 1. General Principles
- **Rewrite the input instruction** to be **concise and technically specific**. Use professional lighting terminology.
- If the original instruction is contradictory, vague, or technically unfeasible, rewrite it to prioritize physically realistic lighting corrections.
- Preserve the core intention of the original instruction while enhancing technical accuracy and visual feasibility.
- All lighting modifications must maintain realistic physics and natural light behavior.
- **Preserve subject integrity**: Keep facial features, clothing, pose, and other non-lighting elements unchanged unless specifically requested in the original instruction.

## 2. Lighting Task Categories

### 1. Light Direction and Positioning
- **Specify precise direction**: front-lit, back-lit, side-lit (left/right), top-lit, bottom-lit, three-quarter lighting
- **Include angle details**: 45-degree side lighting, overhead lighting, low-angle dramatic lighting
- **For vague instructions like "better lighting"**: analyze current lighting issues and specify improvement (e.g., "Add soft front lighting to reduce harsh shadows on face")

### 2. Light Quality and Characteristics
- **Hard vs. Soft**: "hard directional lighting with sharp shadows" vs. "soft diffused lighting with gentle shadows"
- **Intensity**: bright, moderate, dim, dramatic high-contrast, subtle low-contrast
- **Coverage**: full illumination, selective lighting, spotlight effect, rim lighting, fill lighting

### 3. Color Temperature and Mood
- **Temperature specification**: warm (3000K-3500K), neutral (4000K-5000K), cool (5500K-6500K), daylight (6500K+)
- **Mood descriptors**: golden hour warmth, clinical cool lighting, cozy warm ambiance, dramatic cool shadows
- **Mixed lighting**: "warm key light with cool rim lighting," "daylight from window with warm interior lighting"

### 4. Environmental and Context-Specific Lighting
- **Time of day**: morning soft light, midday harsh sun, golden hour, blue hour, night artificial lighting
- **Location-based**: studio lighting setup, natural outdoor lighting, indoor ambient lighting, street lighting
- **Weather conditions**: overcast soft lighting, direct sunlight, sunset glow, stormy dramatic lighting

### 5. Technical Lighting Setups
- **Professional terminology**: key light, fill light, rim/hair light, background light, bounce lighting
- **Studio setups**: Rembrandt lighting, butterfly lighting, split lighting, loop lighting
- **Multiple sources**: "main soft box from camera right, fill light from left, rim light from behind"

## 3. Instruction Rewriting Examples

### For Basic Lighting Changes:
- **Input**: "Make it brighter" → **Rewritten**: "Increase overall lighting with soft front illumination, maintain natural shadows"
- **Input**: "Dramatic lighting" → **Rewritten**: "Add strong side lighting from camera left with deep shadows on right side, high contrast"

### For Direction Changes:
- **Input**: "Light from behind" → **Rewritten**: "Add rim lighting from behind subject, maintain visibility of facial features with subtle fill light"
- **Input**: "Window lighting" → **Rewritten**: "Natural daylight from camera left, soft directional lighting mimicking window light"

### For Mood/Atmosphere:
- **Input**: "Warmer lighting" → **Rewritten**: "Adjust to warm 3200K lighting, golden tone, soft shadows"
- **Input**: "Studio lighting" → **Rewritten**: "Professional three-point lighting: soft key light camera right, fill light camera left, rim light from behind"

## 4. Technical Considerations and Constraints

### Physical Accuracy:
- Ensure shadow directions match light source positions
- Maintain consistent color temperature across the scene
- Respect surface materials (how light interacts with skin, fabric, metal, etc.)
- Consider ambient light contribution and bounce lighting
 

### Preservation Rules:
- **Always specify**: "maintain facial features unchanged," "preserve original pose and expression"
- **For portraits**: "keep skin texture and facial structure identical, only adjust lighting"
- **For scenes**: "preserve all objects and composition, modify lighting only"

### Quality Standards:
- **Include resolution/quality terms**: "realistic lighting physics," "natural light falloff," "smooth gradients"
- **Avoid artifacts**: "no harsh light cutoffs," "natural shadow transitions," "realistic highlight rolloff"

## 5. Common Lighting Scenarios

### Portrait Relighting:
"Apply soft key lighting from camera right at 45-degree angle, add gentle fill light from left to reduce shadow contrast, maintain natural skin tones and facial features"

### Scene Relighting:
"Change to golden hour lighting: warm 3000K directional light from camera right, long soft shadows, enhanced ambient warm bounce light"

### Dramatic Relighting:
"High-contrast lighting setup: strong key light from camera left, minimal fill light, deep shadows on right side, dramatic mood while preserving subject clarity"

### Natural Environment:
"Simulate overcast daylight: soft diffused lighting from above, minimal shadows, cool 6000K color temperature, even illumination across scene"

## 6. Error Prevention
- Never specify impossible lighting (e.g., "shadows pointing toward light source")
- Always include both light addition and shadow consideration
- Specify color temperature changes when requesting "warm" or "cool" lighting
# 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, img)
# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# 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 = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", scheduler=scheduler,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

# Illumination options mapping
ILLUMINATION_OPTIONS = {
# Natural Daylight
    "natural lighting": "Neutral white color temperature with balanced exposure and soft shadows",
    "sunshine from window": "Bright directional sunlight with hard shadows and visible light rays",
    "golden time": "Warm golden hour lighting with enhanced warm colors and soft shadows",
    "sunrise in the mountains": "Warm backlighting with atmospheric haze and lens flare",
    "afternoon light filtering through trees": "Dappled sunlight patterns with green color cast from foliage",
    "early morning rays, forest clearing": "God rays through trees with warm color temperature",
    "golden sunlight streaming through trees": "Golden god rays with atmospheric particles in light beams",
    
    # Sunset & Evening
    "sunset over sea": "Warm sunset light with soft diffused lighting and gentle gradients",
    "golden hour in a meadow": "Golden backlighting with lens flare and rim lighting",
    "golden hour on a city skyline": "Golden lighting on buildings with silhouette effects",
    "evening glow in the desert": "Warm directional lighting with long shadows",
    "dusky evening on a beach": "Cool backlighting with horizon silhouettes",
    "mellow evening glow on a lake": "Warm lighting with water reflections",
    "warm sunset in a rural village": "Golden hour lighting with peaceful warm tones",
    
    # Night & Moonlight
    "moonlight through curtains": "Cool blue lighting with curtain shadow patterns",
    "moonlight in a dark alley": "Cool blue lighting with deep urban shadows",
    "midnight in the forest": "Very low brightness with minimal ambient lighting",
    "midnight sky with bright starlight": "Cool blue lighting with star point sources",
    "fireflies lighting up a summer night": "Small glowing points with warm ambient lighting",
    
    # Indoor & Cozy
    "warm atmosphere, at home, bedroom": "Very warm lighting with soft diffused glow",
    "home atmosphere, cozy bedroom illumination": "Warm table lamp lighting with pools of light",
    "cozy candlelight": "Warm orange flickering light with dramatic shadows",
    "candle-lit room, rustic vibe": "Multiple warm candlelight sources with atmospheric shadows",
    "night, cozy warm light from fireplace": "Warm orange-red firelight with flickering effects",
    "campfire light": "Warm orange flickering light from below with dancing shadows",
    
    # Urban & Neon
    "neon night, city": "Vibrant blue, magenta, and green neon lights with reflections",
    "blue neon light, urban street": "Blue neon lighting with urban glow effects",
    "neon, Wong Kar-wai, warm": "Warm amber and red neon with moody selective lighting",
    "red and blue police lights in rain": "Alternating red and blue strobing with wet reflections",
    "red glow, emergency lights": "Red emergency lighting with harsh shadows and high contrast",

     # Sci-Fi & Fantasy
    "sci-fi RGB glowing, cyberpunk": "Electric blue, pink, and green RGB lighting with glowing effects",
    "rainbow reflections, neon": "Chromatic rainbow patterns with prismatic reflections",
    "magic lit": "Colored rim lighting in purple and blue with soft ethereal glow",
    "mystical glow, enchanted forest": "Supernatural green and blue glowing with floating particles",
    "ethereal glow, magical forest": "Supernatural lighting with blue-green rim lighting",
    "underwater glow, deep sea": "Blue-green lighting with caustic patterns and particles",
    "underwater luminescence": "Blue-green bioluminescent glow with caustic light patterns",
    "aurora borealis glow, arctic landscape": "Green and purple dancing sky lighting",
    "crystal reflections in a cave": "Sparkle effects with prismatic light dispersion",
    
    # Weather & Atmosphere
    "foggy forest at dawn": "Volumetric fog with cool god rays through trees",
    "foggy morning, muted light": "Soft fog effects with reduced contrast throughout",
    "soft, diffused foggy glow": "Heavy fog with soft lighting and no harsh shadows",
    "stormy sky lighting": "Dramatic lighting with high contrast and rim lighting",
    "lightning flash in storm": "Brief intense white light with stark shadows",
    "rain-soaked reflections in city lights": "Wet surface reflections with streaking light effects",
    "gentle snowfall at dusk": "Cool blue lighting with snowflake particle effects",
    "hazy light of a winter morning": "Neutral lighting with atmospheric haze",
    "mysterious twilight, heavy mist": "Heavy fog with cool lighting and atmospheric depth",
    
    # Seasonal & Nature
    "vibrant autumn lighting in a forest": "Enhanced warm autumn colors with dappled sunlight",
    "purple and pink hues at twilight": "Warm lighting with soft purple and pink color grading",
    "desert sunset with mirage-like glow": "Warm orange lighting with heat distortion effects",
    "sunrise through foggy mountains": "Warm lighting through mist with atmospheric perspective",
    
    # Professional & Studio
    "soft studio lighting": "Multiple diffused sources with even illumination and minimal shadows",
    "harsh, industrial lighting": "Bright fluorescent lighting with hard shadows",
    "fluorescent office lighting": "Cool white overhead lighting with slight green tint",
    "harsh spotlight in dark room": "Single intense directional light with dramatic shadows",
    
    # Special Effects & Drama
    "light and shadow": "Maximum contrast with sharp shadow boundaries",
    "shadow from window": "Window frame shadow patterns with geometric shapes",
    "apocalyptic, smoky atmosphere": "Orange-red fire tint with smoke effects",
    "evil, gothic, in a cave": "Low brightness with cool lighting and deep shadows",
    "flickering light in a haunted house": "Unstable flickering with cool and warm mixed lighting",
    "golden beams piercing through storm clouds": "Dramatic god rays with high contrast",
    "dim candlelight in a gothic castle": "Warm orange candlelight with stone texture enhancement",

    # Festival & Celebration
    "colorful lantern light at festival": "Multiple colored lantern sources with bokeh effects",
    "golden glow at a fairground": "Warm carnival lighting with colorful bulb effects",
    "soft glow through stained glass": "Colored light filtering with rainbow surface patterns",
    "glowing embers from a forge": "Orange-red glowing particles with intense heat effects"

    }

# Lighting direction options
DIRECTION_OPTIONS = {
    "auto": "",  
    "left side": "Position the light source from the left side of the frame, creating shadows falling to the right.",
    "right side": "Position the light source from the right side of the frame, creating shadows falling to the left.",
    "top": "Position the light source from directly above, creating downward shadows.",
    "top left": "Position the light source from the top left corner, creating diagonal shadows falling down and to the right.",
    "top right": "Position the light source from the top right corner, creating diagonal shadows falling down and to the left.",
    "bottom": "Position the light source from below, creating upward shadows and dramatic under-lighting.",
    "front": "Position the light source from the front, minimizing shadows and creating even illumination.",
    "back": "Position the light source from behind the subject, creating silhouette effects and rim lighting."
}


# --- Main Inference Function ---
@spaces.GPU(duration=60)
def infer(
    image,
    prompt,
    illumination_dropdown, direction_dropdown, 
    seed=42,
    randomize_seed=True,
    true_guidance_scale=1.0,
    num_inference_steps=8,  # Default to 8 steps for fast inference
    rewrite_prompt=True,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generates an edited image using the Qwen-Image-Edit pipeline with Lightning acceleration.
    """
    # Hardcode the negative prompt as in the original
    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}, Guidance: {true_guidance_scale}")
    
    #If the dropdown isn't custom, and the user didn't specify a prompt, fill the prompt with the correct one from the illumination options
    if illumination_dropdown != "custom" and (prompt == "" or prompt == ILLUMINATION_OPTIONS[illumination_dropdown]):
        prompt = f"change the lighting. add {ILLUMINATION_OPTIONS[illumination_dropdown]}"


    # If direction isn't auto, add the direction suffix
    if direction_dropdown != "auto":
        prompt_with_template = prompt+ f" coming from the {direction_dropdown}"
    else:
        prompt_with_template= prompt
    
    if rewrite_prompt:
        final_prompt = polish_prompt(prompt_with_template, image)
    else:
        final_prompt = prompt_with_template

    print(f"Calling pipeline with prompt: '{final_prompt}'")
    

    # Generate the edited image - always generate just 1 image
    try:
        images = pipe(
            image,
            prompt=final_prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            generator=generator,
            true_cfg_scale=true_guidance_scale,
            num_images_per_prompt=1  # Always generate only 1 image
        ).images
        
        # Return the first (and only) image
        return [image,images[0]], seed, final_prompt
        
    except Exception as e:
        print(f"Error during inference: {e}")
        raise e

def update_prompt_from_dropdown(illumination_option):
    """Update the prompt textbox based on dropdown selection"""
    if illumination_option == "custom":
        return ""  # Clear the prompt for custom input
    else:
        return ILLUMINATION_OPTIONS[illumination_option]

# --- Examples and UI Layout ---
examples = [
    # You can add example pairs of [image_path, prompt] here
    # ["path/to/image1.jpg", "Replace the background with a beach scene"],
    # ["path/to/image2.jpg", "Add a red hat to the person"],
]

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("""
        <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;">
            <h2 style="font-style: italic;color: #5b47d1;margin-top: -27px !important;margin-left: 133px;">Relight [Fast]</h2>
        </div>
        """)
        gr.Markdown("""
        
        Relight images with Qwen Image Edit. [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. 
        
        This demo uses the [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) LoRA with AoT compilation 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():
                input_image = gr.Image(
                    label="Input Image", 
                    show_label=True, 
                    type="pil"
                )
                with gr.Row():
                    
                    illumination_dropdown = gr.Dropdown(
                        choices=["custom"] + list(ILLUMINATION_OPTIONS.keys()),
                        value="sunshine from window",
                        label="Choose Lighting Style",
                        scale=2
                    )
                    
                    direction_dropdown = gr.Dropdown(
                        choices=list(DIRECTION_OPTIONS.keys()),
                        value="auto",
                        label="Light Direction",
                        scale=1
                    )
                prompt = gr.Text(
                    label="Edit Instruction",
                    show_label=False,
                    placeholder="describe the edit instruction",
                    container=False,
                )
                run_button = gr.Button("Edit!", variant="primary")
                
            with gr.Column():
                result = gr.ImageSlider(
                    label="Result", 
                    show_label=True, 
                    type="pil"
                )
                final_prompt = gr.Textbox(label="Processed prompt", visible=False)
                   

        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=4,
                    maximum=28,
                    step=1,
                    value=8
                )
                
            # Removed num_images_per_prompt slider entirely
            rewrite_prompt = gr.Checkbox(
                label="Enhance prompt (using HF Inference)", 
                value=True
            )

        
        gr.Examples(
                examples=[
                    ["./assets/pexels-creationhill-1681010.jpg", "Add multiple colored light sources from lanterns. Create warm festival lighting. Set varied color temperatures. Add bokeh effects.", "colorful lantern light at festival", "auto"],
                    ["./assets/pexels-creationhill-1681010.jpg",  "add futuristic RGB lighting with electric blues, hot pinks, and neon greens creating a high-tech atmosphere with dramatic color separation and glowing effects", "sci-fi RGB glowing, cyberpunk", "left side"],
                    ["./assets/pexels-moose-photos-170195-1587009.jpg",  "Set blue-green color temperature. Add volumetric lighting effects. Reduce red channel significantly. Create particle effects in light beams. Add caustic light patterns.", "underwater glow, deep sea", "top"],
                    ["./assets/pexels-moose-photos-170195-1587009.jpg", "Replace lighting with red sources. Add flashing strobing effects. Increase contrast. Create harsh shadows. Set monochromatic red color scheme.", "red glow, emergency lights", "right side"],
                    ["./assets/pexels-simon-robben-55958-614810.jpg",  "Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.", "sunshine from window", "top right"],
                    ["./assets/pexels-simon-robben-55958-614810.jpg", "add vibrant neon lights in electric blues, magentas, and greens casting colorful reflections on surfaces, creating a cyberpunk urban atmosphere with dramatic color contrasts", "neon night, city", "top left"],
                    ["./assets/pexels-freestockpro-1227513.jpg",  "warm lighting with soft purple and pink color grading", "purple and pink hues at twilight", "auto"],
                    ["./assets/pexels-pixabay-158827.jpg", "Soft fog effects with reduced contrast throughout", "foggy morning, muted light", "auto"],
                    ["./assets/pexels-pixabay-355465.jpg", "daylight, bright sunshine", "custom", "auto" ]           
                ],
                inputs=[input_image, prompt, illumination_dropdown, direction_dropdown],
                outputs=[result, seed, final_prompt],
                fn=infer,
                cache_examples="lazy"
            )
    # update prompt when dropdown changes
    illumination_dropdown.change(
        fn=update_prompt_from_dropdown,
        inputs=[illumination_dropdown],
        outputs=[prompt]
    )
    
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            input_image,
            prompt,
            illumination_dropdown, direction_dropdown, 
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            rewrite_prompt,
            # Removed num_images_per_prompt from inputs
        ],
        outputs=[result, seed, final_prompt],
    )

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