Update app.py

app.py

@@ -8,8 +8,6 @@ import torch
 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
@@ -90,6 +88,7 @@ MAX_SEED       = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 EXAMPLES_DIR   = Path("examples")
 
+# ── Load standard pipeline ──────────────────────────────────────────────────
 print("Loading 4B Distilled model (Standard VAE)...")
 pipe_standard = Flux2KleinPipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
@@ -97,12 +96,14 @@ pipe_standard = Flux2KleinPipeline.from_pretrained(
 )
 pipe_standard.enable_model_cpu_offload()
 
+# ── Load small decoder VAE ───────────────────────────────────────────────────
 print("Loading Small Decoder VAE...")
 vae_small = AutoencoderKLFlux2.from_pretrained(
     "black-forest-labs/FLUX.2-small-decoder",
     torch_dtype=dtype,
 )
 
+# ── Load small-decoder pipeline ──────────────────────────────────────────────
 print("Loading 4B Distilled model (Small Decoder VAE)...")
 pipe_small_decoder = Flux2KleinPipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-4B",
@@ -111,40 +112,27 @@ pipe_small_decoder = Flux2KleinPipeline.from_pretrained(
 )
 pipe_small_decoder.enable_model_cpu_offload()
 
-pipe_lock_standard = threading.Lock()
-pipe_lock_small    = threading.Lock()
-
+# ────────────────────────────────────────────────────────────────────────────
 def calc_dimensions(pil_img: Image.Image):
-    """
-    Given a PIL image return (width, height) snapped to multiples of 8,
-    fitting within 1024 px on the long side, min 256 px on each side.
-    Uses round() so we match the reference app exactly.
-    """
     iw, ih = pil_img.size
     aspect = iw / ih
 
-    if aspect >= 1:          # landscape / square
+    if aspect >= 1:
         new_width  = 1024
         new_height = int(round(1024 / aspect))
-    else:                    # portrait
+    else:
         new_height = 1024
         new_width  = int(round(1024 * aspect))
 
-    # snap to 8-pixel grid with round(), clamp to [256, 1024]
     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 update_dimensions_from_image(image_list):
-    """
-    Called by the gallery .upload() event.
-    Returns updated slider values for width and height.
-    """
     if not image_list:
         return 1024, 1024
 
-    # gallery items arrive as PIL images when type="pil"
     item = image_list[0]
     img  = item[0] if isinstance(item, tuple) else item
 
@@ -155,11 +143,8 @@ def update_dimensions_from_image(image_list):
 
     return calc_dimensions(img)
 
+
 def parse_and_resize_images(input_images, width: int, height: int):
-    """
-    Parse the gallery input and resize every frame to (width, height).
-    Returns a list[PIL.Image] or None.
-    """
     if input_images is None:
         return None
 
@@ -192,12 +177,14 @@ def parse_and_resize_images(input_images, width: int, height: int):
     ]
     return resized
 
-def run_pipeline(pipe, lock, kwargs, seed):
-    with lock:
-        gen    = torch.Generator(device="cpu").manual_seed(seed)
-        result = pipe(**kwargs, generator=gen).images[0]
+
+def run_pipeline(pipe, kwargs, seed):
+    """Run a single pipeline — no locks needed, purely sequential."""
+    gen    = torch.Generator(device="cpu").manual_seed(seed)
+    result = pipe(**kwargs, generator=gen).images[0]
     return result
 
+
 @spaces.GPU(duration=120)
 def infer(
     prompt,
@@ -219,11 +206,9 @@ def infer(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # ── width / height: derive from the first uploaded image if present ──
+    # ── Derive dimensions from the first uploaded image if present ───────────
     image_list = None
     if input_images:
-        # Re-derive dimensions from the actual first image so they are
-        # always consistent with what the pipeline will receive.
         item = (
             input_images[0][0]
             if isinstance(input_images[0], tuple)
@@ -239,10 +224,9 @@ def infer(
         if first_pil is not None:
             width, height = calc_dimensions(first_pil)
 
-        # parse + resize all images to the final (width, height)
         image_list = parse_and_resize_images(input_images, width, height)
 
-    # ensure dims are multiples of 8 even for text-only runs
+    # ensure dims are multiples of 8
     width  = max(256, min(MAX_IMAGE_SIZE, round(int(width)  / 8) * 8))
     height = max(256, min(MAX_IMAGE_SIZE, round(int(height) / 8) * 8))
 
@@ -256,28 +240,22 @@ def infer(
     if image_list is not None:
         shared_kwargs["image"] = image_list
 
-    progress(0.30, desc="Launching both pipelines simultaneously...")
+    # ── Pipeline 1: Standard Decoder ─────────────────────────────────────────
+    progress(0.10, desc="Running Pipeline 1 / 2 — Standard Decoder...")
+    out_standard = run_pipeline(pipe_standard, shared_kwargs, seed)
 
-    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,
-        )
-
-    progress(0.80, desc="✅ Both pipelines done!")
+    gc.collect()
+    torch.cuda.empty_cache()
 
-    out_standard = future_std.result()
-    out_small    = future_small.result()
+    # ── Pipeline 2: Small Decoder ─────────────────────────────────────────────
+    progress(0.55, desc="Running Pipeline 2 / 2 — Small Decoder...")
+    out_small = run_pipeline(pipe_small_decoder, shared_kwargs, seed)
 
     gc.collect()
     torch.cuda.empty_cache()
 
+    progress(1.00, desc="✅ Both pipelines complete!")
+
     return out_standard, out_small, seed
 
 
@@ -347,7 +325,9 @@ with gr.Blocks() as demo:
         )
         gr.Markdown(
             "Compare **FLUX.2-klein-4B** side-by-side with "
-            "[small decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder)."
+            "[small decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder). "
+            "Both pipelines run **one after the other** using the **same seed and latents** — "
+            "only the VAE decoder differs."
         )
 
         with gr.Row(equal_height=True):
@@ -375,7 +355,7 @@ with gr.Blocks() as demo:
                 with gr.Row():
                     with gr.Column():
                         result_standard = gr.Image(
-                            label="Standard Decoder",
+                            label="① Standard Decoder  (runs first)",
                             show_label=True,
                             interactive=False,
                             format="png",
@@ -383,7 +363,7 @@ with gr.Blocks() as demo:
                         )
                     with gr.Column():
                         result_small = gr.Image(
-                            label="Small Decoder",
+                            label="② Small Decoder  (runs second)",
                             show_label=True,
                             interactive=False,
                             format="png",
@@ -392,7 +372,7 @@ with gr.Blocks() as demo:
 
                 seed_output = gr.Number(label="Seed Used", precision=0, visible=False)
 
-                with gr.Accordion("Advanced Settings", open=False, visible=False):
+                with gr.Accordion("Advanced Settings", open=False):
                     seed = gr.Slider(
                         label="Seed",
                         minimum=0,
@@ -478,7 +458,8 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch(
-        theme=orange_red_theme, css=css,
+        theme=orange_red_theme,
+        css=css,
         mcp_server=True,
         ssr_mode=False,
         show_error=True,