update app

app.py

@@ -3,31 +3,39 @@ import numpy as np
 import random
 import torch
 import spaces
-
+import os
 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
+
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
-from typing import Iterable
-
-# --- Custom Theme Definition ---
-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",
+
+colors.steel_blue = colors.Color(
+    name="steel_blue",
+    c50="#EBF3F8",
+    c100="#D3E5F0",
+    c200="#A8CCE1",
+    c300="#7DB3D2",
+    c400="#529AC3",
+    c500="#4682B4",
+    c600="#3E72A0",
+    c700="#36638C",
+    c800="#2E5378",
+    c900="#264364",
+    c950="#1E3450",
 )
 
-class OrangeRedTheme(Soft):
+class SteelBlueTheme(Soft):
     def __init__(
         self,
         *,
         primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red, # Use the new color
+        secondary_hue: colors.Color | str = colors.steel_blue,
         neutral_hue: colors.Color | str = colors.slate,
         text_size: sizes.Size | str = sizes.text_lg,
         font: fonts.Font | str | Iterable[fonts.Font | str] = (
@@ -47,163 +55,264 @@ class OrangeRedTheme(Soft):
         )
         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_text_color="white",
+            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_800)",
+            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_500)",
+            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()
-
+steel_blue_theme = SteelBlueTheme()
 
-# --- Model Loading ---
+# --- Constants and Setup ---
+MAX_SEED = np.iinfo(np.int32).max
 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)
-# Load all LoRA adapters
-pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multiple-angles", weight_name="镜头转换.safetensors", adapter_name="angles")
-pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Light_restoration", weight_name="移除光影.safetensors", adapter_name="light_restoration")
-pipe.load_lora_weights("autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime", weight_name="Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors", adapter_name="photo_to_anime")
-pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight", weight_name="Qwen-Edit-Relight.safetensors", adapter_name="relight")
+# --- Model Loading ---
+# Load the base pipeline and the optimized transformer
+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)
 
+# Load all LoRA adapters with unique names
+pipe.load_lora_weights(
+    "dx8152/Qwen-Image-Edit-2509-Light_restoration", 
+    weight_name="移除光影.safetensors", 
+    adapter_name="light_restoration"
+)
+pipe.load_lora_weights(
+    "dx8152/Qwen-Edit-2509-Multiple-angles",
+    weight_name="镜头转换.safetensors",
+    adapter_name="multiple_angles"
+)
+pipe.load_lora_weights(
+    "autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime",
+    weight_name="Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors",
+    adapter_name="photo_to_anime"
+)
+pipe.load_lora_weights(
+    "dx8152/Qwen-Image-Edit-2509-Relight",
+    weight_name="Qwen-Edit-Relight.safetensors",
+    adapter_name="relight"
+)
 
+# Apply optimizations
 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
-
+# --- Inference Logic ---
 @spaces.GPU
-def infer(input_image, prompt, lora_adapter, seed=42, randomize_seed=True, guidance_scale=1.0, steps=4, progress=gr.Progress(track_tqdm=True)):
-    """
-    Perform image editing based on the selected LoRA adapter and prompt.
-    """
-    if not input_image:
-        raise gr.Error("Please upload an image for editing.")
-
-    # Set the LoRA adapter based on user selection
-    if lora_adapter == "Multiple Angles":
-        pipe.set_adapters(["angles"], adapter_weights=[1.0])
-    elif lora_adapter == "Light Restoration":
+def infer(
+    image,
+    prompt,
+    lora_adapter,
+    seed,
+    randomize_seed,
+    true_guidance_scale,
+    num_inference_steps,
+    height,
+    width,
+    progress=gr.Progress(track_tqdm=True)
+):
+    if image is None:
+        raise gr.Error("Please upload an image to get started.")
+    
+    # Set the active LoRA adapter based on user selection
+    if lora_adapter == "Shadow/Light Restoration":
         pipe.set_adapters(["light_restoration"], adapter_weights=[1.0])
+    elif lora_adapter == "Multiple Angles":
+        pipe.set_adapters(["multiple_angles"], adapter_weights=[1.0])
     elif lora_adapter == "Photo to Anime":
         pipe.set_adapters(["photo_to_anime"], adapter_weights=[1.0])
-    elif lora_adapter == "Relight":
+    elif lora_adapter == "Advanced Relighting":
         pipe.set_adapters(["relight"], adapter_weights=[1.0])
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
     generator = torch.Generator(device=device).manual_seed(seed)
-    
-    original_image = input_image.copy().convert("RGB")
 
     result = pipe(
-        image=original_image,
+        image=image.convert("RGB"),
         prompt=prompt,
-        height=original_image.size[1],
-        width=original_image.size[0],
-        num_inference_steps=steps,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
         generator=generator,
-        true_cfg_scale=guidance_scale,
+        true_cfg_scale=true_guidance_scale,
         num_images_per_prompt=1,
     ).images[0]
 
-    return result, seed, gr.Button(visible=True)
+    return result, seed
+
+# --- UI Helper Functions ---
+def update_dimensions_on_upload(image):
+    """Adjusts the height and width sliders to match the uploaded image's aspect ratio."""
+    if image is None:
+        return 1024, 1024
+    
+    original_width, original_height = image.size
+    
+    if original_width > original_height:
+        new_width = 1024
+        new_height = int(new_width * (original_height / original_width))
+    else:
+        new_height = 1024
+        new_width = int(new_height * (original_width / original_height))
+        
+    # Ensure dimensions are multiples of 8 for model compatibility
+    new_width = (new_width // 8) * 8
+    new_height = (new_height // 8) * 8
+    
+    return new_width, new_height
 
+def update_prompt_on_adapter_change(adapter_name):
+    """Provides a suggested prompt when a new adapter is selected."""
+    prompts = {
+        "Shadow/Light Restoration": "Remove shadows and relight the image using soft lighting.",
+        "Multiple Angles": "A photo of the scene from a top-down view.",
+        "Photo to Anime": "Transform into anime, masterpiece, best quality.",
+        "Advanced Relighting": "Relight the image using soft, diffused lighting that simulates sunlight filtering through curtains."
+    }
+    return prompts.get(adapter_name, "")
 
-# --- UI ---
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 960px;
+# --- Gradio UI ---
+css = '''
+#col-container { 
+    max-width: 960px; 
+    margin: 0 auto; 
+    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
 }
-#main-title h1 {font-size: 2.1em !important;}
-"""
+.dark .progress-text { color: white !important }
+#examples { max-width: 960px; 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-box {
+    border-radius: 16px !important;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1) !important;
+}
+'''
 
-with gr.Blocks(css=css, theme=orange_red_theme) as demo:
+with gr.Blocks(theme=steel_blue_theme, css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# **Qwen-Image-Edit-2509-LoRAs-Fast**", elem_id="main-title")
-        gr.Markdown("Image manipulation with Qwen Image Edit 2509 and various LoRA adapters.")
-
+        gr.Markdown("# Qwen Image Edit - Fast LoRA")
         with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(label="Upload Image", type="pil", height=400)
-                with gr.Row():
-                    prompt = gr.Text(
-                        label="Edit Prompt",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your prompt for editing",
-                        container=False,
-                    )
-                    run_button = gr.Button("Run", variant="primary", scale=0)
-
+            with gr.Column(scale=1):
+                image = gr.Image(label="Upload Image", type="pil", height=450)
+                
                 lora_adapter = gr.Dropdown(
-                    label="Choose LoRA Adapter",
-                    choices=["Multiple Angles", "Light Restoration", "Photo to Anime", "Relight"],
-                    value="Multiple Angles"
+                    label="Choose an Editing Tool",
+                    choices=[
+                        "Shadow/Light Restoration", 
+                        "Multiple Angles", 
+                        "Photo to Anime", 
+                        "Advanced Relighting"
+                    ],
+                    value="Shadow/Light Restoration"
+                )
+
+                prompt = gr.Textbox(
+                    label="Prompt", 
+                    value="Remove shadows and relight the image using soft lighting.",
+                    lines=2
                 )
+                
+                run_btn = gr.Button("Generate", 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)
-                    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=30,
-                        value=4,
-                        step=1
-                    )
-
-            with gr.Column():
-                output_image = gr.Image(label="Output Image", show_label=True, interactive=False, format="png", height=400)
-                reuse_button = gr.Button("Reuse this image", 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)
+                    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="Output Image", interactive=False, height=500, format="png")
 
         gr.Examples(
+            elem_id="examples",
             examples=[
-                ["examples/sea.png", "Rotate the camera 90 degrees to the left.", "Multiple Angles"],
-                ["examples/shadow.jpg", "Remove shadows and relight the image using soft lighting.", "Light Restoration"],
-                ["examples/girl.jpg", "transform into anime", "Photo to Anime"],
-                ["examples/dark.jpg", "Relight the image using soft, diffused lighting that simulates sunlight filtering through curtains.", "Relight"],
+                [
+                    "examples/example1.png",
+                    "A photo of the scene from a low angle shot.",
+                    "Multiple Angles",
+                ],
+                [
+                    "examples/example2.png",
+                    "Remove shadows and relight the image using soft lighting.",
+                    "Shadow/Light Restoration",
+                ],
+                [
+                    "examples/example3.png",
+                    "Transform into anime, masterpiece, best quality, girl with cherry blossoms.",
+                    "Photo to Anime",
+                ],
+                [
+                    "examples/example4.png",
+                    "Relight the image using soft, diffused lighting that simulates sunlight filtering through curtains.",
+                    "Advanced Relighting",
+                ],
             ],
-            inputs=[input_image, prompt, lora_adapter],
-            outputs=[output_image, seed],
+            inputs=[image, prompt, lora_adapter],
+            outputs=[result, seed],
             fn=infer,
-            cache_examples="lazy",
-            label="Examples"
+            cache_examples=False
         )
-
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
-        outputs=[output_image, seed, reuse_button]
+            
+    # --- Event Handlers ---
+    run_btn.click(
+        fn=infer, 
+        inputs=[image, prompt, lora_adapter, seed, randomize_seed, true_guidance_scale, num_inference_steps, height, width], 
+        outputs=[result, seed]
     )
-
-    reuse_button.click(
-        fn=lambda image: image,
-        inputs=[output_image],
-        outputs=[input_image]
+    
+    image.upload(
+        fn=update_dimensions_on_upload,
+        inputs=[image],
+        outputs=[width, height]
+    )
+    
+    lora_adapter.change(
+        fn=update_prompt_on_adapter_change,
+        inputs=[lora_adapter],
+        outputs=[prompt]
     )
 
-demo.launch(mcp_server=True, ssr_mode=False, show_error=True)
+demo.launch(mcp_server=False)