Delete app.py

app.py

@@ -1,466 +0,0 @@
-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
-from typing import Iterable
-
-from gradio.themes import Soft
-from gradio.themes.utils import colors, fonts, sizes
-
-colors.orange_red = colors.Color(
-    name="orange_red",
-    c50="#FFF0E5",
-    c100="#FFE0CC",
-    c200="#FFC299",
-    c300="#FFA366",
-    c400="#FF8533",
-    c500="#FF4500",
-    c600="#E63E00",
-    c700="#CC3700",
-    c800="#B33000",
-    c900="#992900",
-    c950="#802200",
-)
-
-class OrangeRedTheme(Soft):
-    def __init__(
-        self,
-        *,
-        primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red,
-        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("Outfit"), "Arial", "sans-serif",
-        ),
-        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("IBM Plex 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",
-        )
-
-orange_red_theme = OrangeRedTheme()
-
-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")
-
-# ── Load standard pipeline ──────────────────────────────────────────────────
-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()
-
-# ── 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",
-    vae=vae_small,
-    torch_dtype=dtype,
-)
-pipe_small_decoder.enable_model_cpu_offload()
-
-# ────────────────────────────────────────────────────────────────────────────
-def calc_dimensions(pil_img: Image.Image):
-    iw, ih = pil_img.size
-    aspect = iw / ih
-
-    if aspect >= 1:
-        new_width  = 1024
-        new_height = int(round(1024 / aspect))
-    else:
-        new_height = 1024
-        new_width  = int(round(1024 * aspect))
-
-    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):
-    if not image_list:
-        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")
-    elif not isinstance(img, Image.Image):
-        return 1024, 1024
-
-    return calc_dimensions(img)
-
-
-def parse_and_resize_images(input_images, width: int, height: int):
-    if input_images is None:
-        return None
-
-    raw_list = []
-
-    if isinstance(input_images, str):
-        if os.path.exists(input_images):
-            raw_list = [Image.open(input_images).convert("RGB")]
-    elif isinstance(input_images, Image.Image):
-        raw_list = [input_images.convert("RGB")]
-    elif isinstance(input_images, list):
-        for item in input_images:
-            try:
-                src = item[0] if isinstance(item, tuple) else item
-                if isinstance(src, str):
-                    raw_list.append(Image.open(src).convert("RGB"))
-                elif isinstance(src, Image.Image):
-                    raw_list.append(src.convert("RGB"))
-                elif hasattr(src, "name"):
-                    raw_list.append(Image.open(src.name).convert("RGB"))
-            except Exception as e:
-                print(f"Skipping invalid image: {e}")
-
-    if not raw_list:
-        return None
-
-    resized = [
-        img.resize((width, height), Image.LANCZOS)
-        for img in raw_list
-    ]
-    return resized
-
-
-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,
-    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 not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    # ── Derive dimensions from the first uploaded image if present ───────────
-    image_list = None
-    if input_images:
-        item = (
-            input_images[0][0]
-            if isinstance(input_images[0], tuple)
-            else input_images[0]
-        )
-        if isinstance(item, str):
-            first_pil = Image.open(item).convert("RGB")
-        elif isinstance(item, Image.Image):
-            first_pil = item.convert("RGB")
-        else:
-            first_pil = None
-
-        if first_pil is not None:
-            width, height = calc_dimensions(first_pil)
-
-        image_list = parse_and_resize_images(input_images, width, height)
-
-    # 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))
-
-    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
-
-    # ── Pipeline 1: Standard Decoder ─────────────────────────────────────────
-    progress(0.10, desc="Running Pipeline 1 / 2 — Standard Decoder...")
-    out_standard = run_pipeline(pipe_standard, shared_kwargs, seed)
-
-    gc.collect()
-    torch.cuda.empty_cache()
-
-    # ── 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
-
-
-@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
-
-
-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
-
-EXAMPLE_ITEMS = get_example_items()
-
-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;
-}
-"""
-
-with gr.Blocks() as demo:
-
-    with gr.Column(elem_id="col-container"):
-
-        gr.Markdown(
-            "# **Flux.2-4B-Decoder-Comparator**",
-            elem_id="main-title",
-        )
-        gr.Markdown(
-            "Compare **FLUX.2-klein-4B** side-by-side with "
-            "[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):
-
-            with gr.Column():
-                input_images = gr.Gallery(
-                    label="Input Images",
-                    type="pil",
-                    columns=2,
-                    rows=1,
-                    height=300,
-                    allow_preview=True,
-                )
-
-                prompt = gr.Text(
-                    label="Prompt",
-                    max_lines=1,
-                    show_label=True,
-                    placeholder="e.g., A black cat holding a sign that says hello world...",
-                )
-
-                run_button = gr.Button("Run Comparison", variant="primary")
-
-            with gr.Column():
-                with gr.Row():
-                    with gr.Column():
-                        result_standard = gr.Image(
-                            label="① Standard Decoder  (runs first)",
-                            show_label=True,
-                            interactive=False,
-                            format="png",
-                            height=250,
-                        )
-                    with gr.Column():
-                        result_small = gr.Image(
-                            label="② Small Decoder  (runs second)",
-                            show_label=True,
-                            interactive=False,
-                            format="png",
-                            height=250,
-                        )
-
-                seed_output = gr.Number(label="Seed Used", precision=0, visible=False)
-
-                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,
-                        )
-
-        gr.Examples(
-            examples=[
-                [["examples/I1.jpg", "examples/I2.jpg"], "Make her wear these glasses in Image 2."],
-                [["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."],
-            ],
-            inputs=[input_images, prompt],
-            outputs=[result_standard, result_small, seed_output],
-            fn=infer_example,
-            cache_examples=False,
-            label="Examples",
-        )
-
-        gr.Markdown(
-            "[*](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) "
-            "Experimental Space — FLUX.2 [klein] 4B VAE Decoder Comparison."
-        )
-
-    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(
-        theme=orange_red_theme,
-        css=css,
-        mcp_server=True,
-        ssr_mode=False,
-        show_error=True,
-    )