Update app.py

app.py

@@ -1,3 +1,6 @@
+Here's the updated code with sequential generation and individual timing:
+
+```python
 import os
 import gc
 import gradio as gr
@@ -5,10 +8,10 @@ import numpy as np
 import random
 import spaces
 import torch
+import time
 from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
 from PIL import Image
 from pathlib import Path
-import concurrent.futures
 import threading
 from typing import Iterable
 from gradio.themes import Soft
@@ -115,9 +118,6 @@ pipe_small_decoder = Flux2KleinPipeline.from_pretrained(
 )
 pipe_small_decoder.enable_model_cpu_offload()
 
-pipe_lock_standard = threading.Lock()
-pipe_lock_small    = threading.Lock()
-
 # ── Helpers ───────────────────────────────────────────────────────────────────
 
 def update_dimensions_from_image(image_list):
@@ -168,18 +168,18 @@ def parse_input_images(input_images):
     return None
 
 
-# ── Per-pipeline worker ───────────────────────────────────────────────────────
-
-def run_pipeline(pipe, lock, kwargs, seed):
-    with lock:
-        gen    = torch.Generator(device="cpu").manual_seed(seed)
-        result = pipe(**kwargs, generator=gen).images[0]
-    return result
+def format_time(seconds: float) -> str:
+    """Format seconds into a human-readable string."""
+    if seconds < 60:
+        return f"{seconds:.2f}s"
+    minutes = int(seconds // 60)
+    secs    = seconds % 60
+    return f"{minutes}m {secs:.2f}s"
 
 
 # ── Inference ─────────────────────────────────────────────────────────────────
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=240)
 def infer(
     prompt,
     input_images=None,
@@ -212,28 +212,47 @@ def infer(
     if image_list is not None:
         shared_kwargs["image"] = image_list
 
-    progress(0.05, desc="⚡ Launching both pipelines simultaneously...")
+    # ── Step 1: Standard VAE ──────────────────────────────────────────────────
+    progress(0.05, desc="🟦 Running Standard VAE generation...")
+    print("Starting Standard VAE generation...")
 
-    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
-        future_std   = executor.submit(run_pipeline, pipe_standard,     pipe_lock_standard, shared_kwargs, seed)
-        future_small = executor.submit(run_pipeline, pipe_small_decoder, pipe_lock_small,    shared_kwargs, seed)
-        concurrent.futures.wait(
-            [future_std, future_small],
-            return_when=concurrent.futures.ALL_COMPLETED,
-        )
+    t0_std        = time.perf_counter()
+    gen_std       = torch.Generator(device="cpu").manual_seed(seed)
+    out_standard  = pipe_standard(**shared_kwargs, generator=gen_std).images[0]
+    t1_std        = time.perf_counter()
+    time_std      = t1_std - t0_std
+    time_std_str  = format_time(time_std)
+
+    print(f"Standard VAE done in {time_std_str}")
+    progress(0.55, desc=f"🟦 Standard VAE done in {time_std_str} — now running 🟩 Small Decoder VAE...")
+
+    gc.collect()
+    torch.cuda.empty_cache()
 
-    progress(0.95, desc="✅ Both pipelines done!")
+    # ── Step 2: Small Decoder VAE ─────────────────────────────────────────────
+    print("Starting Small Decoder VAE generation...")
 
-    out_standard = future_std.result()
-    out_small    = future_small.result()
+    t0_small       = time.perf_counter()
+    gen_small      = torch.Generator(device="cpu").manual_seed(seed)
+    out_small      = pipe_small_decoder(**shared_kwargs, generator=gen_small).images[0]
+    t1_small       = time.perf_counter()
+    time_small     = t1_small - t0_small
+    time_small_str = format_time(time_small)
+
+    print(f"Small Decoder VAE done in {time_small_str}")
+    progress(1.0, desc=f"✅ Both done! Standard: {time_std_str} | Small Decoder: {time_small_str}")
 
     gc.collect()
     torch.cuda.empty_cache()
 
-    return out_standard, out_small, seed
+    # ── Build timing label strings ────────────────────────────────────────────
+    label_std   = f"🟦 Standard VAE — ⏱ {time_std_str}"
+    label_small = f"🟩 Small Decoder VAE — ⏱ {time_small_str}"
+
+    return out_standard, out_small, seed, label_std, label_small
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=240)
 def infer_example(images, prompt):
     if not images:
         images_list = None
@@ -241,8 +260,8 @@ def infer_example(images, prompt):
         images_list = [images]
     else:
         images_list = images
-        
-    out_std, out_small, seed_used = infer(
+
+    out_std, out_small, seed_used, label_std, label_small = infer(
         prompt=prompt,
         input_images=images_list,
         seed=0,
@@ -252,7 +271,7 @@ def infer_example(images, prompt):
         num_inference_steps=4,
         guidance_scale=1.0,
     )
-    return out_std, out_small, seed_used
+    return out_std, out_small, seed_used, label_std, label_small
 
 
 # ── CSS ───────────────────────────────────────────────────────────────────────
@@ -274,6 +293,16 @@ css = """
     display: block;
     margin-bottom: 6px;
 }
+.timing-label {
+    text-align: center;
+    font-size: 0.92em;
+    font-weight: 600;
+    color: #555;
+    margin-top: 4px;
+    padding: 3px 10px;
+    border-radius: 12px;
+    background: rgba(255,255,255,0.6);
+}
 """
 
 # ── UI ────────────────────────────────────────────────────────────────────────
@@ -287,8 +316,8 @@ with gr.Blocks() as demo:
             elem_id="main-title",
         )
         gr.Markdown(
-            "Compare **FLUX.2-klein-4B** side-by-side with two VAE decoders — generated **simultaneously** from the **same seed**. \n"
-            "🟦 **Standard VAE** vs 🟩 **Small Decoder VAE** ([FLUX.2-small-decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder)) · [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)]"
+            "Compare **FLUX.2-klein-4B** side-by-side with two VAE decoders — generated **sequentially** from the **same seed**. \n"
+            "🟦 **Standard VAE** (generated first) → 🟩 **Small Decoder VAE** ([FLUX.2-small-decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder)) · [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)]"
         )
 
         # ── Main two-column row ───────────────────────────────────────────────
@@ -316,11 +345,12 @@ with gr.Blocks() as demo:
             # ── Right: outputs ────────────────────────────────────────────────
             with gr.Column():
                 with gr.Row():
+                    # Standard VAE output
                     with gr.Column():
                         gr.HTML(
                             '<span class="vae-badge" '
                             'style="background:#FFE0CC;color:#CC3700;">'
-                            '🟦 Standard VAE</span>'
+                            '🟦 Standard VAE &nbsp;·&nbsp; Generated First</span>'
                         )
                         result_standard = gr.Image(
                             label="Standard VAE",
@@ -329,12 +359,17 @@ with gr.Blocks() as demo:
                             format="png",
                             height=280,
                         )
+                        timing_standard = gr.Markdown(
+                            value="⏱ Waiting...",
+                            elem_classes=["timing-label"],
+                        )
 
+                    # Small Decoder VAE output
                     with gr.Column():
                         gr.HTML(
                             '<span class="vae-badge" '
                             'style="background:#FFF0E5;color:#E63E00;">'
-                            '🟩 Small Decoder VAE</span>'
+                            '🟩 Small Decoder VAE &nbsp;·&nbsp; Generated Second</span>'
                         )
                         result_small = gr.Image(
                             label="Small Decoder VAE",
@@ -343,8 +378,12 @@ with gr.Blocks() as demo:
                             format="png",
                             height=280,
                         )
+                        timing_small = gr.Markdown(
+                            value="⏱ Waiting...",
+                            elem_classes=["timing-label"],
+                        )
 
-                with gr.Accordion("Advanced Settings", open=False, visible=False):
+                with gr.Accordion("Advanced Settings", open=False):
                     seed_output = gr.Number(label="Seed Used", precision=0)
                     seed = gr.Slider(
                         label="Seed",
@@ -387,19 +426,20 @@ with gr.Blocks() as demo:
                             value=1.0,
                         )
 
-        # ── Standard App-Style Examples ───────────────────────────────────────
+        # ── Examples ──────────────────────────────────────────────────────────
         gr.Examples(
             examples=[
                 [["examples/1.jpg"], "Change the weather to stormy."],
                 [["examples/2.jpg"], "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition."],
                 [["examples/3.jpg"], "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent."],
                 [["examples/4.jpg"], "Make the texture high-resolution."],
+                [None, "A beautiful cyberpunk cityscape at night, neon lights, highly detailed."],
             ],
             inputs=[input_images, prompt],
-            outputs=[result_standard, result_small, seed_output],
+            outputs=[result_standard, result_small, seed_output, timing_standard, timing_small],
             fn=infer_example,
             cache_examples=False,
-            label="Examples"
+            label="Examples",
         )
 
         gr.Markdown(
@@ -428,13 +468,12 @@ with gr.Blocks() as demo:
             num_inference_steps,
             guidance_scale,
         ],
-        outputs=[result_standard, result_small, seed_output],
+        outputs=[result_standard, result_small, seed_output, timing_standard, timing_small],
     )
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch(
-        css=css,
-        theme=orange_red_theme,
+        css=css, theme=orange_red_theme,
         ssr_mode=False,
         show_error=True,
     )