Update app.py

app.py

@@ -1,13 +1,12 @@
-import os
 import gc
 import random
 import numpy as np
-from typing import List
-
 import torch
+from PIL import Image
+from typing import List, Tuple
+
 import spaces
 import gradio as gr
-from PIL import Image
 from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
 
 # --- App Configuration ---
@@ -17,6 +16,10 @@ MAX_IMAGE_SIZE = 1024
 dtype = torch.bfloat16
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+if torch.cuda.is_available():
+    print("current device:", torch.cuda.current_device())
+    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
+
 # --- Model Loading ---
 print("Loading Small Decoder VAE...")
 vae_small = AutoencoderKLFlux2.from_pretrained(
@@ -30,12 +33,11 @@ pipe = Flux2KleinPipeline.from_pretrained(
     vae=vae_small,
     torch_dtype=dtype,
 ).to(device)
-# Optional memory optimization
-# pipe.enable_model_cpu_offload()
+pipe.enable_model_cpu_offload()
 
 # --- Utility Functions ---
-def calc_dimensions(pil_img: Image.Image):
-    """Calculate target dimensions preserving aspect ratio, capped at 1024."""
+def calc_dimensions(pil_img: Image.Image) -> Tuple[int, int]:
+    """Calculates dimensions preserving aspect ratio, snapped to multiples of 8."""
     iw, ih = pil_img.size
     aspect = iw / ih
 
@@ -46,62 +48,66 @@ def calc_dimensions(pil_img: Image.Image):
         new_height = 1024
         new_width  = int(round(1024 * aspect))
 
-    # Ensure dimensions are multiples of 8 and within bounds
-    new_width  = max(256, min(MAX_IMAGE_SIZE, round(new_width  / 8) * 8))
-    new_height = max(256, min(MAX_IMAGE_SIZE, round(new_height / 8) * 8))
+    new_width  = max(256, min(1024, round(new_width  / 8) * 8))
+    new_height = max(256, min(1024, round(new_height / 8) * 8))
     return new_width, new_height
 
-def parse_and_resize_images(image_paths: List[str], width: int, height: int):
-    """Load and resize all provided images."""
-    if not image_paths:
+def parse_and_resize_images(gallery_items: List, target_width: int, target_height: int) -> List[Image.Image]:
+    """Extracts images from Gradio Gallery and resizes them."""
+    if not gallery_items:
         return None
 
     resized = []
-    for path in image_paths:
+    for item in gallery_items:
         try:
-            img = Image.open(path).convert("RGB")
-            resized.append(img.resize((width, height), Image.LANCZOS))
+            # Gradio Gallery returns a list of tuples: (filepath, label)
+            filepath = item[0] if isinstance(item, (tuple, list)) else item
+            img = Image.open(filepath).convert("RGB")
+            resized.append(img.resize((target_width, target_height), Image.LANCZOS))
         except Exception as e:
-            print(f"Skipping invalid image {path}: {e}")
+            print(f"Skipping invalid image: {e}")
 
     return resized if resized else None
 
 # --- Inference Function ---
-@spaces.GPU(duration=60)
+@spaces.GPU(size="xlarge")
 def generate_image(
+    gallery_inputs,
     prompt: str,
-    image_files: List[str],
     seed: int,
     randomize_seed: bool,
     width: int,
     height: int,
-    num_inference_steps: int,
-    guidance_scale: float,
-    progress=gr.Progress(track_tqdm=True),
+    steps: int,
+    guidance: float,
+    progress=gr.Progress(track_tqdm=True)
 ):
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-    if not prompt or not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
-
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-
+    
     image_list = None
-    if image_files and len(image_files) > 0:
+    if gallery_inputs and len(gallery_inputs) > 0:
         try:
-            # Calculate dynamic dimensions based on the first input image
-            first_pil = Image.open(image_files[0]).convert("RGB")
+            # Get first image to calculate reference dimensions
+            first_item = gallery_inputs[0]
+            first_filepath = first_item[0] if isinstance(first_item, (tuple, list)) else first_item
+            first_pil = Image.open(first_filepath).convert("RGB")
+            
             calc_w, calc_h = calc_dimensions(first_pil)
-            image_list = parse_and_resize_images(image_files, calc_w, calc_h)
+            image_list = parse_and_resize_images(gallery_inputs, calc_w, calc_h)
             
-            # Override manual width/height with the image's aspect-ratio-locked dimensions
+            # Override manual width/height if images are provided to match input aspect ratio
             width, height = calc_w, calc_h
         except Exception as e:
-            print(f"Error processing uploads: {e}")
+            print(f"Error processing gallery uploads: {e}")
 
-    # Ensure constraints if running without an image upload
+    # Ensure dimensions are multiples of 8
     final_width  = max(256, min(MAX_IMAGE_SIZE, round(int(width)  / 8) * 8))
     final_height = max(256, min(MAX_IMAGE_SIZE, round(int(height) / 8) * 8))
 
@@ -109,131 +115,107 @@ def generate_image(
         prompt=prompt,
         height=final_height,
         width=final_width,
-        num_inference_steps=int(num_inference_steps),
-        guidance_scale=float(guidance_scale),
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
     )
     if image_list is not None:
         kwargs["image"] = image_list
 
-    # Generate Image
-    gen = torch.Generator(device="cpu").manual_seed(current_seed)
-    result = pipe(**kwargs, generator=gen).images[0]
+    generator = torch.Generator(device="cpu").manual_seed(current_seed)
+    result = pipe(**kwargs, generator=generator).images[0]
+
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
     return result, current_seed
 
-css = '''
-.output-card {
-    border: 1px solid var(--border-color-primary);
-    border-radius: 12px;
-    padding: 16px;
-    margin-bottom: 12px;
-    background: var(--background-fill-secondary);
-}
-'''
-
-with gr.Blocks() as demo:
-    gr.Markdown("# **Flux.2 Klein — Small Decoder**")
-    gr.Markdown(
-        "Upload an image (optional) and enter a prompt to generate or edit using the **FLUX.2-klein-4B** distilled model paired with the **Small Decoder VAE**."
-    )
+# --- Gradio UI ---
+with gr.Blocks(title="Flux.2 Klein - Small Decoder") as demo:
+    gr.Markdown("# **Flux.2 Klein — Small Decoder VAE**")
+    gr.Markdown("Upload images (optional) and enter a prompt to generate or edit with the 4B distilled model.")
 
     with gr.Row():
-        # -- Left Column: Settings --
         with gr.Column(scale=1):
-            image_files_input = gr.Gallery(
-                file_types=["image"],
-                label="Input Images (Optional)",
-                elem_classes="output-card"
+            gallery_input = gr.Gallery(
+                label="Input Images (Optional)", 
+                type="filepath", 
+                height=300, 
+                allow_preview=True,
+                elem_id="gallery_input"
             )
-            
             prompt_input = gr.Textbox(
-                label="Prompt",
-                placeholder="Describe the edit or generation...",
+                label="Prompt", 
+                placeholder="Describe the edit or generation...", 
                 lines=3
             )
 
             with gr.Accordion("Advanced Settings", open=False):
-                seed_input = gr.Slider(
-                    label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True
-                )
-                randomize_seed_checkbox = gr.Checkbox(
-                    label="Randomize seed", value=True, interactive=True
-                )
-                steps_slider = gr.Slider(
-                    label="Inference Steps", minimum=1, maximum=30, step=1, value=4
-                )
-                width_input = gr.Slider(
-                    label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024
-                )
-                height_input = gr.Slider(
-                    label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024
-                )
-                guidance_input = gr.Slider(
-                    label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0
-                )
+                with gr.Row():
+                    width_slider = gr.Slider(minimum=256, maximum=1024, step=8, value=1024, label="Width")
+                    height_slider = gr.Slider(minimum=256, maximum=1024, step=8, value=1024, label="Height")
+                
+                with gr.Row():
+                    steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps")
+                    guidance_slider = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=1.0, label="Guidance Scale")
+                
+                seed_input = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=42, label="Seed")
+                randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=True)
 
             generate_button = gr.Button("Generate Image", variant="primary")
 
-        # -- Right Column: Outputs --
         with gr.Column(scale=1):
-            with gr.Group(elem_classes="output-card"):
-                output_image = gr.Image(
-                    label="Generated Output", type="pil", interactive=False, format="png",
-                )
-                used_seed_output = gr.Textbox(
-                    label="Used Seed", interactive=False
-                )
-
-    # -- Event Listener --
+            output_image = gr.Image(label="Generated Output", type="pil", interactive=False)
+
+    # Wire up the button
     generate_button.click(
         fn=generate_image,
         inputs=[
-            prompt_input,
-            image_files_input,
-            seed_input,
-            randomize_seed_checkbox,
-            width_input,
-            height_input,
-            steps_slider,
-            guidance_input
+            gallery_input, 
+            prompt_input, 
+            seed_input, 
+            randomize_seed_checkbox, 
+            width_slider, 
+            height_slider, 
+            steps_slider, 
+            guidance_slider
         ],
-        outputs=[output_image, used_seed_output]
+        outputs=[output_image, seed_input]
     )
 
-    # -- Examples --
+    # Examples
     gr.Examples(
         examples=[
             [
-                ["examples/I1.jpg", "examples/I2.jpg"],
-                "Make her wear these glasses in Image 2.",
-                4
+                ["examples/I1.jpg", "examples/I2.jpg"], 
+                "Make her wear these glasses in Image 2."
+            ],
+            [
+                ["examples/1.jpg"], 
+                "Change the weather to stormy."
             ],
             [
-                ["examples/1.jpg"],
-                "Change the weather to stormy.",
-                4
+                ["examples/2.jpg"], 
+                "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition."
             ],
             [
-                ["examples/2.jpg"],
-                "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition.",
-                4
+                ["examples/3.jpg"], 
+                "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent."
             ],
             [
-                ["examples/3.jpg"],
-                "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent.",
-                4
+                ["examples/4.jpg"], 
+                "Make the texture high-resolution."
             ],
             [
-                ["examples/4.jpg"],
-                "Make the texture high-resolution.",
-                4
+                None, 
+                "A futuristic cyberpunk cityscape at night, neon lights reflecting in puddles, flying cars in the background."
             ]
         ],
-        inputs=[image_files_input, prompt_input, steps_slider],
-        outputs=[output_image, used_seed_output],
+        inputs=[gallery_input, prompt_input],
+        outputs=[output_image, seed_input],
         fn=generate_image,
         cache_examples=False,
     )
 
 if __name__ == "__main__":
-    demo.queue().launch(theme=gr.themes.Soft(), css=css, show_error=True, mcp_server=True)
+    demo.queue().launch(theme=gr.themes.Citrus(), show_error=True)