Create app.py

app.py

@@ -0,0 +1,501 @@
+import os
+import gc
+import gradio as gr
+import numpy as np
+import random
+import spaces
+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
+from gradio.themes.utils import colors, fonts, sizes
+
+# ── Theme ─────────────────────────────────────────────────────────────────────
+
+colors.flux_teal = colors.Color(
+    name="flux_teal",
+    c50="#E6FAF8",
+    c100="#CCF5F1",
+    c200="#99EBE3",
+    c300="#66E1D5",
+    c400="#33D7C7",
+    c500="#00CDB9",
+    c600="#00A494",
+    c700="#007B6F",
+    c800="#00524A",
+    c900="#002925",
+    c950="#001412",
+)
+
+
+class FluxTheme(Soft):
+    def __init__(
+        self,
+        *,
+        primary_hue: colors.Color | str = colors.slate,
+        secondary_hue: colors.Color | str = colors.flux_teal,
+        neutral_hue: colors.Color | str = colors.slate,
+        text_size: sizes.Size | str = sizes.text_lg,
+        font: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("Inter"),
+            "Arial",
+            "sans-serif",
+        ),
+        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("JetBrains Mono"),
+            "ui-monospace",
+            "monospace",
+        ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            text_size=text_size,
+            font=font,
+            font_mono=font_mono,
+        )
+        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_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            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",
+        )
+
+
+flux_theme = FluxTheme()
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+dtype   = torch.bfloat16
+device  = "cuda" if torch.cuda.is_available() else "cpu"
+
+MAX_SEED       = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+EXAMPLES_DIR   = Path("examples")
+
+# ── Models ────────────────────────────────────────────────────────────────────
+
+print("Loading 4B Distilled model (Standard VAE)...")
+pipe_standard = Flux2KleinPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-4B",
+    torch_dtype=dtype,
+)
+pipe_standard.enable_model_cpu_offload()
+
+print("Loading Small Decoder VAE...")
+vae_small = AutoencoderKLFlux2.from_pretrained(
+    "black-forest-labs/FLUX.2-small-decoder",
+    torch_dtype=dtype,
+)
+
+print("Loading 4B Distilled model (Small Decoder VAE)...")
+pipe_small_decoder = Flux2KleinPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-4B",
+    vae=vae_small,
+    torch_dtype=dtype,
+)
+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):
+    if image_list is None or len(image_list) == 0:
+        return 1024, 1024
+
+    item = image_list[0]
+    img  = item[0] if isinstance(item, tuple) else item
+
+    if isinstance(img, str):
+        img = Image.open(img).convert("RGB")
+
+    iw, ih       = img.size
+    aspect_ratio = iw / ih
+
+    if aspect_ratio >= 1:
+        new_width  = 1024
+        new_height = int(1024 / aspect_ratio)
+    else:
+        new_height = 1024
+        new_width  = int(1024 * aspect_ratio)
+
+    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 get_example_items():
+    example_prompts = {
+        "1.jpg": "Change the weather to stormy.",
+        "2.jpg": "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition.",
+        "3.jpg": "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent.",
+        "4.jpg": "Make the texture high-resolution.",
+    }
+    items = []
+    if EXAMPLES_DIR.exists():
+        for name in sorted(os.listdir(EXAMPLES_DIR)):
+            if name.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
+                items.append({
+                    "file":   name,
+                    "path":   str(EXAMPLES_DIR / name),
+                    "prompt": example_prompts.get(
+                        name, "Edit this image while preserving composition."
+                    ),
+                })
+    return items
+
+
+def parse_input_images(input_images):
+    """Safely parse gallery / filepath / PIL inputs → list[PIL.Image] or None."""
+    if input_images is None:
+        return None
+    if isinstance(input_images, str):
+        return [Image.open(input_images).convert("RGB")] if os.path.exists(input_images) else None
+    if isinstance(input_images, list) and len(input_images) > 0:
+        parsed = []
+        for item in input_images:
+            try:
+                src = item[0] if isinstance(item, tuple) else item
+                if isinstance(src, str):
+                    parsed.append(Image.open(src).convert("RGB"))
+                elif isinstance(src, Image.Image):
+                    parsed.append(src.convert("RGB"))
+                elif hasattr(src, "name"):
+                    parsed.append(Image.open(src.name).convert("RGB"))
+            except Exception as e:
+                print(f"Skipping invalid image: {e}")
+        return parsed or None
+    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
+
+
+# ── Inference ─────────────────────────────────────────────────────────────────
+
+@spaces.GPU(duration=120)
+def infer(
+    prompt,
+    input_images=None,
+    seed=42,
+    randomize_seed=False,
+    width=1024,
+    height=1024,
+    num_inference_steps=4,
+    guidance_scale=1.0,
+    progress=gr.Progress(track_tqdm=True),
+):
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    if not prompt or prompt.strip() == "":
+        raise gr.Error("Please enter a prompt.")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    image_list = parse_input_images(input_images)
+
+    shared_kwargs = dict(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+    )
+    if image_list is not None:
+        shared_kwargs["image"] = image_list
+
+    progress(0.05, desc="⚡ Launching both pipelines simultaneously...")
+
+    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.95, desc="✅ Both pipelines done!")
+
+    out_standard = future_std.result()
+    out_small    = future_small.result()
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return out_standard, out_small, seed
+
+
+# Text-only wrapper used by gr.Examples (avoids Gallery type issues entirely)
+@spaces.GPU(duration=120)
+def infer_example(prompt):
+    out_std, out_small, seed_used = infer(
+        prompt=prompt,
+        input_images=None,
+        seed=0,
+        randomize_seed=True,
+        width=1024,
+        height=1024,
+        num_inference_steps=4,
+        guidance_scale=1.0,
+    )
+    return out_std, out_small, seed_used
+
+
+# ── Load examples ─────────────────────────────────────────────────────────────
+
+EXAMPLE_ITEMS = get_example_items()
+
+# ── CSS ───────────────────────────────────────────────────────────────────────
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1100px;
+}
+#main-title h1 {
+    font-size: 2.4em !important;
+}
+.vae-badge {
+    font-weight: 700;
+    font-size: 0.95em;
+    text-align: center;
+    padding: 4px 16px;
+    border-radius: 20px;
+    display: block;
+    margin-bottom: 6px;
+}
+"""
+
+# ── UI ────────────────────────────────────────────────────────────────────────
+
+with gr.Blocks() as demo:
+
+    with gr.Column(elem_id="col-container"):
+
+        gr.Markdown(
+            "# ⚡ **Flux.2-4B-Encoder-Comparator**",
+            elem_id="main-title",
+        )
+        gr.Markdown(
+            "Compare **FLUX.2-klein-4B** side-by-side with two VAE decoders — "
+            "generated **simultaneously** from the **same seed**. "
+            "🟦 **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)] · "
+            "[[blog](https://bfl.ai/blog/flux-2)]"
+        )
+
+        # ── Main two-column row ───────────────────────────────────────────────
+        with gr.Row(equal_height=True):
+
+            # ── Left: inputs ─────────────────────────────────────────────────
+            with gr.Column():
+                input_images = gr.Gallery(
+                    label="Input Image(s) for Editing (optional)",
+                    type="pil",
+                    columns=2,
+                    rows=1,
+                    height=280,
+                    allow_preview=True,
+                )
+
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=True,
+                    placeholder="e.g., A black cat holding a sign that says hello world...",
+                )
+
+                run_button = gr.Button("⚡ Run Comparison", variant="primary")
+
+            # ── Right: outputs ────────────────────────────────────────────────
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column():
+                        gr.HTML(
+                            '<span class="vae-badge" '
+                            'style="background:#dbeafe;color:#1d4ed8;">'
+                            '🟦 Standard VAE</span>'
+                        )
+                        result_standard = gr.Image(
+                            label="Standard VAE",
+                            show_label=False,
+                            interactive=False,
+                            format="png",
+                            height=280,
+                        )
+
+                    with gr.Column():
+                        gr.HTML(
+                            '<span class="vae-badge" '
+                            'style="background:#d1fae5;color:#065f46;">'
+                            '🟩 Small Decoder VAE</span>'
+                        )
+                        result_small = gr.Image(
+                            label="Small Decoder VAE",
+                            show_label=False,
+                            interactive=False,
+                            format="png",
+                            height=280,
+                        )
+
+                seed_output = gr.Number(label="Seed Used", precision=0)
+
+                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)
+
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
+                        height_slider = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
+
+                    with gr.Row():
+                        num_inference_steps = gr.Slider(
+                            label="Inference Steps",
+                            minimum=1,
+                            maximum=20,
+                            step=1,
+                            value=4,
+                        )
+                        guidance_scale = gr.Slider(
+                            label="Guidance Scale",
+                            minimum=0.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=1.0,
+                        )
+
+        # ── Examples — prompt-only, no Gallery input ──────────────────────────
+        # Build rows: [prompt_str]  (no image column → no Gallery postprocess)
+        if EXAMPLE_ITEMS:
+            prompt_only_rows = [[item["prompt"]] for item in EXAMPLE_ITEMS]
+
+            gr.Examples(
+                examples=prompt_only_rows,
+                inputs=[prompt],                          # ← only Text, never Gallery
+                outputs=[result_standard, result_small, seed_output],
+                fn=infer_example,                         # ← wrapper with no image arg
+                cache_examples=False,
+                label="Examples",
+            )
+
+        # ── Visual image cards (click to load image + prompt) ─────────────────
+        if EXAMPLE_ITEMS:
+            gr.Markdown("#### 🖼️ Image Editing Examples — click to load")
+            with gr.Row():
+                for item in EXAMPLE_ITEMS:
+                    with gr.Column(scale=1, min_width=180):
+                        gr.Image(
+                            value=item["path"],
+                            show_label=False,
+                            interactive=False,
+                            height=150,
+                        )
+                        card_btn = gr.Button(
+                            (item["prompt"][:48] + "…")
+                            if len(item["prompt"]) > 48
+                            else item["prompt"],
+                            size="sm",
+                        )
+
+                        def _make_loader(p, path):
+                            def _load():
+                                pil = Image.open(path).convert("RGB")
+                                return p, [(pil, None)]
+                            return _load
+
+                        card_btn.click(
+                            fn=_make_loader(item["prompt"], item["path"]),
+                            inputs=[],
+                            outputs=[prompt, input_images],
+                        )
+
+        gr.Markdown(
+            "[*](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) "
+            "Experimental Space — FLUX.2 [klein] 4B VAE Decoder Comparison."
+        )
+
+    # ── Events ────────────────────────────────────────────────────────────────
+
+    input_images.upload(
+        fn=update_dimensions_from_image,
+        inputs=[input_images],
+        outputs=[width, height_slider],
+    )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[
+            prompt,
+            input_images,
+            seed,
+            randomize_seed,
+            width,
+            height_slider,
+            num_inference_steps,
+            guidance_scale,
+        ],
+        outputs=[result_standard, result_small, seed_output],
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(
+        css=css,
+        theme=flux_theme,
+        ssr_mode=False,
+        show_error=True,
+    )