update app

app.py

@@ -1,4 +1,3 @@
-
 import os
 import gc
 import gradio as gr
@@ -6,7 +5,6 @@ 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
@@ -87,8 +85,8 @@ orange_red_theme = OrangeRedTheme()
 
 # ── Config ────────────────────────────────────────────────────────────────────
 
-dtype   = torch.bfloat16
-device  = "cuda" if torch.cuda.is_available() else "cpu"
+dtype  = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 MAX_SEED       = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -166,15 +164,6 @@ def parse_input_images(input_images):
     return None
 
 
-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=240)
@@ -214,15 +203,11 @@ def infer(
     progress(0.05, desc="🟦 Running Standard VAE generation...")
     print("Starting Standard VAE generation...")
 
-    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)
+    gen_std      = torch.Generator(device="cpu").manual_seed(seed)
+    out_standard = pipe_standard(**shared_kwargs, generator=gen_std).images[0]
 
-    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...")
+    print("Standard VAE generation complete.")
+    progress(0.55, desc="🟦 Standard VAE done — now running 🟩 Small Decoder VAE...")
 
     gc.collect()
     torch.cuda.empty_cache()
@@ -230,24 +215,16 @@ def infer(
     # ── Step 2: Small Decoder VAE ─────────────────────────────────────────────
     print("Starting Small Decoder VAE generation...")
 
-    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)
+    gen_small = torch.Generator(device="cpu").manual_seed(seed)
+    out_small = pipe_small_decoder(**shared_kwargs, generator=gen_small).images[0]
 
-    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}")
+    print("Small Decoder VAE generation complete.")
+    progress(1.0, desc="✅ Both generations complete!")
 
     gc.collect()
     torch.cuda.empty_cache()
 
-    # ── 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
+    return out_standard, out_small, seed
 
 
 @spaces.GPU(duration=240)
@@ -259,7 +236,7 @@ def infer_example(images, prompt):
     else:
         images_list = images
 
-    out_std, out_small, seed_used, label_std, label_small = infer(
+    out_std, out_small, seed_used = infer(
         prompt=prompt,
         input_images=images_list,
         seed=0,
@@ -269,7 +246,7 @@ def infer_example(images, prompt):
         num_inference_steps=4,
         guidance_scale=1.0,
     )
-    return out_std, out_small, seed_used, label_std, label_small
+    return out_std, out_small, seed_used
 
 
 # ── CSS ───────────────────────────────────────────────────────────────────────
@@ -291,15 +268,13 @@ css = """
     display: block;
     margin-bottom: 6px;
 }
-.timing-label {
+.output-order-note {
     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);
+    font-size: 0.85em;
+    color: #777;
+    font-style: italic;
+    margin-top: 2px;
+    margin-bottom: 4px;
 }
 """
 
@@ -310,12 +285,13 @@ with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
 
         gr.Markdown(
-            "# **Flux.2-4B-Encoder-Comparator**",
+            "# ⚡ **Flux.2-4B-Encoder-Comparator**",
             elem_id="main-title",
         )
         gr.Markdown(
-            "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)]"
+            "Compare **FLUX.2-klein-4B** side-by-side with two VAE decoders — generated **sequentially** from the **same seed**.\n\n"
+            "🟦 **Standard VAE** is generated **first**, then 🟩 **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 ───────────────────────────────────────────────
@@ -343,12 +319,14 @@ with gr.Blocks() as demo:
             # ── Right: outputs ────────────────────────────────────────────────
             with gr.Column():
                 with gr.Row():
-                    # Standard VAE output
+
+                    # Standard VAE — Generated First
                     with gr.Column():
                         gr.HTML(
                             '<span class="vae-badge" '
                             'style="background:#FFE0CC;color:#CC3700;">'
-                            '🟦 Standard VAE &nbsp;·&nbsp; Generated First</span>'
+                            '🟦 Standard VAE</span>'
+                            '<p class="output-order-note">① Generated First</p>'
                         )
                         result_standard = gr.Image(
                             label="Standard VAE",
@@ -357,17 +335,14 @@ with gr.Blocks() as demo:
                             format="png",
                             height=280,
                         )
-                        timing_standard = gr.Markdown(
-                            value="⏱ Waiting...",
-                            elem_classes=["timing-label"],
-                        )
 
-                    # Small Decoder VAE output
+                    # Small Decoder VAE — Generated Second
                     with gr.Column():
                         gr.HTML(
                             '<span class="vae-badge" '
                             'style="background:#FFF0E5;color:#E63E00;">'
-                            '🟩 Small Decoder VAE &nbsp;·&nbsp; Generated Second</span>'
+                            '🟩 Small Decoder VAE</span>'
+                            '<p class="output-order-note">② Generated Second</p>'
                         )
                         result_small = gr.Image(
                             label="Small Decoder VAE",
@@ -376,10 +351,6 @@ 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):
                     seed_output = gr.Number(label="Seed Used", precision=0)
@@ -431,10 +402,9 @@ with gr.Blocks() as demo:
                 [["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, timing_standard, timing_small],
+            outputs=[result_standard, result_small, seed_output],
             fn=infer_example,
             cache_examples=False,
             label="Examples",
@@ -466,7 +436,7 @@ with gr.Blocks() as demo:
             num_inference_steps,
             guidance_scale,
         ],
-        outputs=[result_standard, result_small, seed_output, timing_standard, timing_small],
+        outputs=[result_standard, result_small, seed_output],
     )
 
 if __name__ == "__main__":