| """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 hub_model_url(self) -> str: |
| return f"https://huggingface.co/{self.hub_repo}/tree/main/{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 |
| return _infer_resolution_from_variant(self.variant) |
|
|
|
|
| def _infer_resolution_from_variant(variant: str) -> int: |
| name = variant.lower() |
| if "1024" in name: |
| return 1024 |
| if "512" in name or "img512" in name or name.endswith("-32") or "-32-" in name: |
| return 512 |
| 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: |
| res = _infer_resolution_from_variant(variant) |
| 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] |
|
|
| NATIVE_SCHEDULER_COLLECTIONS = frozenset( |
| { |
| "EDM2-diffusers", |
| "iMF-diffusers", |
| "pMF-diffusers", |
| "PixelFlow-diffusers", |
| "Self-Flow-diffusers", |
| } |
| ) |
|
|
| DIFFUSION_SCHEDULERS: list[str] = [ |
| "DDPMScheduler", |
| "DDIMScheduler", |
| "PNDMScheduler", |
| "LMSDiscreteScheduler", |
| "EulerDiscreteScheduler", |
| "EulerAncestralDiscreteScheduler", |
| "HeunDiscreteScheduler", |
| "DPMSolverMultistepScheduler", |
| "DPMSolverSinglestepScheduler", |
| "KDPM2DiscreteScheduler", |
| "KDPM2AncestralDiscreteScheduler", |
| "DEISMultistepScheduler", |
| "UniPCMultistepScheduler", |
| ] |
|
|
| FLOW_SCHEDULERS: list[str] = [ |
| "FlowMatchEulerDiscreteScheduler", |
| "FlowMatchHeunDiscreteScheduler", |
| ] |
|
|
| SWAPPABLE_SCHEDULERS: list[str] = DIFFUSION_SCHEDULERS + [ |
| scheduler for scheduler in FLOW_SCHEDULERS if scheduler not in DIFFUSION_SCHEDULERS |
| ] |
|
|
|
|
| def uses_native_scheduler(profile: ModelProfile) -> bool: |
| return profile.collection in NATIVE_SCHEDULER_COLLECTIONS |
|
|
|
|
| def scheduler_family_for_profile(profile: ModelProfile) -> Literal["native", "flow", "diffusion"]: |
| if uses_native_scheduler(profile): |
| return "native" |
| if profile.scheduler and profile.scheduler.startswith("FlowMatch"): |
| return "flow" |
| if profile.collection in {"JiT-diffusers", "SiT-diffusers"}: |
| return "flow" |
| return "diffusion" |
|
|
|
|
| def scheduler_choices_for_profile(profile: ModelProfile) -> list[str]: |
| if scheduler_family_for_profile(profile) == "native": |
| return [] |
| return list(SWAPPABLE_SCHEDULERS) |
|
|
