Add class-conditional BiliSakura diffusers playground UI.

Unified model selector, inference config widgets, and ZeroGPU-ready generation for all *-diffusers checkpoints.

README.md

@@ -1,6 +1,6 @@
 ---
 title: Visual Generative Model Playground
-emoji: 🏃
+emoji: 🎨
 colorFrom: blue
 colorTo: indigo
 sdk: gradio
@@ -8,6 +8,19 @@ sdk_version: 6.15.2
 python_version: '3.12'
 app_file: app.py
 pinned: false
+short_description: Class-conditional BiliSakura diffusers playground on ZeroGPU
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Visual Generative Model Playground
+
+Class-conditional image generation demo for [`BiliSakura/*-diffusers`](https://huggingface.co/BiliSakura) on **ZeroGPU**.
+
+Select a model, configure inference args (`class_labels`, `num_inference_steps`, `guidance_scale`, resolution), and generate.
+
+## Hardware
+
+Set Space hardware to **ZeroGPU** in Settings.
+
+## Models
+
+ADM, DiT, DiT-MoE, EDM2, FiT, iMF, JiT, LightningDiT, NiT, PixelFlow, PixNerd, pMF, Self-Flow, and SiT (47 variants).

app.py

@@ -0,0 +1,405 @@
+"""Gradio demo for BiliSakura visual generative foundation models on Hugging Face ZeroGPU."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import gradio as gr
+import torch
+
+from model_catalog import (
+    COLLECTIONS,
+    MODEL_LABELS,
+    get_profile_by_label,
+    parse_model_label,
+)
+from model_loader import PIPELINE_MANAGER, run_inference
+
+try:
+    import spaces
+except ImportError:  # Local development without the Spaces runtime.
+    class _SpacesStub:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            def decorator(fn):
+                return fn
+
+            if args and callable(args[0]):
+                return args[0]
+            return decorator
+
+    spaces = _SpacesStub()  # type: ignore[assignment]
+
+
+DEFAULT_MODEL = MODEL_LABELS[0]
+DEFAULT_PROFILE = get_profile_by_label(DEFAULT_MODEL)
+
+INTERVAL_COLLECTIONS = {"iMF-diffusers", "NiT-diffusers", "PixelFlow-diffusers"}
+PMF_COLLECTION = "pMF-diffusers"
+
+
+def _model_info_markdown(profile) -> str:
+    extras = profile.extra_call_kwargs
+    extra_lines = ""
+    if extras:
+        extra_lines = "\n".join(f"- `{key}`: `{value}`" for key, value in extras.items())
+        extra_lines = f"\n\n**Default extra args**\n{extra_lines}"
+    return (
+        f"**Hub repo:** [`{profile.hub_model_id}`](https://huggingface.co/{profile.hub_model_id})\n\n"
+        f"- dtype: `{profile.dtype}`\n"
+        f"- default resolution: `{profile.default_height}x{profile.default_width}`\n"
+        f"- GPU size: `{profile.gpu_size}`"
+        f"{extra_lines}"
+    )
+
+
+def _interval_defaults(profile) -> tuple[float, float]:
+    extras = profile.extra_call_kwargs
+    if "guidance_interval_start" in extras:
+        return float(extras["guidance_interval_start"]), float(extras["guidance_interval_end"])
+    interval = extras.get("guidance_interval", (0.0, 0.7))
+    return float(interval[0]), float(interval[1])
+
+
+def _build_extra_kwargs(
+    profile,
+    guidance_interval_start: float,
+    guidance_interval_end: float,
+    guidance_interval_min: float,
+    guidance_interval_max: float,
+    noise_scale: float,
+) -> dict[str, Any]:
+    if profile.collection == "iMF-diffusers":
+        return {
+            "guidance_interval_start": guidance_interval_start,
+            "guidance_interval_end": guidance_interval_end,
+        }
+    if profile.collection in {"NiT-diffusers", "PixelFlow-diffusers"}:
+        return {"guidance_interval": (guidance_interval_start, guidance_interval_end)}
+    if profile.collection == PMF_COLLECTION:
+        return {
+            "guidance_interval_min": guidance_interval_min,
+            "guidance_interval_max": guidance_interval_max,
+            "noise_scale": noise_scale,
+        }
+    return dict(profile.extra_call_kwargs)
+
+
+def _config_from_profile(profile):
+    g_start, g_end = _interval_defaults(profile)
+    extras = profile.extra_call_kwargs
+    show_interval = profile.collection in INTERVAL_COLLECTIONS
+    show_pmf = profile.collection == PMF_COLLECTION
+    return (
+        _model_info_markdown(profile),
+        gr.update(value=profile.default_class_label),
+        gr.update(value=profile.default_seed),
+        gr.update(value=profile.default_steps, maximum=profile.max_steps),
+        gr.update(value=profile.default_guidance),
+        gr.update(value=profile.default_height or profile.infer_resolution()),
+        gr.update(value=profile.default_width or profile.infer_resolution()),
+        gr.update(value=g_start),
+        gr.update(value=g_end),
+        gr.update(value=float(extras.get("guidance_interval_min", 0.2))),
+        gr.update(value=float(extras.get("guidance_interval_max", 0.6))),
+        gr.update(value=float(extras.get("noise_scale", 4.0))),
+        gr.update(visible=show_interval, open=show_interval),
+        gr.update(visible=show_pmf, open=show_pmf),
+    )
+
+
+def on_model_change(model_label: str):
+    return _config_from_profile(get_profile_by_label(model_label))
+
+
+def _gpu_duration(
+    model_label: str,
+    num_steps: int,
+    guidance_scale: float,
+    seed: int,
+    class_label: str,
+    height: int,
+    width: int,
+    guidance_interval_start: float,
+    guidance_interval_end: float,
+    guidance_interval_min: float,
+    guidance_interval_max: float,
+    noise_scale: float,
+) -> int:
+    profile = get_profile_by_label(model_label)
+    step_budget = num_steps if not profile.steps_are_list else max(num_steps, 40)
+    base = 45 if profile.gpu_size == "large" else 90
+    return int(min(300, max(base, step_budget * 0.6 + 30)))
+
+
+def _load_model_core(model_label: str) -> str:
+    collection, variant = parse_model_label(model_label)
+    message, _ = PIPELINE_MANAGER.load(collection, variant)
+    return message
+
+
+def load_model(model_label: str):
+    try:
+        message = _load_model_core(model_label)
+    except Exception as exc:
+        raise gr.Error(f"Failed to load `{model_label}`: {exc}") from exc
+    return (message, *_config_from_profile(get_profile_by_label(model_label)))
+
+
+@spaces.GPU(size="xlarge", duration=_gpu_duration)
+def _generate_on_gpu(
+    model_label: str,
+    class_label: str,
+    seed: int,
+    num_steps: int,
+    guidance_scale: float,
+    height: int,
+    width: int,
+    guidance_interval_start: float,
+    guidance_interval_end: float,
+    guidance_interval_min: float,
+    guidance_interval_max: float,
+    noise_scale: float,
+):
+    profile = get_profile_by_label(model_label)
+    pipe = PIPELINE_MANAGER.pipe
+    if pipe is None or PIPELINE_MANAGER.loaded_label != model_label:
+        raise gr.Error(f"Model `{model_label}` is not loaded.")
+
+    extra_kwargs = _build_extra_kwargs(
+        profile,
+        guidance_interval_start,
+        guidance_interval_end,
+        guidance_interval_min,
+        guidance_interval_max,
+        noise_scale,
+    )
+    return run_inference(
+        profile,
+        pipe,
+        class_label=class_label,
+        seed=seed,
+        num_steps=num_steps,
+        guidance_scale=guidance_scale,
+        height=height,
+        width=width,
+        extra_kwargs=extra_kwargs,
+    )
+
+
+def generate(
+    model_label: str,
+    class_label: str,
+    seed: int,
+    num_steps: int,
+    guidance_scale: float,
+    height: int,
+    width: int,
+    guidance_interval_start: float,
+    guidance_interval_end: float,
+    guidance_interval_min: float,
+    guidance_interval_max: float,
+    noise_scale: float,
+):
+    try:
+        status = _load_model_core(model_label)
+    except Exception as exc:
+        raise gr.Error(f"Failed to load `{model_label}`: {exc}") from exc
+    image = _generate_on_gpu(
+        model_label,
+        class_label,
+        seed,
+        num_steps,
+        guidance_scale,
+        height,
+        width,
+        guidance_interval_start,
+        guidance_interval_end,
+        guidance_interval_min,
+        guidance_interval_max,
+        noise_scale,
+    )
+    return status, image
+
+
+def build_demo() -> gr.Blocks:
+    g_start, g_end = _interval_defaults(DEFAULT_PROFILE)
+    extras = DEFAULT_PROFILE.extra_call_kwargs
+
+    with gr.Blocks(title="BiliSakura Visual Generation Models") as demo:
+        gr.Markdown(
+            """
+            # BiliSakura Visual Generative Foundation Models
+
+            Class-conditional image generation for [`BiliSakura/*-diffusers`](https://huggingface.co/BiliSakura)
+            on Hugging Face **ZeroGPU**.
+            """
+        )
+
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=5):
+                model = gr.Dropdown(
+                    MODEL_LABELS,
+                    value=DEFAULT_MODEL,
+                    label="Model",
+                    info="Select a checkpoint, then configure inference args below",
+                )
+                model_info = gr.Markdown(_model_info_markdown(DEFAULT_PROFILE))
+
+                with gr.Accordion("Inference config", open=True):
+                    class_label = gr.Textbox(
+                        label="class_labels",
+                        value=DEFAULT_PROFILE.default_class_label,
+                        info="ImageNet class name (e.g. golden retriever) or id (e.g. 207)",
+                    )
+                    with gr.Row():
+                        seed = gr.Number(label="seed", value=DEFAULT_PROFILE.default_seed, precision=0)
+                        num_steps = gr.Slider(
+                            label="num_inference_steps",
+                            minimum=1,
+                            maximum=DEFAULT_PROFILE.max_steps,
+                            step=1,
+                            value=DEFAULT_PROFILE.default_steps,
+                        )
+                    guidance_scale = gr.Slider(
+                        label="guidance_scale",
+                        minimum=0.0,
+                        maximum=20.0,
+                        step=0.1,
+                        value=DEFAULT_PROFILE.default_guidance,
+                    )
+                    with gr.Row():
+                        height = gr.Slider(
+                            label="height",
+                            minimum=128,
+                            maximum=1024,
+                            step=16,
+                            value=DEFAULT_PROFILE.default_height or 256,
+                        )
+                        width = gr.Slider(
+                            label="width",
+                            minimum=128,
+                            maximum=1024,
+                            step=16,
+                            value=DEFAULT_PROFILE.default_width or 256,
+                        )
+
+                with gr.Accordion(
+                    "Advanced: guidance interval",
+                    open=False,
+                    visible=DEFAULT_PROFILE.collection in INTERVAL_COLLECTIONS,
+                ) as interval_accordion:
+                    with gr.Row():
+                        guidance_interval_start = gr.Slider(
+                            label="guidance_interval_start / [0]",
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.05,
+                            value=g_start,
+                        )
+                        guidance_interval_end = gr.Slider(
+                            label="guidance_interval_end / [1]",
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.05,
+                            value=g_end,
+                        )
+
+                with gr.Accordion(
+                    "Advanced: pMF args",
+                    open=False,
+                    visible=DEFAULT_PROFILE.collection == PMF_COLLECTION,
+                ) as pmf_accordion:
+                    with gr.Row():
+                        guidance_interval_min = gr.Slider(
+                            label="guidance_interval_min",
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.05,
+                            value=float(extras.get("guidance_interval_min", 0.2)),
+                        )
+                        guidance_interval_max = gr.Slider(
+                            label="guidance_interval_max",
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.05,
+                            value=float(extras.get("guidance_interval_max", 0.6)),
+                        )
+                    noise_scale = gr.Slider(
+                        label="noise_scale",
+                        minimum=0.0,
+                        maximum=10.0,
+                        step=0.1,
+                        value=float(extras.get("noise_scale", 4.0)),
+                    )
+
+                with gr.Row():
+                    load_btn = gr.Button("Load model", variant="secondary")
+                    generate_btn = gr.Button("Generate", variant="primary")
+
+                status = gr.Textbox(label="Status", interactive=False, lines=2)
+                gr.Markdown(
+                    f"**Catalog:** {len(MODEL_LABELS)} variants · {len(COLLECTIONS)} families"
+                )
+
+            with gr.Column(scale=6):
+                output = gr.Image(
+                    label="Generated image",
+                    type="pil",
+                    height=720,
+                    elem_classes=["output-image"],
+                )
+                gr.Examples(
+                    examples=[
+                        [DEFAULT_MODEL, "golden retriever", 42],
+                        [DEFAULT_MODEL, "207", 0],
+                        [DEFAULT_MODEL, "tabby, tabby cat", 123],
+                    ],
+                    inputs=[model, class_label, seed],
+                    label="Quick examples",
+                )
+
+        inference_inputs = [
+            model,
+            class_label,
+            seed,
+            num_steps,
+            guidance_scale,
+            height,
+            width,
+            guidance_interval_start,
+            guidance_interval_end,
+            guidance_interval_min,
+            guidance_interval_max,
+            noise_scale,
+        ]
+        config_outputs = [
+            model_info,
+            class_label,
+            seed,
+            num_steps,
+            guidance_scale,
+            height,
+            width,
+            guidance_interval_start,
+            guidance_interval_end,
+            guidance_interval_min,
+            guidance_interval_max,
+            noise_scale,
+            interval_accordion,
+            pmf_accordion,
+        ]
+
+        model.change(on_model_change, inputs=model, outputs=config_outputs)
+        load_btn.click(load_model, inputs=model, outputs=[status, *config_outputs])
+        generate_btn.click(generate, inputs=inference_inputs, outputs=[status, output])
+        demo.load(on_model_change, inputs=model, outputs=config_outputs)
+
+    return demo
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        print("CUDA is not available locally; ZeroGPU Spaces will provide GPU at inference time.")
+    demo = build_demo()
+    demo.queue(max_size=8).launch()

model_catalog.py

@@ -0,0 +1,480 @@
+"""Catalog of BiliSakura *-diffusers models hosted on Hugging Face Hub."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+DtypeName = Literal["bfloat16", "float32"]
+GpuSize = Literal["large", "xlarge"]
+
+
+@dataclass(frozen=True)
+class ModelProfile:
+    collection: str
+    variant: str
+    dtype: DtypeName = "bfloat16"
+    use_custom_pipeline: bool = True
+    default_class_label: str = "golden retriever"
+    default_steps: int = 50
+    default_guidance: float = 4.0
+    default_height: int | None = None
+    default_width: int | None = None
+    default_seed: int = 42
+    gpu_size: GpuSize = "large"
+    scheduler: str | None = None
+    scheduler_kwargs: dict[str, Any] = field(default_factory=dict)
+    extra_call_kwargs: dict[str, Any] = field(default_factory=dict)
+    steps_are_list: bool = False
+    max_steps: int = 250
+
+    @property
+    def hub_repo(self) -> str:
+        return f"BiliSakura/{self.collection}"
+
+    @property
+    def hub_model_id(self) -> str:
+        return f"{self.hub_repo}/{self.variant}"
+
+    @property
+    def label(self) -> str:
+        return f"{self.collection}/{self.variant}"
+
+    def infer_resolution(self) -> int:
+        if self.default_height:
+            return self.default_height
+        name = self.variant.lower()
+        if "512" in name or "img512" in name:
+            return 512
+        if "1024" in name:
+            return 1024
+        return 256
+
+
+def _p(
+    collection: str,
+    variant: str,
+    *,
+    dtype: DtypeName = "bfloat16",
+    use_custom_pipeline: bool = True,
+    default_class_label: str = "golden retriever",
+    default_steps: int = 50,
+    default_guidance: float = 4.0,
+    default_height: int | None = None,
+    default_width: int | None = None,
+    gpu_size: GpuSize = "large",
+    scheduler: str | None = None,
+    scheduler_kwargs: dict[str, Any] | None = None,
+    extra_call_kwargs: dict[str, Any] | None = None,
+    steps_are_list: bool = False,
+    max_steps: int = 250,
+) -> ModelProfile:
+    if default_height is None:
+        name = variant.lower()
+        if "512" in name or "img512" in name:
+            res = 512
+        elif "1024" in name:
+            res = 1024
+        else:
+            res = 256
+        default_height = res
+        default_width = res
+
+    return ModelProfile(
+        collection=collection,
+        variant=variant,
+        dtype=dtype,
+        use_custom_pipeline=use_custom_pipeline,
+        default_class_label=default_class_label,
+        default_steps=default_steps,
+        default_guidance=default_guidance,
+        default_height=default_height,
+        default_width=default_width,
+        gpu_size=gpu_size,
+        scheduler=scheduler,
+        scheduler_kwargs=scheduler_kwargs or {},
+        extra_call_kwargs=extra_call_kwargs or {},
+        steps_are_list=steps_are_list,
+        max_steps=max_steps,
+    )
+
+
+MODEL_PROFILES: list[ModelProfile] = [
+    _p("ADM-diffusers", "ADM-G-256", default_steps=50, default_guidance=0.0, scheduler="DDIMScheduler"),
+    _p("ADM-diffusers", "ADM-G-512", default_steps=50, default_guidance=0.0, scheduler="DDIMScheduler"),
+    _p("DiT-diffusers", "DiT-XL-2-256", default_steps=250, default_guidance=4.0),
+    _p("DiT-diffusers", "DiT-XL-2-512", default_steps=250, default_guidance=4.0, gpu_size="xlarge"),
+    _p("DiT-MoE-diffusers", "DiT-MoE-S-8E2A", default_steps=50, default_guidance=4.0),
+    _p("DiT-MoE-diffusers", "DiT-MoE-B-8E2A", default_steps=50, default_guidance=4.0),
+    _p("DiT-MoE-diffusers", "DiT-MoE-XL-8E2A", default_steps=50, default_guidance=4.0, gpu_size="xlarge"),
+    _p(
+        "EDM2-diffusers",
+        "edm2-img512-xs-fid",
+        use_custom_pipeline=False,
+        default_steps=32,
+        default_guidance=1.0,
+        default_height=512,
+        default_width=512,
+    ),
+    _p(
+        "EDM2-diffusers",
+        "edm2-img512-s-fid",
+        use_custom_pipeline=False,
+        default_steps=32,
+        default_guidance=1.0,
+        default_height=512,
+        default_width=512,
+    ),
+    _p(
+        "EDM2-diffusers",
+        "edm2-img512-m-fid",
+        use_custom_pipeline=False,
+        default_steps=32,
+        default_guidance=1.0,
+        default_height=512,
+        default_width=512,
+    ),
+    _p(
+        "EDM2-diffusers",
+        "edm2-img512-l-fid",
+        use_custom_pipeline=False,
+        default_steps=32,
+        default_guidance=1.0,
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+    _p(
+        "EDM2-diffusers",
+        "edm2-img512-l-dino",
+        use_custom_pipeline=False,
+        default_steps=32,
+        default_guidance=1.0,
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+    _p(
+        "EDM2-diffusers",
+        "edm2-img512-xl-fid",
+        use_custom_pipeline=False,
+        default_steps=32,
+        default_guidance=1.0,
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+    _p(
+        "EDM2-diffusers",
+        "edm2-img512-xxl-fid",
+        use_custom_pipeline=False,
+        default_steps=32,
+        default_guidance=1.0,
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+    _p(
+        "FiT-diffusers",
+        "FiTv1-XL-2-256",
+        default_steps=50,
+        default_guidance=1.5,
+        scheduler="DDIMScheduler",
+    ),
+    _p(
+        "FiT-diffusers",
+        "FiTv2-XL-2-256",
+        default_steps=50,
+        default_guidance=1.5,
+        scheduler="FlowMatchEulerDiscreteScheduler",
+    ),
+    _p(
+        "FiT-diffusers",
+        "FiTv2-XL-2-512",
+        default_steps=50,
+        default_guidance=1.5,
+        scheduler="FlowMatchEulerDiscreteScheduler",
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+    _p(
+        "FiT-diffusers",
+        "FiTv2-3B-2-256",
+        default_steps=50,
+        default_guidance=1.5,
+        scheduler="FlowMatchEulerDiscreteScheduler",
+        gpu_size="xlarge",
+    ),
+    _p(
+        "FiT-diffusers",
+        "FiTv2-3B-2-512",
+        default_steps=50,
+        default_guidance=1.5,
+        scheduler="FlowMatchEulerDiscreteScheduler",
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+    _p(
+        "iMF-diffusers",
+        "iMF-B-2",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=1.8,
+        extra_call_kwargs={
+            "guidance_interval_start": 0.0,
+            "guidance_interval_end": 1.0,
+        },
+    ),
+    _p(
+        "iMF-diffusers",
+        "iMF-L-2",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=1.8,
+        extra_call_kwargs={
+            "guidance_interval_start": 0.0,
+            "guidance_interval_end": 1.0,
+        },
+    ),
+    _p(
+        "iMF-diffusers",
+        "iMF-XL-2",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=1.8,
+        extra_call_kwargs={
+            "guidance_interval_start": 0.0,
+            "guidance_interval_end": 1.0,
+        },
+    ),
+    _p(
+        "JiT-diffusers",
+        "JiT-B-16",
+        dtype="float32",
+        default_steps=250,
+        default_guidance=2.3,
+        scheduler="FlowMatchHeunDiscreteScheduler",
+        scheduler_kwargs={"shift": 4.0},
+    ),
+    _p(
+        "JiT-diffusers",
+        "JiT-B-32",
+        dtype="float32",
+        default_steps=250,
+        default_guidance=2.3,
+        scheduler="FlowMatchHeunDiscreteScheduler",
+        scheduler_kwargs={"shift": 4.0},
+    ),
+    _p(
+        "JiT-diffusers",
+        "JiT-L-16",
+        dtype="float32",
+        default_steps=250,
+        default_guidance=2.3,
+        scheduler="FlowMatchHeunDiscreteScheduler",
+        scheduler_kwargs={"shift": 4.0},
+    ),
+    _p(
+        "JiT-diffusers",
+        "JiT-L-32",
+        dtype="float32",
+        default_steps=250,
+        default_guidance=2.3,
+        scheduler="FlowMatchHeunDiscreteScheduler",
+        scheduler_kwargs={"shift": 4.0},
+    ),
+    _p(
+        "JiT-diffusers",
+        "JiT-H-16",
+        dtype="float32",
+        default_steps=250,
+        default_guidance=2.3,
+        scheduler="FlowMatchHeunDiscreteScheduler",
+        scheduler_kwargs={"shift": 4.0},
+        gpu_size="xlarge",
+    ),
+    _p(
+        "JiT-diffusers",
+        "JiT-H-32",
+        dtype="float32",
+        default_steps=250,
+        default_guidance=2.3,
+        scheduler="FlowMatchHeunDiscreteScheduler",
+        scheduler_kwargs={"shift": 4.0},
+        gpu_size="xlarge",
+    ),
+    _p(
+        "LightningDiT-diffusers",
+        "LightningDit-XL-1-256",
+        default_steps=50,
+        default_guidance=6.7,
+    ),
+    _p(
+        "NiT-diffusers",
+        "NiT-S",
+        default_steps=250,
+        default_guidance=2.25,
+        extra_call_kwargs={"guidance_interval": (0.0, 0.7)},
+    ),
+    _p(
+        "NiT-diffusers",
+        "NiT-B",
+        default_steps=250,
+        default_guidance=2.25,
+        extra_call_kwargs={"guidance_interval": (0.0, 0.7)},
+    ),
+    _p(
+        "NiT-diffusers",
+        "NiT-L",
+        default_steps=250,
+        default_guidance=2.25,
+        extra_call_kwargs={"guidance_interval": (0.0, 0.7)},
+    ),
+    _p(
+        "NiT-diffusers",
+        "NiT-XL",
+        default_steps=250,
+        default_guidance=2.25,
+        extra_call_kwargs={"guidance_interval": (0.0, 0.7)},
+        gpu_size="xlarge",
+    ),
+    _p(
+        "PixelFlow-diffusers",
+        "PixelFlow-256",
+        default_steps=40,
+        default_guidance=4.0,
+        steps_are_list=True,
+        extra_call_kwargs={"guidance_interval": (0.0, 0.7)},
+    ),
+    _p("PixNerd-diffusers", "PixNerd-XL-16-256", default_steps=25, default_guidance=4.0),
+    _p(
+        "PixNerd-diffusers",
+        "PixNerd-XL-16-512",
+        default_steps=25,
+        default_guidance=4.0,
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+    _p(
+        "pMF-diffusers",
+        "pMF-B-16",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=7.5,
+        extra_call_kwargs={
+            "guidance_interval_min": 0.2,
+            "guidance_interval_max": 0.6,
+            "noise_scale": 4.0,
+        },
+    ),
+    _p(
+        "pMF-diffusers",
+        "pMF-B-32",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=7.5,
+        extra_call_kwargs={
+            "guidance_interval_min": 0.2,
+            "guidance_interval_max": 0.6,
+            "noise_scale": 4.0,
+        },
+    ),
+    _p(
+        "pMF-diffusers",
+        "pMF-L-16",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=7.5,
+        extra_call_kwargs={
+            "guidance_interval_min": 0.2,
+            "guidance_interval_max": 0.6,
+            "noise_scale": 4.0,
+        },
+    ),
+    _p(
+        "pMF-diffusers",
+        "pMF-L-32",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=7.5,
+        extra_call_kwargs={
+            "guidance_interval_min": 0.2,
+            "guidance_interval_max": 0.6,
+            "noise_scale": 4.0,
+        },
+    ),
+    _p(
+        "pMF-diffusers",
+        "pMF-H-16",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=7.5,
+        extra_call_kwargs={
+            "guidance_interval_min": 0.2,
+            "guidance_interval_max": 0.6,
+            "noise_scale": 4.0,
+        },
+        gpu_size="xlarge",
+    ),
+    _p(
+        "pMF-diffusers",
+        "pMF-H-32",
+        dtype="float32",
+        default_steps=1,
+        default_guidance=7.5,
+        extra_call_kwargs={
+            "guidance_interval_min": 0.2,
+            "guidance_interval_max": 0.6,
+            "noise_scale": 4.0,
+        },
+        gpu_size="xlarge",
+    ),
+    _p("Self-Flow-diffusers", "Self-Flow-XL-2-256", default_steps=250, default_guidance=3.5),
+    _p("SiT-diffusers", "SiT-S-2-256", default_steps=250, default_guidance=4.0, scheduler="FlowMatchEulerDiscreteScheduler"),
+    _p("SiT-diffusers", "SiT-B-2-256", default_steps=250, default_guidance=4.0, scheduler="FlowMatchEulerDiscreteScheduler"),
+    _p("SiT-diffusers", "SiT-L-2-256", default_steps=250, default_guidance=4.0, scheduler="FlowMatchEulerDiscreteScheduler"),
+    _p("SiT-diffusers", "SiT-XL-2-256", default_steps=250, default_guidance=4.0, scheduler="FlowMatchEulerDiscreteScheduler"),
+    _p(
+        "SiT-diffusers",
+        "SiT-XL-2-512",
+        default_steps=250,
+        default_guidance=4.0,
+        scheduler="FlowMatchEulerDiscreteScheduler",
+        default_height=512,
+        default_width=512,
+        gpu_size="xlarge",
+    ),
+]
+
+
+PROFILE_BY_LABEL: dict[str, ModelProfile] = {profile.label: profile for profile in MODEL_PROFILES}
+COLLECTIONS: list[str] = sorted({profile.collection for profile in MODEL_PROFILES})
+VARIANTS_BY_COLLECTION: dict[str, list[str]] = {
+    collection: [profile.variant for profile in MODEL_PROFILES if profile.collection == collection]
+    for collection in COLLECTIONS
+}
+
+
+def get_profile(collection: str, variant: str) -> ModelProfile:
+    key = f"{collection}/{variant}"
+    if key not in PROFILE_BY_LABEL:
+        raise KeyError(f"Unknown model: {key}")
+    return PROFILE_BY_LABEL[key]
+
+
+def get_profile_by_label(label: str) -> ModelProfile:
+    if label not in PROFILE_BY_LABEL:
+        raise KeyError(f"Unknown model: {label}")
+    return PROFILE_BY_LABEL[label]
+
+
+def parse_model_label(label: str) -> tuple[str, str]:
+    collection, variant = label.split("/", 1)
+    return collection, variant
+
+
+MODEL_LABELS: list[str] = [profile.label for profile in MODEL_PROFILES]

model_loader.py

@@ -0,0 +1,194 @@
+"""Notebook-style diffusers loader for BiliSakura Hub models."""
+
+from __future__ import annotations
+
+import gc
+import inspect
+import os
+from pathlib import Path
+from typing import Any, get_args, get_origin
+
+import torch
+from diffusers import DiffusionPipeline
+import diffusers.pipelines.pipeline_utils as pipeline_utils
+
+from model_catalog import ModelProfile, get_profile
+
+
+def _patch_diffusers_custom_pipeline_type_check() -> None:
+    """Work around diffusers 0.36 KeyError when custom pipelines omit parsed annotations."""
+
+    if getattr(pipeline_utils, "_bilisakura_type_check_patch", False):
+        return
+
+    @classmethod
+    def patched_get_signature_types(cls):
+        signature_types = {}
+        for name, param in inspect.signature(cls.__init__).parameters.items():
+            if name == "self":
+                continue
+            annotation = param.annotation
+            if annotation is inspect.Parameter.empty:
+                signature_types[name] = (inspect.Signature.empty,)
+                continue
+            origin = get_origin(annotation)
+            if inspect.isclass(annotation):
+                signature_types[name] = (annotation,)
+            elif origin is not None:
+                args = get_args(annotation)
+                signature_types[name] = args if args else (annotation,)
+            else:
+                signature_types[name] = (inspect.Signature.empty,)
+        return signature_types
+
+    original_from_pretrained = DiffusionPipeline.from_pretrained.__func__
+
+    @classmethod
+    def from_pretrained_patched(cls, pretrained_model_name_or_path, *args, **kwargs):
+        original_get_signature_types = DiffusionPipeline._get_signature_types
+        DiffusionPipeline._get_signature_types = patched_get_signature_types
+        try:
+            return original_from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs)
+        finally:
+            DiffusionPipeline._get_signature_types = original_get_signature_types
+
+    DiffusionPipeline.from_pretrained = from_pretrained_patched
+    pipeline_utils._bilisakura_type_check_patch = True
+
+
+_patch_diffusers_custom_pipeline_type_check()
+
+
+LOCAL_MODELS_ROOT = Path(os.environ.get("LOCAL_MODELS_ROOT", "")).expanduser()
+USE_LOCAL_MODELS = LOCAL_MODELS_ROOT.is_dir()
+HF_ORG = os.environ.get("HF_MODEL_ORG", "BiliSakura")
+
+SCHEDULER_CLASSES = {
+    "DDIMScheduler": "diffusers.DDIMScheduler",
+    "FlowMatchEulerDiscreteScheduler": "diffusers.FlowMatchEulerDiscreteScheduler",
+    "FlowMatchHeunDiscreteScheduler": "diffusers.FlowMatchHeunDiscreteScheduler",
+}
+
+
+class PipelineManager:
+    def __init__(self) -> None:
+        self._pipe: DiffusionPipeline | None = None
+        self._loaded_label: str | None = None
+        self._loaded_profile: ModelProfile | None = None
+
+    @property
+    def loaded_label(self) -> str | None:
+        return self._loaded_label
+
+    @property
+    def loaded_profile(self) -> ModelProfile | None:
+        return self._loaded_profile
+
+    @property
+    def pipe(self) -> DiffusionPipeline | None:
+        return self._pipe
+
+    def _resolve_model_source(self, profile: ModelProfile) -> str:
+        if USE_LOCAL_MODELS:
+            local_path = LOCAL_MODELS_ROOT / profile.collection / profile.variant
+            if local_path.is_dir():
+                return str(local_path)
+        return f"{HF_ORG}/{profile.collection}/{profile.variant}"
+
+    def _resolve_dtype(self, profile: ModelProfile) -> torch.dtype:
+        return torch.bfloat16 if profile.dtype == "bfloat16" else torch.float32
+
+    def _apply_scheduler(self, pipe: DiffusionPipeline, profile: ModelProfile) -> None:
+        if not profile.scheduler:
+            return
+        import diffusers
+
+        scheduler_cls = getattr(diffusers, profile.scheduler)
+        pipe.scheduler = scheduler_cls.from_config(pipe.scheduler.config, **profile.scheduler_kwargs)
+
+    def _apply_post_load(self, pipe: DiffusionPipeline) -> None:
+        pipe.set_progress_bar_config(disable=True)
+
+    def unload(self) -> None:
+        if self._pipe is not None:
+            del self._pipe
+            self._pipe = None
+        self._loaded_label = None
+        self._loaded_profile = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def load(self, collection: str, variant: str) -> tuple[str, ModelProfile]:
+        profile = get_profile(collection, variant)
+        label = profile.label
+        if self._loaded_label == label and self._pipe is not None:
+            return f"Model already loaded: `{label}`", profile
+
+        self.unload()
+        model_source = self._resolve_model_source(profile)
+        dtype = self._resolve_dtype(profile)
+
+        load_kwargs: dict[str, Any] = {
+            "trust_remote_code": True,
+            "torch_dtype": dtype,
+        }
+        if profile.use_custom_pipeline:
+            source_path = Path(model_source)
+            if source_path.is_dir():
+                load_kwargs["custom_pipeline"] = str(source_path / "pipeline.py")
+            else:
+                load_kwargs["custom_pipeline"] = model_source
+
+        if USE_LOCAL_MODELS and not model_source.startswith(HF_ORG):
+            load_kwargs["local_files_only"] = True
+
+        pipe = DiffusionPipeline.from_pretrained(model_source, **load_kwargs)
+        self._apply_scheduler(pipe, profile)
+        pipe = pipe.to("cuda")
+        self._apply_post_load(pipe)
+
+        self._pipe = pipe
+        self._loaded_label = label
+        self._loaded_profile = profile
+        return f"Loaded `{label}` from `{model_source}`.", profile
+
+
+PIPELINE_MANAGER = PipelineManager()
+
+
+def build_inference_steps(profile: ModelProfile, steps: int) -> int | list[int]:
+    if profile.steps_are_list:
+        per_stage = max(1, steps // 4)
+        return [per_stage, per_stage, per_stage, per_stage]
+    return int(steps)
+
+
+def run_inference(
+    profile: ModelProfile,
+    pipe: DiffusionPipeline,
+    *,
+    class_label: str,
+    seed: int,
+    num_steps: int,
+    guidance_scale: float,
+    height: int,
+    width: int,
+    extra_kwargs: dict[str, Any] | None = None,
+) -> Any:
+    generator = torch.Generator(device="cuda").manual_seed(int(seed))
+    call_kwargs: dict[str, Any] = {
+        "num_inference_steps": build_inference_steps(profile, num_steps),
+        "guidance_scale": float(guidance_scale),
+        "generator": generator,
+    }
+    call_kwargs.update(extra_kwargs if extra_kwargs is not None else profile.extra_call_kwargs)
+
+    label = class_label.strip() or profile.default_class_label
+    call_kwargs["class_labels"] = int(label) if label.isdigit() else label
+
+    if height > 0 and width > 0:
+        call_kwargs["height"] = int(height)
+        call_kwargs["width"] = int(width)
+
+    return pipe(**call_kwargs).images[0]

requirements.txt

@@ -0,0 +1,9 @@
+gradio>=5.0.0
+torch
+diffusers>=0.36.0
+transformers
+accelerate
+safetensors
+huggingface_hub
+einops
+spaces