app.py

import os
import json
import time
import requests
import random
import numpy as np
import spaces
import torch
from PIL import Image
import gradio as gr

# --- Qwen Specific Imports ---
from diffusers import FlowMatchEulerDiscreteScheduler
# Assuming the qwenimage package is available in the environment
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

from huggingface_hub import (
    hf_hub_download,
    HfFileSystem,
    ModelCard
)

from typing import Iterable
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# =========================================
# THEME CONFIGURATION
# =========================================

colors.orange_red = colors.Color(
    name="orange_red",
    c50="#FFF0E5",
    c100="#FFE0CC",
    c200="#FFC299",
    c300="#FFA366",
    c400="#FF8533",
    c500="#FF4500",
    c600="#E63E00",
    c700="#CC3700",
    c800="#B33000",
    c900="#992900",
    c950="#802200",
)

class OrangeRedTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.orange_red,
        neutral_hue: colors.Color | str = colors.slate,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_secondary_text_color="black",
            button_secondary_text_color_hover="white",
            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
            button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px",
            color_accent_soft="*primary_100",
            block_label_background_fill="*primary_200",
        )

orange_red_theme = OrangeRedTheme()

# =========================================
# LORA CONFIGURATION (The "DLC" List)
# =========================================

loras = [
    {
        "image": "https://huggingface.co/autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime/resolve/main/images/example.jpg", 
        "title": "Photo to Anime",
        "repo": "autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime",
        "weights": "Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors",
        "trigger_word": "Transform into anime"
    },
    {
        "image": "https://huggingface.co/dx8152/Qwen-Edit-2509-Multiple-angles/resolve/main/images/example.jpg", 
        "title": "Multiple Angles",
        "repo": "dx8152/Qwen-Edit-2509-Multiple-angles",
        "weights": "镜头转换.safetensors",
        "trigger_word": "Rotate camera"
    },
    {
        "image": "https://huggingface.co/dx8152/Qwen-Image-Edit-2509-Light_restoration/resolve/main/images/example.jpg", 
        "title": "Light Restoration",
        "repo": "dx8152/Qwen-Image-Edit-2509-Light_restoration",
        "weights": "移除光影.safetensors",
        "trigger_word": "Remove shadows"
    },
    {
        "image": "https://huggingface.co/dx8152/Qwen-Image-Edit-2509-Relight/resolve/main/images/example.jpg", 
        "title": "Relight",
        "repo": "dx8152/Qwen-Image-Edit-2509-Relight",
        "weights": "Qwen-Edit-Relight.safetensors",
        "trigger_word": "Relight the image"
    },
    {
        "image": "https://huggingface.co/dx8152/Qwen-Edit-2509-Multi-Angle-Lighting/resolve/main/images/example.jpg",
        "title": "Multi-Angle Lighting",
        "repo": "dx8152/Qwen-Edit-2509-Multi-Angle-Lighting",
        "weights": "多角度灯光-251116.safetensors",
        "trigger_word": "Light source from"
    },
    {
        "image": "https://huggingface.co/tlennon-ie/qwen-edit-skin/resolve/main/images/example.jpg",
        "title": "Edit Skin",
        "repo": "tlennon-ie/qwen-edit-skin",
        "weights": "qwen-edit-skin_1.1_000002750.safetensors",
        "trigger_word": "Make skin details prominent"
    },
    {
        "image": "https://huggingface.co/lovis93/next-scene-qwen-image-lora-2509/resolve/main/images/example.jpg",
        "title": "Next Scene",
        "repo": "lovis93/next-scene-qwen-image-lora-2509",
        "weights": "next-scene_lora-v2-3000.safetensors",
        "trigger_word": "Next scene cinematic"
    },
    {
        "image": "https://huggingface.co/vafipas663/Qwen-Edit-2509-Upscale-LoRA/resolve/main/images/example.jpg",
        "title": "Upscale Image",
        "repo": "vafipas663/Qwen-Edit-2509-Upscale-LoRA",
        "weights": "qwen-edit-enhance_64-v3_000001000.safetensors",
        "trigger_word": "Upscale the image"
    },
]

# =========================================
# MODEL SETUP
# =========================================

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
base_model = "Qwen/Qwen-Image-Edit-2509"

print(f"Loading {base_model} pipeline...")

# Initialize Pipeline
pipe = QwenImageEditPlusPipeline.from_pretrained(
    base_model,
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "linoyts/Qwen-Image-Edit-Rapid-AIO",
        subfolder='transformer',
        torch_dtype=dtype,
        device_map=device
    ),
    torch_dtype=dtype,
).to(device)

# Apply Optimization
try:
    print("Applying FA3 Processor...")
    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
    print("Optimization applied successfully.")
except Exception as e:
    print(f"Optimization warning: {e}. Continuing with standard pipeline.")

MAX_SEED = np.iinfo(np.int32).max

# =========================================
# HELPER FUNCTIONS
# =========================================

class calculateDuration:
    def __init__(self, activity_name=""):
        self.activity_name = activity_name

    def __enter__(self):
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.elapsed_time = self.end_time - self.start_time
        if self.activity_name:
            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
        else:
            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")

def update_dimensions_on_upload(image):
    if image is None:
        return 1024, 1024
    
    original_width, original_height = image.size
    
    if original_width > original_height:
        new_width = 1024
        aspect_ratio = original_height / original_width
        new_height = int(new_width * aspect_ratio)
    else:
        new_height = 1024
        aspect_ratio = original_width / original_height
        new_width = int(new_height * aspect_ratio)
        
    # Ensure dimensions are multiples of 8 (Qwen requirement)
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    
    return new_width, new_height

def update_selection(evt: gr.SelectData, current_prompt):
    selected_lora = loras[evt.index]
    trigger = selected_lora.get("trigger_word", "")
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    lora_repo = selected_lora["repo"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
    
    # Append trigger word to prompt if not present
    new_prompt = current_prompt
    if trigger and trigger not in current_prompt:
        if current_prompt:
            new_prompt = f"{trigger} {current_prompt}"
        else:
            new_prompt = trigger

    return (
        gr.update(placeholder=new_placeholder, value=new_prompt),
        updated_text,
        evt.index
    )

def check_custom_model(link):
    if link.startswith("https://"):
        if "huggingface.co" in link:
            parts = link.split("huggingface.co/")
            repo_part = parts[1].strip()
            return repo_part, link
    return link, link

# =========================================
# INFERENCE LOGIC
# =========================================

@spaces.GPU(duration=60)
def run_lora(
    input_image, 
    prompt, 
    steps, 
    guidance_scale, 
    selected_index, 
    randomize_seed, 
    seed, 
    custom_lora_path,
    progress=gr.Progress(track_tqdm=True)
):
    if input_image is None:
        raise gr.Error("Input image is required for Qwen Image Edit.")

    # 1. Clean up previous LoRAs
    with calculateDuration("Unloading LoRA"):
        try:
            pipe.unload_lora_weights()
        except Exception:
            pass 
    
    # 2. Determine which LoRA to load
    lora_repo = None
    weight_name = None
    adapter_name = "default"

    if custom_lora_path and custom_lora_path.strip() != "":
        repo, link = check_custom_model(custom_lora_path)
        lora_repo = repo
        print(f"Attempting to load custom LoRA: {lora_repo}")
    elif selected_index is not None and selected_index < len(loras):
        selected_lora = loras[selected_index]
        lora_repo = selected_lora["repo"]
        weight_name = selected_lora.get("weights", None)
        print(f"Loading Gallery LoRA: {selected_lora['title']}")
    else:
        print("No LoRA selected. Running Base Model.")

    # 3. Load LoRA
    if lora_repo:
        with calculateDuration(f"Loading LoRA weights"):
            try:
                pipe.load_lora_weights(
                    lora_repo, 
                    weight_name=weight_name, 
                    adapter_name=adapter_name
                )
                pipe.set_adapters([adapter_name], adapter_weights=[1.0])
            except Exception as e:
                print(f"Error loading LoRA: {e}")
                gr.Warning(f"Failed to load LoRA {lora_repo}. Generating with base model.")

    # 4. Prepare Seed
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)

    # 5. Process Image
    original_image = input_image.convert("RGB")
    width, height = update_dimensions_on_upload(original_image)
    
    negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"

    # 6. Generate
    with calculateDuration("Generating image"):
        final_image = pipe(
            image=original_image,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_inference_steps=int(steps),
            true_cfg_scale=guidance_scale, 
            generator=generator,
        ).images[0]
        
    return final_image, seed, gr.update(visible=False)

def add_custom_lora(custom_path):
    if not custom_path:
        return gr.update(visible=False), None, "No custom LoRA"
    
    repo, _ = check_custom_model(custom_path)
    card_html = f'''
    <div class="custom_lora_card" style="border:1px solid #ccc; padding:10px; border-radius:8px; margin-top:10px;">
        <h3>Custom LoRA Active</h3>
        <code>{repo}</code>
    </div>
    '''
    return gr.update(visible=True, value=card_html), None, repo

def remove_custom_lora():
    return gr.update(visible=False), None, ""

# =========================================
# UI CONSTRUCTION
# =========================================

css = '''
#gen_btn{height: 100%}
#gen_column{align-self: stretch}
#title{text-align: center}
#title h1{font-size: 3em; display:inline-flex; align-items:center}
#title img{width: 100px; margin-right: 0.5em}
#gallery .grid-wrap{height: 15vh}
#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
.card_internal{display: flex;height: 100px;margin-top: .5em}
.card_internal img{margin-right: 1em}
#progress{height:30px}
'''

with gr.Blocks(delete_cache=(60, 60)) as demo:
    title = gr.HTML(
        """<h1>Qwen Image Edit 2.5 DLC 🧪</h1>""",
        elem_id="title",
    )
    
    selected_index = gr.State(None)
    custom_lora_state = gr.State("") 

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(label="Edit Prompt", lines=2, placeholder="✦︎ Select a LoRA from the gallery below and describe the edit...")
            
        with gr.Column(scale=1, elem_id="gen_column"):
            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
            
    with gr.Row():
        input_image = gr.Image(label="Upload Input Image (Required)", type="pil", height=300)
            
    with gr.Row():
        with gr.Column():
            selected_info = gr.Markdown("### No LoRA Selected (Base Model)")
            
            gallery = gr.Gallery(
                [(item.get("image", ""), item["title"]) for item in loras],
                label="Available LoRAs",
                allow_preview=False,
                columns=4,
                elem_id="gallery",
            )
            
            with gr.Group():
                custom_lora_input = gr.Textbox(label="Load Custom LoRA", placeholder="Enter HuggingFace Repo ID (e.g. autoweeb/Qwen-Image-Edit-...)")
                gr.Markdown("[See compatible Qwen LoRAs](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image-Edit-2509)", elem_id="lora_list")
            
            custom_lora_info = gr.HTML(visible=False)
            remove_custom_btn = gr.Button("Remove Custom LoRA", visible=False)

        with gr.Column():
            progress_bar = gr.Markdown(elem_id="progress", visible=False)
            result = gr.Image(label="Edited Image", format="png", height=630)

    with gr.Row():
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Column():
                with gr.Row():
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=4)
                
                with gr.Row():
                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)

    # Event Wiring
    gallery.select(
        update_selection,
        inputs=[prompt],
        outputs=[prompt, selected_info, selected_index]
    )

    custom_lora_input.submit(
        add_custom_lora,
        inputs=[custom_lora_input],
        outputs=[custom_lora_info, selected_index, custom_lora_state]
    ).then(
        lambda: gr.update(visible=True), outputs=[remove_custom_btn]
    )

    remove_custom_btn.click(
        remove_custom_lora,
        outputs=[custom_lora_info, selected_index, custom_lora_state]
    ).then(
        lambda: gr.update(visible=False), outputs=[remove_custom_btn]
    ).then(
        lambda: "", outputs=[custom_lora_input]
    )

    gr.on(
        triggers=[generate_button.click, prompt.submit],
        fn=run_lora,
        inputs=[input_image, prompt, steps, guidance_scale, selected_index, randomize_seed, seed, custom_lora_state],
        outputs=[result, seed, progress_bar]
    )

demo.queue()
demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)