app.py

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

from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler
from optimization import optimize_pipeline_
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

import math
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from PIL import Image
import os
import gradio as gr
from gradio_client import Client, handle_file
import tempfile


# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

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

pipe.load_lora_weights(
        "dx8152/Qwen-Image-Edit-2509-Light_restoration", 
        weight_name="移除光影.safetensors", adapter_name="light_restoration"
    )

pipe.set_adapters(["light_restoration"], adapter_weights=[1.])
pipe.fuse_lora(adapter_names=["light_restoration"], lora_scale=1.0)
pipe.unload_lora_weights()

pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")


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

def build_light_restoration_prompt(brightness, contrast, exposure, saturation):
    prompt_parts = []

    # Brightness adjustment
    if brightness > 0:
        prompt_parts.append(f"Increase brightness by {brightness}%")
    elif brightness < 0:
        prompt_parts.append(f"Decrease brightness by {abs(brightness)}%")

    # Contrast adjustment
    if contrast > 0:
        prompt_parts.append(f"Increase contrast by {contrast}%")
    elif contrast < 0:
        prompt_parts.append(f"Decrease contrast by {abs(contrast)}%")

    # Exposure adjustment
    if exposure > 0:
        prompt_parts.append(f"Increase exposure by {exposure}%")
    elif exposure < 0:
        prompt_parts.append(f"Decrease exposure by {abs(exposure)}%")

    # Saturation adjustment
    if saturation > 0:
        prompt_parts.append(f"Increase saturation by {saturation}%")
    elif saturation < 0:
        prompt_parts.append(f"Decrease saturation by {abs(saturation)}%")

    final_prompt = ", ".join(prompt_parts).strip()
    return final_prompt if final_prompt else "Restore image lighting"


@spaces.GPU
def infer_light_restoration(
    image,
    brightness,
    contrast,
    exposure,
    saturation,
    seed,
    randomize_seed,
    true_guidance_scale,
    num_inference_steps,
    height,
    width,
    prev_output = None,
    progress=gr.Progress(track_tqdm=True)
):
    prompt = build_light_restoration_prompt(brightness, contrast, exposure, saturation)
    print(f"Generated Prompt: {prompt}")

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

    # Choose input image (prefer uploaded, else last output)
    pil_images = []
    if image is not None:
        if isinstance(image, Image.Image):
            pil_images.append(image.convert("RGB"))
        elif hasattr(image, "name"):
            pil_images.append(Image.open(image.name).convert("RGB"))
    elif prev_output:
        pil_images.append(prev_output.convert("RGB"))

    if len(pil_images) == 0:
        raise gr.Error("Please upload an image first.")

    if prompt == "Restore image lighting":
        return image, seed, prompt
    
    result = pipe(
        image=pil_images,
        prompt=prompt,
        height=height if height != 0 else None,
        width=width if width != 0 else None,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=1,
    ).images[0]

    return result, seed, prompt


# --- UI ---
css = '''
#col-container { 
    max-width: 900px; 
    margin: 0 auto; 
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
}
.dark .progress-text{color: white !important}
#examples{max-width: 900px; margin: 0 auto; }
.gradio-container {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
}
.gr-button-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 12px 24px !important;
    font-weight: 600 !important;
}
.gr-button {
    border-radius: 12px !important;
    padding: 10px 20px !important;
}
.gr-box {
    border-radius: 16px !important;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1) !important;
}
'''

def reset_all():
    return [0, 0, 0, 0]

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
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    
    return new_width, new_height


with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# ✨ Light Restoration Studio")
        gr.Markdown("""
            Professional image light restoration powered by Qwen Image Edit 2509<br>
            Using [dx8152's Light Restoration LoRA](https://huggingface.co/dx8152/Qwen-Image-Edit-2509-Light_restoration) 
            and [Phr00t/Qwen-Image-Edit-Rapid-AIO](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) for fast inference 💨<br>
            Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(label="📸 Upload Image", type="pil", height=400)
                prev_output = gr.Image(value=None, visible=False)

                gr.Markdown("### 🎨 Light Adjustments")
                brightness = gr.Slider(label="☀️ Brightness", minimum=-50, maximum=50, step=5, value=0)
                contrast = gr.Slider(label="🌓 Contrast", minimum=-50, maximum=50, step=5, value=0)
                exposure = gr.Slider(label="💡 Exposure", minimum=-50, maximum=50, step=5, value=0)
                saturation = gr.Slider(label="🎨 Saturation", minimum=-50, maximum=50, step=5, value=0)
                
                with gr.Row():
                    reset_btn = gr.Button("🔄 Reset", size="lg")
                    run_btn = gr.Button("✨ Restore", variant="primary", size="lg")

                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)
                    true_guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                    num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=40, step=1, value=4)
                    height = gr.Slider(label="Height", minimum=256, maximum=2048, step=8, value=1024)
                    width = gr.Slider(label="Width", minimum=256, maximum=2048, step=8, value=1024)

            with gr.Column(scale=1):
                result = gr.Image(label="✨ Restored Image", interactive=False, height=400)
                prompt_preview = gr.Textbox(label="Generated Prompt", interactive=False)
                    
    inputs = [
        image, brightness, contrast, exposure, saturation,
        seed, randomize_seed, true_guidance_scale, num_inference_steps, height, width, prev_output
    ]
    outputs = [result, seed, prompt_preview]

    # Reset behavior
    reset_btn.click(
        fn=reset_all,
        inputs=None,
        outputs=[brightness, contrast, exposure, saturation],
        queue=False
    )

    # Manual generation
    run_btn.click(
        fn=infer_light_restoration, 
        inputs=inputs, 
        outputs=outputs
    ).then(lambda img, *_: img, inputs=[result], outputs=[prev_output])
    
    # Image upload triggers dimension update and control reset
    image.upload(
        fn=update_dimensions_on_upload,
        inputs=[image],
        outputs=[width, height]
    ).then(
        fn=reset_all,
        inputs=None,
        outputs=[brightness, contrast, exposure, saturation],
        queue=False
    )

    # Live updates
    for control in [brightness, contrast, exposure, saturation]:
        control.release(fn=infer_light_restoration, inputs=inputs, outputs=outputs)

demo.launch()