diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..94a7bb48dab7dbc170cbd09b877a2644f313b31b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,15 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +AFM-B-2-1NFE-guided/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-B-2-1NFE-noguide/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-L-2-1NFE-guided/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-L-2-1NFE-noguide/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-M-2-1NFE-guided/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-M-2-1NFE-noguide/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-XL-2-112layer-1NFE-guided/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-XL-2-1NFE-guided/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-XL-2-1NFE-noguide/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-XL-2-2NFE-guided/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-XL-2-2NFE-noguide/demo.png filter=lfs diff=lfs merge=lfs -text +AFM-XL-2-56layer-1NFE-guided/demo.png filter=lfs diff=lfs merge=lfs -text diff --git a/AFM-B-2-1NFE-guided/demo.png b/AFM-B-2-1NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..4d7f9bbf6ef8895006e9b83335fbef8f06f7aadf --- /dev/null +++ b/AFM-B-2-1NFE-guided/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a9b31babbcfbd6fe28ae301f4cdb6c193fd5621b4348e7a30432d41d8e962808 +size 130162 diff --git a/AFM-B-2-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc b/AFM-B-2-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5403634b9cc6df84725144ab1f78a2cac2168840 Binary files /dev/null and b/AFM-B-2-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc differ diff --git a/AFM-B-2-1NFE-guided/generator/config.json b/AFM-B-2-1NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..39aa2f048301932543b9c210687f3b7b3b627bb1 --- /dev/null +++ b/AFM-B-2-1NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-B/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 12, + "hidden_size": 768, + "patch_size": 2, + "num_heads": 12 +} diff --git a/AFM-B-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-B-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d55e1c4e76d2754d0cbebd44b6e28e4829be16e5 --- /dev/null +++ b/AFM-B-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8288ea86cfca336987ea264c9ec315c683a39fdd7ec9740efb32bbcd463a4875 +size 589740400 diff --git a/AFM-B-2-1NFE-guided/generator/modeling_afm.py b/AFM-B-2-1NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-B-2-1NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-B-2-1NFE-guided/model_index.json b/AFM-B-2-1NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-B-2-1NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-B-2-1NFE-guided/pipeline.py b/AFM-B-2-1NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-B-2-1NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-B-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-B-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..061f6a72f5367911f30c297e313d6038c04eccb6 Binary files /dev/null and b/AFM-B-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-B-2-1NFE-guided/scheduler/scheduler_config.json b/AFM-B-2-1NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-B-2-1NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-B-2-1NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-B-2-1NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-B-2-1NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-B-2-1NFE-guided/vae/config.json b/AFM-B-2-1NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-B-2-1NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-B-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-B-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-B-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-B-2-1NFE-noguide/demo.png b/AFM-B-2-1NFE-noguide/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..8e2b9c71fb64af4b2b59ad8698afd2a5760dba63 --- /dev/null +++ b/AFM-B-2-1NFE-noguide/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0bfabcede55121fac1b540e5ef9807e29fd9e74572799c361313ed373cde8f25 +size 141287 diff --git a/AFM-B-2-1NFE-noguide/generator/config.json b/AFM-B-2-1NFE-noguide/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..39aa2f048301932543b9c210687f3b7b3b627bb1 --- /dev/null +++ b/AFM-B-2-1NFE-noguide/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-B/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 12, + "hidden_size": 768, + "patch_size": 2, + "num_heads": 12 +} diff --git a/AFM-B-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors b/AFM-B-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bcf9313335451e3fe6b769bcb5fbdc01a08c060f --- /dev/null +++ b/AFM-B-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eb5297181aebfeb15fd7bf69781c80a2ae15aef5c24a2479991cba2fd17ba14d +size 589740400 diff --git a/AFM-B-2-1NFE-noguide/generator/modeling_afm.py b/AFM-B-2-1NFE-noguide/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-B-2-1NFE-noguide/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-B-2-1NFE-noguide/model_index.json b/AFM-B-2-1NFE-noguide/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-B-2-1NFE-noguide/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-B-2-1NFE-noguide/pipeline.py b/AFM-B-2-1NFE-noguide/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-B-2-1NFE-noguide/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-B-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-B-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a12ad73ea21ab64d57c42f3686823037c722a775 Binary files /dev/null and b/AFM-B-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-B-2-1NFE-noguide/scheduler/scheduler_config.json b/AFM-B-2-1NFE-noguide/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-B-2-1NFE-noguide/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-B-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py b/AFM-B-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-B-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-B-2-1NFE-noguide/vae/config.json b/AFM-B-2-1NFE-noguide/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-B-2-1NFE-noguide/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-B-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors b/AFM-B-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-B-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-L-2-1NFE-guided/demo.png b/AFM-L-2-1NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..7b1e96c3b0821247027488c11097689f68166627 --- /dev/null +++ b/AFM-L-2-1NFE-guided/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b4a722696e47cfab9a967c46526d4740eafdb6c29beb605b3d140a2fc3caa885 +size 112806 diff --git a/AFM-L-2-1NFE-guided/generator/config.json b/AFM-L-2-1NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..beb9b0ae33b98e968ca99d0e45eb6b2fd5ab08ac --- /dev/null +++ b/AFM-L-2-1NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-L/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 24, + "hidden_size": 1024, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-L-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-L-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b9d342650a2dcf9d95c6982d738967205fc6a2b7 --- /dev/null +++ b/AFM-L-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f95d199f5759ef66f561d7d9f08f31489b1794272f450713a844551c2c816682 +size 2046434688 diff --git a/AFM-L-2-1NFE-guided/generator/modeling_afm.py b/AFM-L-2-1NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-L-2-1NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-L-2-1NFE-guided/model_index.json b/AFM-L-2-1NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-L-2-1NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-L-2-1NFE-guided/pipeline.py b/AFM-L-2-1NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-L-2-1NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-L-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-L-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2092c1dac962c21d126a52cce3040aa4c9256fc1 Binary files /dev/null and b/AFM-L-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-L-2-1NFE-guided/scheduler/scheduler_config.json b/AFM-L-2-1NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-L-2-1NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-L-2-1NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-L-2-1NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-L-2-1NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-L-2-1NFE-guided/vae/config.json b/AFM-L-2-1NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-L-2-1NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-L-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-L-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-L-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-L-2-1NFE-noguide/demo.png b/AFM-L-2-1NFE-noguide/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..301442f0428121625cea9ccb2499326ff60217dd --- /dev/null +++ b/AFM-L-2-1NFE-noguide/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:acaf057cb18ab2cb8c4f09cf9cdc0029d344c0cd1cfb1d0e76dace6c03506d34 +size 115504 diff --git a/AFM-L-2-1NFE-noguide/generator/config.json b/AFM-L-2-1NFE-noguide/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..beb9b0ae33b98e968ca99d0e45eb6b2fd5ab08ac --- /dev/null +++ b/AFM-L-2-1NFE-noguide/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-L/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 24, + "hidden_size": 1024, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-L-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors b/AFM-L-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..08e408c59db8aca06920a60e9d73b131093d5ea8 --- /dev/null +++ b/AFM-L-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f8519fdc60ab76ee443981ec7a2601933ae54e13b67a247a0cc03e706c239c22 +size 2046434688 diff --git a/AFM-L-2-1NFE-noguide/generator/modeling_afm.py b/AFM-L-2-1NFE-noguide/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-L-2-1NFE-noguide/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-L-2-1NFE-noguide/model_index.json b/AFM-L-2-1NFE-noguide/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-L-2-1NFE-noguide/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-L-2-1NFE-noguide/pipeline.py b/AFM-L-2-1NFE-noguide/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-L-2-1NFE-noguide/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-L-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-L-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93f7865fe210b50556f83397d1ddf11ff95030f7 Binary files /dev/null and b/AFM-L-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-L-2-1NFE-noguide/scheduler/scheduler_config.json b/AFM-L-2-1NFE-noguide/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-L-2-1NFE-noguide/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-L-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py b/AFM-L-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-L-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-L-2-1NFE-noguide/vae/config.json b/AFM-L-2-1NFE-noguide/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-L-2-1NFE-noguide/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-L-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors b/AFM-L-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-L-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-M-2-1NFE-guided/demo.png b/AFM-M-2-1NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..0309b97d94ba2e8438bc72125b6ee5d24c8de869 --- /dev/null +++ b/AFM-M-2-1NFE-guided/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:033a1fc178676cf614ca0e2a68bbedd349094dd2e6235e97ac0523aadccd558e +size 108149 diff --git a/AFM-M-2-1NFE-guided/generator/config.json b/AFM-M-2-1NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..2144112097d7f70cae0ac33bfca600a5796401a1 --- /dev/null +++ b/AFM-M-2-1NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-M/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 16, + "hidden_size": 1024, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-M-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-M-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..750709ac5cb9092f9a12f0d86086f4a7534f7e0a --- /dev/null +++ b/AFM-M-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:12c57e80b3d69a132a46821571db1d7559f7b5df75fc74454b99c394c5c3385c +size 1367133768 diff --git a/AFM-M-2-1NFE-guided/generator/modeling_afm.py b/AFM-M-2-1NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-M-2-1NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-M-2-1NFE-guided/model_index.json b/AFM-M-2-1NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-M-2-1NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-M-2-1NFE-guided/pipeline.py b/AFM-M-2-1NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-M-2-1NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-M-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-M-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67389aafe31185b4ad5d735cc5224c7ecceb4751 Binary files /dev/null and b/AFM-M-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-M-2-1NFE-guided/scheduler/scheduler_config.json b/AFM-M-2-1NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-M-2-1NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-M-2-1NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-M-2-1NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-M-2-1NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-M-2-1NFE-guided/vae/config.json b/AFM-M-2-1NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-M-2-1NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-M-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-M-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-M-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-M-2-1NFE-noguide/demo.png b/AFM-M-2-1NFE-noguide/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..97430ff117d0ad8be3fd0c447c8afece66c09f35 --- /dev/null +++ b/AFM-M-2-1NFE-noguide/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2f8a3c191cdfe22bd4e0102c916057fcffd6323dc231c8e1fb4de9ada268b9e4 +size 110398 diff --git a/AFM-M-2-1NFE-noguide/generator/config.json b/AFM-M-2-1NFE-noguide/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..2144112097d7f70cae0ac33bfca600a5796401a1 --- /dev/null +++ b/AFM-M-2-1NFE-noguide/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-M/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 16, + "hidden_size": 1024, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-M-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors b/AFM-M-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9f28b521d09c11e6418c8dfd5bb642616715778e --- /dev/null +++ b/AFM-M-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ccdd512d0451e0fb435e8e0c16add4b4e949ecdb8996d6abf4e0e933b43d6e1c +size 1367133768 diff --git a/AFM-M-2-1NFE-noguide/generator/modeling_afm.py b/AFM-M-2-1NFE-noguide/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-M-2-1NFE-noguide/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-M-2-1NFE-noguide/model_index.json b/AFM-M-2-1NFE-noguide/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-M-2-1NFE-noguide/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-M-2-1NFE-noguide/pipeline.py b/AFM-M-2-1NFE-noguide/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-M-2-1NFE-noguide/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-M-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-M-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87701589bbccbe8fb45056a31a96b1c6082af6f5 Binary files /dev/null and b/AFM-M-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-M-2-1NFE-noguide/scheduler/scheduler_config.json b/AFM-M-2-1NFE-noguide/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-M-2-1NFE-noguide/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-M-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py b/AFM-M-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-M-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-M-2-1NFE-noguide/vae/config.json b/AFM-M-2-1NFE-noguide/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-M-2-1NFE-noguide/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-M-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors b/AFM-M-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-M-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-XL-2-112layer-1NFE-guided/demo.png b/AFM-XL-2-112layer-1NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..39987d4a46b21d4f95ea756aa5c5a0dad8da6ba1 --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:02d75444fd83294f9652f22e040fd59394c5b5774f1971ed02bf909aefac9435 +size 110421 diff --git a/AFM-XL-2-112layer-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc b/AFM-XL-2-112layer-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c1cf92c74702a51c40c383cec291c52484b2e56 Binary files /dev/null and b/AFM-XL-2-112layer-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc differ diff --git a/AFM-XL-2-112layer-1NFE-guided/generator/config.json b/AFM-XL-2-112layer-1NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..954f5e25429d411dafe9889875bd4457db64373f --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGeneratorDeep2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-XL/2", + "architecture": "deep", + "repeat": 4, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 28, + "hidden_size": 1152, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-XL-2-112layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-XL-2-112layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fecb17de5915335b357d49d09b85899f94aeb3b3 --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f49b04fcd6f77d6fb7039781d0a872bb80ed4a07001b11dc3cd340492965c1a3 +size 2999303040 diff --git a/AFM-XL-2-112layer-1NFE-guided/generator/modeling_afm.py b/AFM-XL-2-112layer-1NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-XL-2-112layer-1NFE-guided/model_index.json b/AFM-XL-2-112layer-1NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..9a67cb7c995f6f0bb64e3fbbca961304f2499c87 --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGeneratorDeep2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-XL-2-112layer-1NFE-guided/pipeline.py b/AFM-XL-2-112layer-1NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-XL-2-112layer-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-XL-2-112layer-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df591395077092fd351ba830a84056ba091d8096 Binary files /dev/null and b/AFM-XL-2-112layer-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-XL-2-112layer-1NFE-guided/scheduler/scheduler_config.json b/AFM-XL-2-112layer-1NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-XL-2-112layer-1NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-XL-2-112layer-1NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-XL-2-112layer-1NFE-guided/vae/config.json b/AFM-XL-2-112layer-1NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-XL-2-112layer-1NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-XL-2-112layer-1NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-XL-2-112layer-1NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-XL-2-1NFE-guided/demo.png b/AFM-XL-2-1NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..8fd4c275db6dca628f34368c5b23352ced7532f9 --- /dev/null +++ b/AFM-XL-2-1NFE-guided/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b4042235fe2f2718eeb0fef7d9ec9a56bee7c01415b031ee27d64cfc1db0d4e3 +size 121177 diff --git a/AFM-XL-2-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc b/AFM-XL-2-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1252a236c1b4ff52025f7fe50d25222ee63cc0af Binary files /dev/null and b/AFM-XL-2-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc differ diff --git a/AFM-XL-2-1NFE-guided/generator/config.json b/AFM-XL-2-1NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..041c855cf441f459f946d9bf7416013abbc37854 --- /dev/null +++ b/AFM-XL-2-1NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-XL/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 28, + "hidden_size": 1152, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-XL-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-XL-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e1b64c8546cdf47e71a3316ee6dc2d54921e4d85 --- /dev/null +++ b/AFM-XL-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c4b03219ec7f002d176494286d7bed760ff916e5b1518b3f9ce3e2bb9d0c25d1 +size 2999303040 diff --git a/AFM-XL-2-1NFE-guided/generator/modeling_afm.py b/AFM-XL-2-1NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-XL-2-1NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-XL-2-1NFE-guided/model_index.json b/AFM-XL-2-1NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-XL-2-1NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-XL-2-1NFE-guided/pipeline.py b/AFM-XL-2-1NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-XL-2-1NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-XL-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-XL-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09c797e16ed23ba779bc2cdbb8af63485beb00df Binary files /dev/null and b/AFM-XL-2-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-XL-2-1NFE-guided/scheduler/scheduler_config.json b/AFM-XL-2-1NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-XL-2-1NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-XL-2-1NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-XL-2-1NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-XL-2-1NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-XL-2-1NFE-guided/vae/config.json b/AFM-XL-2-1NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-XL-2-1NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-XL-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-XL-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-XL-2-1NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-XL-2-1NFE-noguide/demo.png b/AFM-XL-2-1NFE-noguide/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..194676385afba4586422adfe0e9d4f0223bcbaf3 --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e290d78bb4d1e750153711c50fa3899b1cc7d33abee522a4bfd812d6c8c2a02c +size 121448 diff --git a/AFM-XL-2-1NFE-noguide/generator/config.json b/AFM-XL-2-1NFE-noguide/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..041c855cf441f459f946d9bf7416013abbc37854 --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-XL/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 28, + "hidden_size": 1152, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-XL-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors b/AFM-XL-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..65bc42a0d63076b1a81b16a78c7b33bb4fd6050d --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:acb2f8302ed8e316e57a73a6fd34dd4da9b14d752e94658dfc1d6f86f580ebfc +size 2999303040 diff --git a/AFM-XL-2-1NFE-noguide/generator/modeling_afm.py b/AFM-XL-2-1NFE-noguide/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-XL-2-1NFE-noguide/model_index.json b/AFM-XL-2-1NFE-noguide/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..636af8d2b8dff852fe92809f8b09cc6158065cca --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-XL-2-1NFE-noguide/pipeline.py b/AFM-XL-2-1NFE-noguide/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-XL-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-XL-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c98c184d667dbe459045eeb7eaa3ac68d9221218 Binary files /dev/null and b/AFM-XL-2-1NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-XL-2-1NFE-noguide/scheduler/scheduler_config.json b/AFM-XL-2-1NFE-noguide/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-XL-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py b/AFM-XL-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-XL-2-1NFE-noguide/vae/config.json b/AFM-XL-2-1NFE-noguide/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-XL-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors b/AFM-XL-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-XL-2-1NFE-noguide/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-XL-2-2NFE-guided/demo.png b/AFM-XL-2-2NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..34c691967ea41bcd692bf902a65f6ca8cf1f1f0d --- /dev/null +++ b/AFM-XL-2-2NFE-guided/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4237a61a9b94e0c39e38dd589d56a1f8a1dab613c12171c68627a8683f6ef80 +size 110905 diff --git a/AFM-XL-2-2NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc b/AFM-XL-2-2NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ae40d39b40732d953379b65ab63056cb9b0e386 Binary files /dev/null and b/AFM-XL-2-2NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc differ diff --git a/AFM-XL-2-2NFE-guided/generator/config.json b/AFM-XL-2-2NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..059ac33c5509c2c68ff07fce7d3dc5167e2c83ef --- /dev/null +++ b/AFM-XL-2-2NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-XL/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": true, + "use_t_tgt": false, + "pred_type": "v", + "num_inference_steps": 2, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 28, + "hidden_size": 1152, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-XL-2-2NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-XL-2-2NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..04a9f7e013059c7c171f8d1595bdca713c05378d --- /dev/null +++ b/AFM-XL-2-2NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c4ba62222809142771f48d4281e98299d6c11b83ff2b9b75163f98c7a23a9167 +size 2999303040 diff --git a/AFM-XL-2-2NFE-guided/generator/modeling_afm.py b/AFM-XL-2-2NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-XL-2-2NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-XL-2-2NFE-guided/model_index.json b/AFM-XL-2-2NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..5aa478748145cae43ee0a8415ac8424996a234aa --- /dev/null +++ b/AFM-XL-2-2NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "v", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-XL-2-2NFE-guided/pipeline.py b/AFM-XL-2-2NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-XL-2-2NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-XL-2-2NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-XL-2-2NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3b32a5aab36499ae173622c8d1d4a00e7c439e8 Binary files /dev/null and b/AFM-XL-2-2NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-XL-2-2NFE-guided/scheduler/scheduler_config.json b/AFM-XL-2-2NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-XL-2-2NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-XL-2-2NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-XL-2-2NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-XL-2-2NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-XL-2-2NFE-guided/vae/config.json b/AFM-XL-2-2NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-XL-2-2NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-XL-2-2NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-XL-2-2NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-XL-2-2NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-XL-2-2NFE-noguide/demo.png b/AFM-XL-2-2NFE-noguide/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..069364081b823e1feb4ba663ee406e788bfdc966 --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8891f9f105422b5fd16f8350450c90165af5e3c1e2fbfc5c4ee8879f19892069 +size 120730 diff --git a/AFM-XL-2-2NFE-noguide/generator/config.json b/AFM-XL-2-2NFE-noguide/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..059ac33c5509c2c68ff07fce7d3dc5167e2c83ef --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-XL/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": true, + "use_t_tgt": false, + "pred_type": "v", + "num_inference_steps": 2, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 28, + "hidden_size": 1152, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-XL-2-2NFE-noguide/generator/diffusion_pytorch_model.safetensors b/AFM-XL-2-2NFE-noguide/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9fc46e68bcb20ac69b6d9e265d4c308baf8ea484 --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8a646088766ee3de076b7731984bad8fce5105d4a7166e84872dd0d41a6f5a1a +size 2999303040 diff --git a/AFM-XL-2-2NFE-noguide/generator/modeling_afm.py b/AFM-XL-2-2NFE-noguide/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-XL-2-2NFE-noguide/model_index.json b/AFM-XL-2-2NFE-noguide/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..5aa478748145cae43ee0a8415ac8424996a234aa --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "v", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-XL-2-2NFE-noguide/pipeline.py b/AFM-XL-2-2NFE-noguide/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-XL-2-2NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-XL-2-2NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b4207954d4ee299050614453b35a9bbfaccb2ee Binary files /dev/null and b/AFM-XL-2-2NFE-noguide/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-XL-2-2NFE-noguide/scheduler/scheduler_config.json b/AFM-XL-2-2NFE-noguide/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-XL-2-2NFE-noguide/scheduler/scheduling_continuous_flow.py b/AFM-XL-2-2NFE-noguide/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-XL-2-2NFE-noguide/vae/config.json b/AFM-XL-2-2NFE-noguide/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-XL-2-2NFE-noguide/vae/diffusion_pytorch_model.safetensors b/AFM-XL-2-2NFE-noguide/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-XL-2-2NFE-noguide/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-XL-2-4NFE-guided/demo.png b/AFM-XL-2-4NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..e3aab45196edbefa539bea90c11c1b445a82cbc2 Binary files /dev/null and b/AFM-XL-2-4NFE-guided/demo.png differ diff --git a/AFM-XL-2-4NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc b/AFM-XL-2-4NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1023f2efa4826f2ec5570426b344138a7e7973b3 Binary files /dev/null and b/AFM-XL-2-4NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc differ diff --git a/AFM-XL-2-4NFE-guided/generator/config.json b/AFM-XL-2-4NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a9501ed13286692aac5584bd887a6a78b99f3059 --- /dev/null +++ b/AFM-XL-2-4NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGenerator2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-XL/2", + "architecture": "standard", + "repeat": 1, + "use_t_src": true, + "use_t_tgt": false, + "pred_type": "v", + "num_inference_steps": 4, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 28, + "hidden_size": 1152, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-XL-2-4NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-XL-2-4NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5a7d1e2c4090f4f94b2f74670308ae5f8fa00139 --- /dev/null +++ b/AFM-XL-2-4NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0ba89ed56a0d1e898ec71dd953a7c28060fd04b440e85a41ccb00b385a93daa3 +size 2999303040 diff --git a/AFM-XL-2-4NFE-guided/generator/modeling_afm.py b/AFM-XL-2-4NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-XL-2-4NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-XL-2-4NFE-guided/model_index.json b/AFM-XL-2-4NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..5aa478748145cae43ee0a8415ac8424996a234aa --- /dev/null +++ b/AFM-XL-2-4NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGenerator2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "v", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-XL-2-4NFE-guided/pipeline.py b/AFM-XL-2-4NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-XL-2-4NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-XL-2-4NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-XL-2-4NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ba652ef2499bd511bd0a705acb418bea5458c3f Binary files /dev/null and b/AFM-XL-2-4NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-XL-2-4NFE-guided/scheduler/scheduler_config.json b/AFM-XL-2-4NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-XL-2-4NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-XL-2-4NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-XL-2-4NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-XL-2-4NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-XL-2-4NFE-guided/vae/config.json b/AFM-XL-2-4NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-XL-2-4NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-XL-2-4NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-XL-2-4NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-XL-2-4NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/AFM-XL-2-56layer-1NFE-guided/demo.png b/AFM-XL-2-56layer-1NFE-guided/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..accb91083ebe38085e5ce2094fcd4501c1cdcf3f --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:083ea9794ae26890e097f8c10da5bf441e113a078803a369eb5fc16cffd93142 +size 103547 diff --git a/AFM-XL-2-56layer-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc b/AFM-XL-2-56layer-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..983f10e4f337ebd05e1990818835b96a6e20c8f6 Binary files /dev/null and b/AFM-XL-2-56layer-1NFE-guided/generator/__pycache__/modeling_afm.cpython-312.pyc differ diff --git a/AFM-XL-2-56layer-1NFE-guided/generator/config.json b/AFM-XL-2-56layer-1NFE-guided/generator/config.json new file mode 100644 index 0000000000000000000000000000000000000000..80af7dcfad80a927c3d33415b8a6bbaf5b0c05cf --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/generator/config.json @@ -0,0 +1,20 @@ +{ + "_class_name": "AFMGeneratorDeep2DModel", + "_diffusers_version": "0.36.0", + "model_type": "AFM-XL/2", + "architecture": "deep", + "repeat": 2, + "use_t_src": false, + "use_t_tgt": false, + "pred_type": "x", + "num_inference_steps": 1, + "learn_sigma": false, + "class_dropout_prob": 0.0, + "input_size": 32, + "num_classes": 1000, + "in_channels": 4, + "depth": 28, + "hidden_size": 1152, + "patch_size": 2, + "num_heads": 16 +} diff --git a/AFM-XL-2-56layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors b/AFM-XL-2-56layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..c1ba7437052ae4c656438d6dba4b221973db7f09 --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a35d5ec2bb5d92a4ff5bc98bd923ec5d34818d21213446c055010a3a7224a2ee +size 2999303040 diff --git a/AFM-XL-2-56layer-1NFE-guided/generator/modeling_afm.py b/AFM-XL-2-56layer-1NFE-guided/generator/modeling_afm.py new file mode 100644 index 0000000000000000000000000000000000000000..1a57f57e2e2dbd5ad38b3b8fc24ada7ef6bb6c13 --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/generator/modeling_afm.py @@ -0,0 +1,511 @@ +"""Self-contained AFM generator for Hugging Face diffusers checkpoints.""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from torch import nn + + +def get_dit_transformer_class(): + from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel + + return DiTTransformer2DModel + + +def build_dit_config( + *, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, +) -> Dict[str, Any]: + attention_head_dim = hidden_size // num_heads + num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0) + out_channels = in_channels * 2 if learn_sigma else in_channels + return { + "sample_size": input_size, + "num_layers": depth, + "num_attention_heads": num_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "patch_size": patch_size, + "attention_bias": True, + "activation_fn": "gelu-approximate", + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_type": "ada_norm_zero", + "norm_elementwise_affine": False, + "dropout": 0.0, + "norm_num_groups": 32, + "norm_eps": 1e-5, + "upcast_attention": False, + } + + +def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: + converted = copy.deepcopy(state_dict) + + converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous() + converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous() + + if "t_embedder.mlp.0.weight" in converted: + timestep_weights = { + "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"), + "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"), + "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"), + "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"), + } + else: + timestep_weights = None + class_embedding = converted.pop("y_embedder.embedding_table.weight") + if class_embedding.shape[0] == 1000: + null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype) + class_embedding = torch.cat([class_embedding, null_class], dim=0) + + for block_idx in range(depth): + if timestep_weights is not None: + for key, tensor in timestep_weights.items(): + converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone() + converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = ( + class_embedding.clone() + ) + converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.weight" + ] + converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[ + f"blocks.{block_idx}.adaLN_modulation.1.bias" + ] + + q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0) + q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0) + converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q + converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k + converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v + converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[ + f"blocks.{block_idx}.attn.proj.weight" + ] + converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[ + f"blocks.{block_idx}.attn.proj.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc1.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc1.bias" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[ + f"blocks.{block_idx}.mlp.fc2.weight" + ] + converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[ + f"blocks.{block_idx}.mlp.fc2.bias" + ] + + for suffix in ( + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.fc1.weight", + "mlp.fc1.bias", + "mlp.fc2.weight", + "mlp.fc2.bias", + "adaLN_modulation.1.weight", + "adaLN_modulation.1.bias", + ): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight") + converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias") + converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight") + converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias") + + converted.pop("pos_embed", None) + for block_idx in range(depth): + for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"): + converted.pop(f"blocks.{block_idx}.{suffix}", None) + + return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()} + + +def load_legacy_dit_state_dict( + state_dict: Dict[str, torch.Tensor], + depth: int, + patch_size: int, + hidden_size: int, + num_heads: int, +) -> Dict[str, torch.Tensor]: + if "transformer_blocks.0.attn1.to_q.weight" in state_dict: + return state_dict + if any(key.startswith("blocks.") for key in state_dict): + return _convert_legacy_dit_state_dict(state_dict, depth) + return state_dict + + +def forward_dit_output( + transformer, + hidden_states: torch.Tensor, + timestep: Optional[torch.Tensor], + class_labels: torch.Tensor, + extra_cond: Optional[torch.Tensor] = None, +) -> torch.Tensor: + hidden_states = transformer.pos_embed(hidden_states) + + for block in transformer.transformer_blocks: + if torch.is_grad_enabled() and transformer.gradient_checkpointing: + hidden_states = transformer._gradient_checkpointing_func( + block, + hidden_states, + None, + None, + None, + timestep, + None, + class_labels, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=class_labels, + ) + + conditioning = transformer.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + if extra_cond is not None: + conditioning = conditioning + extra_cond + shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + return hidden_states.reshape( + shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size) + ) + + +def remap_transformer_state_dict( + state_dict: Dict[str, torch.Tensor], + legacy_kwargs: Dict[str, int], + extra_prefixes: tuple[str, ...] = (), +) -> Dict[str, torch.Tensor]: + state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs) + remapped: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if any(key.startswith(prefix) for prefix in extra_prefixes): + remapped[key] = value + elif key.startswith("transformer."): + remapped[key] = value + else: + remapped[f"transformer.{key}"] = value + return remapped + + +def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor: + if output.shape[1] == in_channels: + return output + output, _ = output.chunk(2, dim=1) + return output + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = self.timestep_embedding(t.float(), self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + return self.mlp(emb.to(dtype=weight_dtype)) + + +class Generator(nn.Module): + def __init__( + self, + *, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.use_t_src = use_t_src + self.use_t_tgt = use_t_tgt + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + if use_t_tgt: + self.t_tgt_embedder = TimestepEmbedder(hidden_size) + nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02) + + def load_state_dict(self, state_dict, strict: bool = True): + extra = ("t_tgt_embedder.",) + has_extra = any(key.startswith(extra) for key in state_dict) + has_t_embedder = any("t_embedder" in key for key in state_dict) + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra) + if self.use_t_tgt and not has_extra: + strict = False + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, t_src=None, t_tgt=None): + if self.use_t_src and t_src is not None: + timestep = (t_src * 1000).long() + else: + timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long) + + extra_cond = None + if self.use_t_tgt and t_tgt is not None: + extra_cond = self.t_tgt_embedder(t_tgt * 1000) + + output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond) + return split_output_channels(output, self.in_channels) + + +class GeneratorDeep(nn.Module): + def __init__( + self, + *, + repeat: int, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + **_, + ): + super().__init__() + self.repeat = repeat + self._legacy_kwargs = { + "depth": depth, + "hidden_size": hidden_size, + "patch_size": patch_size, + "num_heads": num_heads, + } + dit_cls = get_dit_transformer_class() + self.transformer = dit_cls( + **build_dit_config( + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + ) + self.in_channels = in_channels + + def load_state_dict(self, state_dict, strict: bool = True): + remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs) + has_t_embedder = any("t_embedder" in key for key in state_dict) + if not has_t_embedder: + strict = False + out = super().load_state_dict(remapped, strict=strict) + if not has_t_embedder: + for name, param in self.named_parameters(): + if "timestep_embedder" in name: + param.data.zero_() + return out + + def forward(self, x, y, *args, **kwargs): + hidden_states = self.transformer.pos_embed(x) + for t in torch.arange(1, 0, -1 / self.repeat).tolist(): + timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long) + for block in self.transformer.transformer_blocks: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=timestep, + cross_attention_kwargs=None, + class_labels=y, + ) + + conditioning = self.transformer.transformer_blocks[0].norm1.emb( + torch.zeros(x.shape[0], device=x.device, dtype=torch.long), + y, + hidden_dtype=hidden_states.dtype, + ) + shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.transformer.proj_out_2(hidden_states) + + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.transformer.patch_size, + self.transformer.patch_size, + self.transformer.out_channels, + ) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size) + ) + return split_output_channels(output, self.in_channels) + + +class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "standard", + repeat: int = 1, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + if architecture == "deep": + raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.") + Generator.__init__( + self, + use_t_src=use_t_src, + use_t_tgt=use_t_tgt, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin): + config_name = "config.json" + + @register_to_config + def __init__( + self, + model_type: str = "AFM-XL/2", + architecture: str = "deep", + repeat: int = 2, + use_t_src: bool = False, + use_t_tgt: bool = False, + depth: int = 28, + hidden_size: int = 1152, + patch_size: int = 2, + num_heads: int = 16, + learn_sigma: bool = False, + class_dropout_prob: float = 0.0, + input_size: int = 32, + num_classes: int = 1000, + in_channels: int = 4, + pred_type: str = "x", + num_inference_steps: int = 1, + ) -> None: + GeneratorDeep.__init__( + self, + repeat=repeat, + depth=depth, + hidden_size=hidden_size, + patch_size=patch_size, + num_heads=num_heads, + learn_sigma=learn_sigma, + class_dropout_prob=class_dropout_prob, + input_size=input_size, + num_classes=num_classes, + in_channels=in_channels, + ) + + +__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"] diff --git a/AFM-XL-2-56layer-1NFE-guided/model_index.json b/AFM-XL-2-56layer-1NFE-guided/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..9a67cb7c995f6f0bb64e3fbbca961304f2499c87 --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/model_index.json @@ -0,0 +1,1022 @@ +{ + "_class_name": [ + "pipeline", + "AFMPipeline" + ], + "_diffusers_version": "0.36.0", + "generator": [ + "modeling_afm", + "AFMGeneratorDeep2DModel" + ], + "vae": [ + "diffusers", + "AutoencoderKL" + ], + "scheduler": [ + "scheduling_continuous_flow", + "ContinuousFlowMatchScheduler" + ], + "pred_type": "x", + "id2label": { + "0": "tench, Tinca tinca", + "1": "goldfish, Carassius auratus", + "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", + "3": "tiger shark, Galeocerdo cuvieri", + "4": "hammerhead, hammerhead shark", + "5": "electric ray, crampfish, numbfish, torpedo", + "6": "stingray", + "7": "cock", + "8": "hen", + "9": "ostrich, Struthio camelus", + "10": "brambling, Fringilla montifringilla", + "11": "goldfinch, Carduelis carduelis", + "12": "house finch, linnet, Carpodacus mexicanus", + "13": "junco, snowbird", + "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea", + "15": "robin, American robin, Turdus migratorius", + "16": "bulbul", + "17": "jay", + "18": "magpie", + "19": "chickadee", + "20": "water ouzel, dipper", + "21": "kite", + "22": "bald eagle, American eagle, Haliaeetus leucocephalus", + "23": "vulture", + "24": "great grey owl, great gray owl, Strix nebulosa", + "25": "European fire salamander, Salamandra salamandra", + "26": "common newt, Triturus vulgaris", + "27": "eft", + "28": "spotted salamander, Ambystoma maculatum", + "29": "axolotl, mud puppy, Ambystoma mexicanum", + "30": "bullfrog, Rana catesbeiana", + "31": "tree frog, tree-frog", + "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui", + "33": "loggerhead, loggerhead turtle, Caretta caretta", + "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea", + "35": "mud turtle", + "36": "terrapin", + "37": "box turtle, box tortoise", + "38": "banded gecko", + "39": "common iguana, iguana, Iguana iguana", + "40": "American chameleon, anole, Anolis carolinensis", + "41": "whiptail, whiptail lizard", + "42": "agama", + "43": "frilled lizard, Chlamydosaurus kingi", + "44": "alligator lizard", + "45": "Gila monster, Heloderma suspectum", + "46": "green lizard, Lacerta viridis", + "47": "African chameleon, Chamaeleo chamaeleon", + "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis", + "49": "African crocodile, Nile crocodile, Crocodylus niloticus", + "50": "American alligator, Alligator mississipiensis", + "51": "triceratops", + "52": "thunder snake, worm snake, Carphophis amoenus", + "53": "ringneck snake, ring-necked snake, ring snake", + "54": "hognose snake, puff adder, sand viper", + "55": "green snake, grass snake", + "56": "king snake, kingsnake", + "57": "garter snake, grass snake", + "58": "water snake", + "59": "vine snake", + "60": "night snake, Hypsiglena torquata", + "61": "boa constrictor, Constrictor constrictor", + "62": "rock python, rock snake, Python sebae", + "63": "Indian cobra, Naja naja", + "64": "green mamba", + "65": "sea snake", + "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus", + "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus", + "68": "sidewinder, horned rattlesnake, Crotalus cerastes", + "69": "trilobite", + "70": "harvestman, daddy longlegs, Phalangium opilio", + "71": "scorpion", + "72": "black and gold garden spider, Argiope aurantia", + "73": "barn spider, Araneus cavaticus", + "74": "garden spider, Aranea diademata", + "75": "black widow, Latrodectus mactans", + "76": "tarantula", + "77": "wolf spider, hunting spider", + "78": "tick", + "79": "centipede", + "80": "black grouse", + "81": "ptarmigan", + "82": "ruffed grouse, partridge, Bonasa umbellus", + "83": "prairie chicken, prairie grouse, prairie fowl", + "84": "peacock", + "85": "quail", + "86": "partridge", + "87": "African grey, African gray, Psittacus erithacus", + "88": "macaw", + "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita", + "90": "lorikeet", + "91": "coucal", + "92": "bee eater", + "93": "hornbill", + "94": "hummingbird", + "95": "jacamar", + "96": "toucan", + "97": "drake", + "98": "red-breasted merganser, Mergus serrator", + "99": "goose", + "100": "black swan, Cygnus atratus", + "101": "tusker", + "102": "echidna, spiny anteater, anteater", + "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus", + "104": "wallaby, brush kangaroo", + "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus", + "106": "wombat", + "107": "jellyfish", + "108": "sea anemone, anemone", + "109": "brain coral", + "110": "flatworm, platyhelminth", + "111": "nematode, nematode worm, roundworm", + "112": "conch", + "113": "snail", + "114": "slug", + "115": "sea slug, nudibranch", + "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore", + "117": "chambered nautilus, pearly nautilus, nautilus", + "118": "Dungeness crab, Cancer magister", + "119": "rock crab, Cancer irroratus", + "120": "fiddler crab", + "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica", + "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus", + "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish", + "124": "crayfish, crawfish, crawdad, crawdaddy", + "125": "hermit crab", + "126": "isopod", + "127": "white stork, Ciconia ciconia", + "128": "black stork, Ciconia nigra", + "129": "spoonbill", + "130": "flamingo", + "131": "little blue heron, Egretta caerulea", + "132": "American egret, great white heron, Egretta albus", + "133": "bittern", + "134": "crane", + "135": "limpkin, Aramus pictus", + "136": "European gallinule, Porphyrio porphyrio", + "137": "American coot, marsh hen, mud hen, water hen, Fulica americana", + "138": "bustard", + "139": "ruddy turnstone, Arenaria interpres", + "140": "red-backed sandpiper, dunlin, Erolia alpina", + "141": "redshank, Tringa totanus", + "142": "dowitcher", + "143": "oystercatcher, oyster catcher", + "144": "pelican", + "145": "king penguin, Aptenodytes patagonica", + "146": "albatross, mollymawk", + "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus", + "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca", + "149": "dugong, Dugong dugon", + "150": "sea lion", + "151": "Chihuahua", + "152": "Japanese spaniel", + "153": "Maltese dog, Maltese terrier, Maltese", + "154": "Pekinese, Pekingese, Peke", + "155": "Shih-Tzu", + "156": "Blenheim spaniel", + "157": "papillon", + "158": "toy terrier", + "159": "Rhodesian ridgeback", + "160": "Afghan hound, Afghan", + "161": "basset, basset hound", + "162": "beagle", + "163": "bloodhound, sleuthhound", + "164": "bluetick", + "165": "black-and-tan coonhound", + "166": "Walker hound, Walker foxhound", + "167": "English foxhound", + "168": "redbone", + "169": "borzoi, Russian wolfhound", + "170": "Irish wolfhound", + "171": "Italian greyhound", + "172": "whippet", + "173": "Ibizan hound, Ibizan Podenco", + "174": "Norwegian elkhound, elkhound", + "175": "otterhound, otter hound", + "176": "Saluki, gazelle hound", + "177": "Scottish deerhound, deerhound", + "178": "Weimaraner", + "179": "Staffordshire bullterrier, Staffordshire bull terrier", + "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier", + "181": "Bedlington terrier", + "182": "Border terrier", + "183": "Kerry blue terrier", + "184": "Irish terrier", + "185": "Norfolk terrier", + "186": "Norwich terrier", + "187": "Yorkshire terrier", + "188": "wire-haired fox terrier", + "189": "Lakeland terrier", + "190": "Sealyham terrier, Sealyham", + "191": "Airedale, Airedale terrier", + "192": "cairn, cairn terrier", + "193": "Australian terrier", + "194": "Dandie Dinmont, Dandie Dinmont terrier", + "195": "Boston bull, Boston terrier", + "196": "miniature schnauzer", + "197": "giant schnauzer", + "198": "standard schnauzer", + "199": "Scotch terrier, Scottish terrier, Scottie", + "200": "Tibetan terrier, chrysanthemum dog", + "201": "silky terrier, Sydney silky", + "202": "soft-coated wheaten terrier", + "203": "West Highland white terrier", + "204": "Lhasa, Lhasa apso", + "205": "flat-coated retriever", + "206": "curly-coated retriever", + "207": "golden retriever", + "208": "Labrador retriever", + "209": "Chesapeake Bay retriever", + "210": "German short-haired pointer", + "211": "vizsla, Hungarian pointer", + "212": "English setter", + "213": "Irish setter, red setter", + "214": "Gordon setter", + "215": "Brittany spaniel", + "216": "clumber, clumber spaniel", + "217": "English springer, English springer spaniel", + "218": "Welsh springer spaniel", + "219": "cocker spaniel, English cocker spaniel, cocker", + "220": "Sussex spaniel", + "221": "Irish water spaniel", + "222": "kuvasz", + "223": "schipperke", + "224": "groenendael", + "225": "malinois", + "226": "briard", + "227": "kelpie", + "228": "komondor", + "229": "Old English sheepdog, bobtail", + "230": "Shetland sheepdog, Shetland sheep dog, Shetland", + "231": "collie", + "232": "Border collie", + "233": "Bouvier des Flandres, Bouviers des Flandres", + "234": "Rottweiler", + "235": "German shepherd, German shepherd dog, German police dog, alsatian", + "236": "Doberman, Doberman pinscher", + "237": "miniature pinscher", + "238": "Greater Swiss Mountain dog", + "239": "Bernese mountain dog", + "240": "Appenzeller", + "241": "EntleBucher", + "242": "boxer", + "243": "bull mastiff", + "244": "Tibetan mastiff", + "245": "French bulldog", + "246": "Great Dane", + "247": "Saint Bernard, St Bernard", + "248": "Eskimo dog, husky", + "249": "malamute, malemute, Alaskan malamute", + "250": "Siberian husky", + "251": "dalmatian, coach dog, carriage dog", + "252": "affenpinscher, monkey pinscher, monkey dog", + "253": "basenji", + "254": "pug, pug-dog", + "255": "Leonberg", + "256": "Newfoundland, Newfoundland dog", + "257": "Great Pyrenees", + "258": "Samoyed, Samoyede", + "259": "Pomeranian", + "260": "chow, chow chow", + "261": "keeshond", + "262": "Brabancon griffon", + "263": "Pembroke, Pembroke Welsh corgi", + "264": "Cardigan, Cardigan Welsh corgi", + "265": "toy poodle", + "266": "miniature poodle", + "267": "standard poodle", + "268": "Mexican hairless", + "269": "timber wolf, grey wolf, gray wolf, Canis lupus", + "270": "white wolf, Arctic wolf, Canis lupus tundrarum", + "271": "red wolf, maned wolf, Canis rufus, Canis niger", + "272": "coyote, prairie wolf, brush wolf, Canis latrans", + "273": "dingo, warrigal, warragal, Canis dingo", + "274": "dhole, Cuon alpinus", + "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus", + "276": "hyena, hyaena", + "277": "red fox, Vulpes vulpes", + "278": "kit fox, Vulpes macrotis", + "279": "Arctic fox, white fox, Alopex lagopus", + "280": "grey fox, gray fox, Urocyon cinereoargenteus", + "281": "tabby, tabby cat", + "282": "tiger cat", + "283": "Persian cat", + "284": "Siamese cat, Siamese", + "285": "Egyptian cat", + "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor", + "287": "lynx, catamount", + "288": "leopard, Panthera pardus", + "289": "snow leopard, ounce, Panthera uncia", + "290": "jaguar, panther, Panthera onca, Felis onca", + "291": "lion, king of beasts, Panthera leo", + "292": "tiger, Panthera tigris", + "293": "cheetah, chetah, Acinonyx jubatus", + "294": "brown bear, bruin, Ursus arctos", + "295": "American black bear, black bear, Ursus americanus, Euarctos americanus", + "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus", + "297": "sloth bear, Melursus ursinus, Ursus ursinus", + "298": "mongoose", + "299": "meerkat, mierkat", + "300": "tiger beetle", + "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle", + "302": "ground beetle, carabid beetle", + "303": "long-horned beetle, longicorn, longicorn beetle", + "304": "leaf beetle, chrysomelid", + "305": "dung beetle", + "306": "rhinoceros beetle", + "307": "weevil", + "308": "fly", + "309": "bee", + "310": "ant, emmet, pismire", + "311": "grasshopper, hopper", + "312": "cricket", + "313": "walking stick, walkingstick, stick insect", + "314": "cockroach, roach", + "315": "mantis, mantid", + "316": "cicada, cicala", + "317": "leafhopper", + "318": "lacewing, lacewing fly", + "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk", + "320": "damselfly", + "321": "admiral", + "322": "ringlet, ringlet butterfly", + "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus", + "324": "cabbage butterfly", + "325": "sulphur butterfly, sulfur butterfly", + "326": "lycaenid, lycaenid butterfly", + "327": "starfish, sea star", + "328": "sea urchin", + "329": "sea cucumber, holothurian", + "330": "wood rabbit, cottontail, cottontail rabbit", + "331": "hare", + "332": "Angora, Angora rabbit", + "333": "hamster", + "334": "porcupine, hedgehog", + "335": "fox squirrel, eastern fox squirrel, Sciurus niger", + "336": "marmot", + "337": "beaver", + "338": "guinea pig, Cavia cobaya", + "339": "sorrel", + "340": "zebra", + "341": "hog, pig, grunter, squealer, Sus scrofa", + "342": "wild boar, boar, Sus scrofa", + "343": "warthog", + "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius", + "345": "ox", + "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis", + "347": "bison", + "348": "ram, tup", + "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis", + "350": "ibex, Capra ibex", + "351": "hartebeest", + "352": "impala, Aepyceros melampus", + "353": "gazelle", + "354": "Arabian camel, dromedary, Camelus dromedarius", + "355": "llama", + "356": "weasel", + "357": "mink", + "358": "polecat, fitch, foulmart, foumart, Mustela putorius", + "359": "black-footed ferret, ferret, Mustela nigripes", + "360": "otter", + "361": "skunk, polecat, wood pussy", + "362": "badger", + "363": "armadillo", + "364": "three-toed sloth, ai, Bradypus tridactylus", + "365": "orangutan, orang, orangutang, Pongo pygmaeus", + "366": "gorilla, Gorilla gorilla", + "367": "chimpanzee, chimp, Pan troglodytes", + "368": "gibbon, Hylobates lar", + "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus", + "370": "guenon, guenon monkey", + "371": "patas, hussar monkey, Erythrocebus patas", + "372": "baboon", + "373": "macaque", + "374": "langur", + "375": "colobus, colobus monkey", + "376": "proboscis monkey, Nasalis larvatus", + "377": "marmoset", + "378": "capuchin, ringtail, Cebus capucinus", + "379": "howler monkey, howler", + "380": "titi, titi monkey", + "381": "spider monkey, Ateles geoffroyi", + "382": "squirrel monkey, Saimiri sciureus", + "383": "Madagascar cat, ring-tailed lemur, Lemur catta", + "384": "indri, indris, Indri indri, Indri brevicaudatus", + "385": "Indian elephant, Elephas maximus", + "386": "African elephant, Loxodonta africana", + "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens", + "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca", + "389": "barracouta, snoek", + "390": "eel", + "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch", + "392": "rock beauty, Holocanthus tricolor", + "393": "anemone fish", + "394": "sturgeon", + "395": "gar, garfish, garpike, billfish, Lepisosteus osseus", + "396": "lionfish", + "397": "puffer, pufferfish, blowfish, globefish", + "398": "abacus", + "399": "abaya", + "400": "academic gown, academic robe, judge robe", + "401": "accordion, piano accordion, squeeze box", + "402": "acoustic guitar", + "403": "aircraft carrier, carrier, flattop, attack aircraft carrier", + "404": "airliner", + "405": "airship, dirigible", + "406": "altar", + "407": "ambulance", + "408": "amphibian, amphibious vehicle", + "409": "analog clock", + "410": "apiary, bee house", + "411": "apron", + "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin", + "413": "assault rifle, assault gun", + "414": "backpack, back pack, knapsack, packsack, rucksack, haversack", + "415": "bakery, bakeshop, bakehouse", + "416": "balance beam, beam", + "417": "balloon", + "418": "ballpoint, ballpoint pen, ballpen, Biro", + "419": "Band Aid", + "420": "banjo", + "421": "bannister, banister, balustrade, balusters, handrail", + "422": "barbell", + "423": "barber chair", + "424": "barbershop", + "425": "barn", + "426": "barometer", + "427": "barrel, cask", + "428": "barrow, garden cart, lawn cart, wheelbarrow", + "429": "baseball", + "430": "basketball", + "431": "bassinet", + "432": "bassoon", + "433": "bathing cap, swimming cap", + "434": "bath towel", + "435": "bathtub, bathing tub, bath, tub", + "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon", + "437": "beacon, lighthouse, beacon light, pharos", + "438": "beaker", + "439": "bearskin, busby, shako", + "440": "beer bottle", + "441": "beer glass", + "442": "bell cote, bell cot", + "443": "bib", + "444": "bicycle-built-for-two, tandem bicycle, tandem", + "445": "bikini, two-piece", + "446": "binder, ring-binder", + "447": "binoculars, field glasses, opera glasses", + "448": "birdhouse", + "449": "boathouse", + "450": "bobsled, bobsleigh, bob", + "451": "bolo tie, bolo, bola tie, bola", + "452": "bonnet, poke bonnet", + "453": "bookcase", + "454": "bookshop, bookstore, bookstall", + "455": "bottlecap", + "456": "bow", + "457": "bow tie, bow-tie, bowtie", + "458": "brass, memorial tablet, plaque", + "459": "brassiere, bra, bandeau", + "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty", + "461": "breastplate, aegis, egis", + "462": "broom", + "463": "bucket, pail", + "464": "buckle", + "465": "bulletproof vest", + "466": "bullet train, bullet", + "467": "butcher shop, meat market", + "468": "cab, hack, taxi, taxicab", + "469": "caldron, cauldron", + "470": "candle, taper, wax light", + "471": "cannon", + "472": "canoe", + "473": "can opener, tin opener", + "474": "cardigan", + "475": "car mirror", + "476": "carousel, carrousel, merry-go-round, roundabout, whirligig", + "477": "carpenters kit, tool kit", + "478": "carton", + "479": "car wheel", + "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM", + "481": "cassette", + "482": "cassette player", + "483": "castle", + "484": "catamaran", + "485": "CD player", + "486": "cello, violoncello", + "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone", + "488": "chain", + "489": "chainlink fence", + "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour", + "491": "chain saw, chainsaw", + "492": "chest", + "493": "chiffonier, commode", + "494": "chime, bell, gong", + "495": "china cabinet, china closet", + "496": "Christmas stocking", + "497": "church, church building", + "498": "cinema, movie theater, movie theatre, movie house, picture palace", + "499": "cleaver, meat cleaver, chopper", + "500": "cliff dwelling", + "501": "cloak", + "502": "clog, geta, patten, sabot", + "503": "cocktail shaker", + "504": "coffee mug", + "505": "coffeepot", + "506": "coil, spiral, volute, whorl, helix", + "507": "combination lock", + "508": "computer keyboard, keypad", + "509": "confectionery, confectionary, candy store", + "510": "container ship, containership, container vessel", + "511": "convertible", + "512": "corkscrew, bottle screw", + "513": "cornet, horn, trumpet, trump", + "514": "cowboy boot", + "515": "cowboy hat, ten-gallon hat", + "516": "cradle", + "517": "crane", + "518": "crash helmet", + "519": "crate", + "520": "crib, cot", + "521": "Crock Pot", + "522": "croquet ball", + "523": "crutch", + "524": "cuirass", + "525": "dam, dike, dyke", + "526": "desk", + "527": "desktop computer", + "528": "dial telephone, dial phone", + "529": "diaper, nappy, napkin", + "530": "digital clock", + "531": "digital watch", + "532": "dining table, board", + "533": "dishrag, dishcloth", + "534": "dishwasher, dish washer, dishwashing machine", + "535": "disk brake, disc brake", + "536": "dock, dockage, docking facility", + "537": "dogsled, dog sled, dog sleigh", + "538": "dome", + "539": "doormat, welcome mat", + "540": "drilling platform, offshore rig", + "541": "drum, membranophone, tympan", + "542": "drumstick", + "543": "dumbbell", + "544": "Dutch oven", + "545": "electric fan, blower", + "546": "electric guitar", + "547": "electric locomotive", + "548": "entertainment center", + "549": "envelope", + "550": "espresso maker", + "551": "face powder", + "552": "feather boa, boa", + "553": "file, file cabinet, filing cabinet", + "554": "fireboat", + "555": "fire engine, fire truck", + "556": "fire screen, fireguard", + "557": "flagpole, flagstaff", + "558": "flute, transverse flute", + "559": "folding chair", + "560": "football helmet", + "561": "forklift", + "562": "fountain", + "563": "fountain pen", + "564": "four-poster", + "565": "freight car", + "566": "French horn, horn", + "567": "frying pan, frypan, skillet", + "568": "fur coat", + "569": "garbage truck, dustcart", + "570": "gasmask, respirator, gas helmet", + "571": "gas pump, gasoline pump, petrol pump, island dispenser", + "572": "goblet", + "573": "go-kart", + "574": "golf ball", + "575": "golfcart, golf cart", + "576": "gondola", + "577": "gong, tam-tam", + "578": "gown", + "579": "grand piano, grand", + "580": "greenhouse, nursery, glasshouse", + "581": "grille, radiator grille", + "582": "grocery store, grocery, food market, market", + "583": "guillotine", + "584": "hair slide", + "585": "hair spray", + "586": "half track", + "587": "hammer", + "588": "hamper", + "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier", + "590": "hand-held computer, hand-held microcomputer", + "591": "handkerchief, hankie, hanky, hankey", + "592": "hard disc, hard disk, fixed disk", + "593": "harmonica, mouth organ, harp, mouth harp", + "594": "harp", + "595": "harvester, reaper", + "596": "hatchet", + "597": "holster", + "598": "home theater, home theatre", + "599": "honeycomb", + "600": "hook, claw", + "601": "hoopskirt, crinoline", + "602": "horizontal bar, high bar", + "603": "horse cart, horse-cart", + "604": "hourglass", + "605": "iPod", + "606": "iron, smoothing iron", + "607": "jack-o-lantern", + "608": "jean, blue jean, denim", + "609": "jeep, landrover", + "610": "jersey, T-shirt, tee shirt", + "611": "jigsaw puzzle", + "612": "jinrikisha, ricksha, rickshaw", + "613": "joystick", + "614": "kimono", + "615": "knee pad", + "616": "knot", + "617": "lab coat, laboratory coat", + "618": "ladle", + "619": "lampshade, lamp shade", + "620": "laptop, laptop computer", + "621": "lawn mower, mower", + "622": "lens cap, lens cover", + "623": "letter opener, paper knife, paperknife", + "624": "library", + "625": "lifeboat", + "626": "lighter, light, igniter, ignitor", + "627": "limousine, limo", + "628": "liner, ocean liner", + "629": "lipstick, lip rouge", + "630": "Loafer", + "631": "lotion", + "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system", + "633": "loupe, jewelers loupe", + "634": "lumbermill, sawmill", + "635": "magnetic compass", + "636": "mailbag, postbag", + "637": "mailbox, letter box", + "638": "maillot", + "639": "maillot, tank suit", + "640": "manhole cover", + "641": "maraca", + "642": "marimba, xylophone", + "643": "mask", + "644": "matchstick", + "645": "maypole", + "646": "maze, labyrinth", + "647": "measuring cup", + "648": "medicine chest, medicine cabinet", + "649": "megalith, megalithic structure", + "650": "microphone, mike", + "651": "microwave, microwave oven", + "652": "military uniform", + "653": "milk can", + "654": "minibus", + "655": "miniskirt, mini", + "656": "minivan", + "657": "missile", + "658": "mitten", + "659": "mixing bowl", + "660": "mobile home, manufactured home", + "661": "Model T", + "662": "modem", + "663": "monastery", + "664": "monitor", + "665": "moped", + "666": "mortar", + "667": "mortarboard", + "668": "mosque", + "669": "mosquito net", + "670": "motor scooter, scooter", + "671": "mountain bike, all-terrain bike, off-roader", + "672": "mountain tent", + "673": "mouse, computer mouse", + "674": "mousetrap", + "675": "moving van", + "676": "muzzle", + "677": "nail", + "678": "neck brace", + "679": "necklace", + "680": "nipple", + "681": "notebook, notebook computer", + "682": "obelisk", + "683": "oboe, hautboy, hautbois", + "684": "ocarina, sweet potato", + "685": "odometer, hodometer, mileometer, milometer", + "686": "oil filter", + "687": "organ, pipe organ", + "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO", + "689": "overskirt", + "690": "oxcart", + "691": "oxygen mask", + "692": "packet", + "693": "paddle, boat paddle", + "694": "paddlewheel, paddle wheel", + "695": "padlock", + "696": "paintbrush", + "697": "pajama, pyjama, pjs, jammies", + "698": "palace", + "699": "panpipe, pandean pipe, syrinx", + "700": "paper towel", + "701": "parachute, chute", + "702": "parallel bars, bars", + "703": "park bench", + "704": "parking meter", + "705": "passenger car, coach, carriage", + "706": "patio, terrace", + "707": "pay-phone, pay-station", + "708": "pedestal, plinth, footstall", + "709": "pencil box, pencil case", + "710": "pencil sharpener", + "711": "perfume, essence", + "712": "Petri dish", + "713": "photocopier", + "714": "pick, plectrum, plectron", + "715": "pickelhaube", + "716": "picket fence, paling", + "717": "pickup, pickup truck", + "718": "pier", + "719": "piggy bank, penny bank", + "720": "pill bottle", + "721": "pillow", + "722": "ping-pong ball", + "723": "pinwheel", + "724": "pirate, pirate ship", + "725": "pitcher, ewer", + "726": "plane, carpenters plane, woodworking plane", + "727": "planetarium", + "728": "plastic bag", + "729": "plate rack", + "730": "plow, plough", + "731": "plunger, plumbers helper", + "732": "Polaroid camera, Polaroid Land camera", + "733": "pole", + "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria", + "735": "poncho", + "736": "pool table, billiard table, snooker table", + "737": "pop bottle, soda bottle", + "738": "pot, flowerpot", + "739": "potters wheel", + "740": "power drill", + "741": "prayer rug, prayer mat", + "742": "printer", + "743": "prison, prison house", + "744": "projectile, missile", + "745": "projector", + "746": "puck, hockey puck", + "747": "punching bag, punch bag, punching ball, punchball", + "748": "purse", + "749": "quill, quill pen", + "750": "quilt, comforter, comfort, puff", + "751": "racer, race car, racing car", + "752": "racket, racquet", + "753": "radiator", + "754": "radio, wireless", + "755": "radio telescope, radio reflector", + "756": "rain barrel", + "757": "recreational vehicle, RV, R.V.", + "758": "reel", + "759": "reflex camera", + "760": "refrigerator, icebox", + "761": "remote control, remote", + "762": "restaurant, eating house, eating place, eatery", + "763": "revolver, six-gun, six-shooter", + "764": "rifle", + "765": "rocking chair, rocker", + "766": "rotisserie", + "767": "rubber eraser, rubber, pencil eraser", + "768": "rugby ball", + "769": "rule, ruler", + "770": "running shoe", + "771": "safe", + "772": "safety pin", + "773": "saltshaker, salt shaker", + "774": "sandal", + "775": "sarong", + "776": "sax, saxophone", + "777": "scabbard", + "778": "scale, weighing machine", + "779": "school bus", + "780": "schooner", + "781": "scoreboard", + "782": "screen, CRT screen", + "783": "screw", + "784": "screwdriver", + "785": "seat belt, seatbelt", + "786": "sewing machine", + "787": "shield, buckler", + "788": "shoe shop, shoe-shop, shoe store", + "789": "shoji", + "790": "shopping basket", + "791": "shopping cart", + "792": "shovel", + "793": "shower cap", + "794": "shower curtain", + "795": "ski", + "796": "ski mask", + "797": "sleeping bag", + "798": "slide rule, slipstick", + "799": "sliding door", + "800": "slot, one-armed bandit", + "801": "snorkel", + "802": "snowmobile", + "803": "snowplow, snowplough", + "804": "soap dispenser", + "805": "soccer ball", + "806": "sock", + "807": "solar dish, solar collector, solar furnace", + "808": "sombrero", + "809": "soup bowl", + "810": "space bar", + "811": "space heater", + "812": "space shuttle", + "813": "spatula", + "814": "speedboat", + "815": "spider web, spiders web", + "816": "spindle", + "817": "sports car, sport car", + "818": "spotlight, spot", + "819": "stage", + "820": "steam locomotive", + "821": "steel arch bridge", + "822": "steel drum", + "823": "stethoscope", + "824": "stole", + "825": "stone wall", + "826": "stopwatch, stop watch", + "827": "stove", + "828": "strainer", + "829": "streetcar, tram, tramcar, trolley, trolley car", + "830": "stretcher", + "831": "studio couch, day bed", + "832": "stupa, tope", + "833": "submarine, pigboat, sub, U-boat", + "834": "suit, suit of clothes", + "835": "sundial", + "836": "sunglass", + "837": "sunglasses, dark glasses, shades", + "838": "sunscreen, sunblock, sun blocker", + "839": "suspension bridge", + "840": "swab, swob, mop", + "841": "sweatshirt", + "842": "swimming trunks, bathing trunks", + "843": "swing", + "844": "switch, electric switch, electrical switch", + "845": "syringe", + "846": "table lamp", + "847": "tank, army tank, armored combat vehicle, armoured combat vehicle", + "848": "tape player", + "849": "teapot", + "850": "teddy, teddy bear", + "851": "television, television system", + "852": "tennis ball", + "853": "thatch, thatched roof", + "854": "theater curtain, theatre curtain", + "855": "thimble", + "856": "thresher, thrasher, threshing machine", + "857": "throne", + "858": "tile roof", + "859": "toaster", + "860": "tobacco shop, tobacconist shop, tobacconist", + "861": "toilet seat", + "862": "torch", + "863": "totem pole", + "864": "tow truck, tow car, wrecker", + "865": "toyshop", + "866": "tractor", + "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi", + "868": "tray", + "869": "trench coat", + "870": "tricycle, trike, velocipede", + "871": "trimaran", + "872": "tripod", + "873": "triumphal arch", + "874": "trolleybus, trolley coach, trackless trolley", + "875": "trombone", + "876": "tub, vat", + "877": "turnstile", + "878": "typewriter keyboard", + "879": "umbrella", + "880": "unicycle, monocycle", + "881": "upright, upright piano", + "882": "vacuum, vacuum cleaner", + "883": "vase", + "884": "vault", + "885": "velvet", + "886": "vending machine", + "887": "vestment", + "888": "viaduct", + "889": "violin, fiddle", + "890": "volleyball", + "891": "waffle iron", + "892": "wall clock", + "893": "wallet, billfold, notecase, pocketbook", + "894": "wardrobe, closet, press", + "895": "warplane, military plane", + "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin", + "897": "washer, automatic washer, washing machine", + "898": "water bottle", + "899": "water jug", + "900": "water tower", + "901": "whiskey jug", + "902": "whistle", + "903": "wig", + "904": "window screen", + "905": "window shade", + "906": "Windsor tie", + "907": "wine bottle", + "908": "wing", + "909": "wok", + "910": "wooden spoon", + "911": "wool, woolen, woollen", + "912": "worm fence, snake fence, snake-rail fence, Virginia fence", + "913": "wreck", + "914": "yawl", + "915": "yurt", + "916": "web site, website, internet site, site", + "917": "comic book", + "918": "crossword puzzle, crossword", + "919": "street sign", + "920": "traffic light, traffic signal, stoplight", + "921": "book jacket, dust cover, dust jacket, dust wrapper", + "922": "menu", + "923": "plate", + "924": "guacamole", + "925": "consomme", + "926": "hot pot, hotpot", + "927": "trifle", + "928": "ice cream, icecream", + "929": "ice lolly, lolly, lollipop, popsicle", + "930": "French loaf", + "931": "bagel, beigel", + "932": "pretzel", + "933": "cheeseburger", + "934": "hotdog, hot dog, red hot", + "935": "mashed potato", + "936": "head cabbage", + "937": "broccoli", + "938": "cauliflower", + "939": "zucchini, courgette", + "940": "spaghetti squash", + "941": "acorn squash", + "942": "butternut squash", + "943": "cucumber, cuke", + "944": "artichoke, globe artichoke", + "945": "bell pepper", + "946": "cardoon", + "947": "mushroom", + "948": "Granny Smith", + "949": "strawberry", + "950": "orange", + "951": "lemon", + "952": "fig", + "953": "pineapple, ananas", + "954": "banana", + "955": "jackfruit, jak, jack", + "956": "custard apple", + "957": "pomegranate", + "958": "hay", + "959": "carbonara", + "960": "chocolate sauce, chocolate syrup", + "961": "dough", + "962": "meat loaf, meatloaf", + "963": "pizza, pizza pie", + "964": "potpie", + "965": "burrito", + "966": "red wine", + "967": "espresso", + "968": "cup", + "969": "eggnog", + "970": "alp", + "971": "bubble", + "972": "cliff, drop, drop-off", + "973": "coral reef", + "974": "geyser", + "975": "lakeside, lakeshore", + "976": "promontory, headland, head, foreland", + "977": "sandbar, sand bar", + "978": "seashore, coast, seacoast, sea-coast", + "979": "valley, vale", + "980": "volcano", + "981": "ballplayer, baseball player", + "982": "groom, bridegroom", + "983": "scuba diver", + "984": "rapeseed", + "985": "daisy", + "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum", + "987": "corn", + "988": "acorn", + "989": "hip, rose hip, rosehip", + "990": "buckeye, horse chestnut, conker", + "991": "coral fungus", + "992": "agaric", + "993": "gyromitra", + "994": "stinkhorn, carrion fungus", + "995": "earthstar", + "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa", + "997": "bolete", + "998": "ear, spike, capitulum", + "999": "toilet tissue, toilet paper, bathroom tissue" + } +} diff --git a/AFM-XL-2-56layer-1NFE-guided/pipeline.py b/AFM-XL-2-56layer-1NFE-guided/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c53b4af7a387dab1499ea365cba4895ae59377ee --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/pipeline.py @@ -0,0 +1,276 @@ +"""Hub custom pipeline: AFMPipeline. + +Load with native Hugging Face diffusers and trust_remote_code=True. +""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.utils.torch_utils import randn_tensor + + +class AFMPipeline(DiffusionPipeline): + model_cpu_offload_seq = "generator->vae" + + @staticmethod + def _coerce_scheduler(scheduler, generator): + if scheduler is not None and not isinstance(scheduler, (list, tuple)): + return scheduler + variant_path = getattr(generator.config, "_name_or_path", None) + if variant_path: + scheduler_dir = Path(variant_path).resolve().parent / "scheduler" + module_path = scheduler_dir / "scheduling_continuous_flow.py" + config_path = scheduler_dir / "scheduler_config.json" + if module_path.is_file() and config_path.is_file(): + spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path) + if spec is not None and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir)) + raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.") + + def __init__( + self, + generator, + vae, + scheduler=None, + pred_type: str = "x", + id2label: Optional[Dict[Union[int, str], str]] = None, + ) -> None: + super().__init__() + scheduler = self._coerce_scheduler(scheduler, generator) + self.register_modules(generator=generator, vae=vae, scheduler=scheduler) + self.register_to_config(pred_type=pred_type) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self._id2label = self._normalize_id2label(id2label) + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = bool(self._id2label) + + @staticmethod + def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]: + if not id2label: + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]: + if not variant_path: + return {} + model_index_path = Path(variant_path).resolve() / "model_index.json" + if not model_index_path.exists(): + return {} + raw = json.loads(model_index_path.read_text(encoding="utf-8")) + id2label = raw.get("id2label") + if not isinstance(id2label, dict): + return {} + return {int(key): value for key, value in id2label.items()} + + @staticmethod + def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]: + label2id: Dict[str, int] = {} + for class_id, value in id2label.items(): + for synonym in value.split(","): + synonym = synonym.strip() + if synonym: + label2id[synonym] = int(class_id) + return dict(sorted(label2id.items())) + + @property + def vae_scale_factor(self) -> int: + block_out_channels = getattr(self.vae.config, "block_out_channels", None) + if block_out_channels: + return int(2 ** (len(block_out_channels) - 1)) + return 8 + + @property + def id2label(self) -> Dict[int, str]: + self._ensure_labels_loaded() + return self._id2label + + def _ensure_labels_loaded(self) -> None: + if self._labels_loaded_from_model_index: + return + loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None)) + if loaded: + self._id2label = loaded + self.labels = self._build_label2id(self._id2label) + self._labels_loaded_from_model_index = True + + def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + self._ensure_labels_loaded() + labels = [label] if isinstance(label, str) else label + if not self.labels: + raise ValueError("No id2label mapping is available in this checkpoint.") + missing = [item for item in labels if item not in self.labels] + if missing: + preview = ", ".join(list(self.labels.keys())[:8]) + raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...") + return [self.labels[item] for item in labels] + + def _normalize_class_labels( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + device: torch.device, + ) -> torch.LongTensor: + if torch.is_tensor(class_labels): + return class_labels.to(device=device, dtype=torch.long).reshape(-1) + if isinstance(class_labels, int): + class_label_ids = [class_labels] + elif isinstance(class_labels, str): + class_label_ids = self.get_label_ids(class_labels) + elif class_labels and isinstance(class_labels[0], str): + class_label_ids = self.get_label_ids(class_labels) + else: + class_label_ids = list(class_labels) + return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1) + + def _default_image_size(self) -> int: + transformer = getattr(self.generator, "transformer", None) + if transformer is not None: + return int(transformer.config.sample_size) * self.vae_scale_factor + return 256 + + def check_inputs( + self, + height: int, + width: int, + num_inference_steps: int, + output_type: str, + sampler: str, + ) -> None: + if num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be >= 1.") + if output_type not in {"pil", "np", "pt", "latent"}: + raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.") + if sampler not in {"euler", "heun"}: + raise ValueError("sampler must be one of: 'euler', 'heun'.") + if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})." + ) + + def prepare_latents( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + in_channels = int(getattr(self.generator, "in_channels", 4)) + return randn_tensor( + (batch_size, in_channels, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + + def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"): + if output_type == "latent": + return latents + scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) + image = self.vae.decode(latents / scaling_factor).sample + if output_type == "pt": + return image + return self.image_processor.postprocess(image, output_type=output_type) + + def _run_denoising_loop( + self, + latents: torch.Tensor, + class_labels_tensor: torch.Tensor, + batch_size: int, + device: torch.device, + num_inference_steps: int, + sampler: str, + pred_type: str, + ) -> torch.Tensor: + self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler) + timesteps = self.scheduler.timesteps + self.generator.eval() + for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))): + t_src_batch = t_src.expand(batch_size) + t_tgt_batch = t_tgt.expand(batch_size) + outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch) + if pred_type == "v": + if sampler == "heun": + latents_next = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch) + latents = self.scheduler.step( + outputs, + t_src, + t_tgt, + latents, + model_output_next=outputs_next, + prediction_type="v", + ).prev_sample + else: + latents = self.scheduler.step( + outputs, t_src, t_tgt, latents, prediction_type="v" + ).prev_sample + else: + latents = outputs + return latents + + @torch.inference_mode() + def __call__( + self, + class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + sampler: str = "euler", + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + output_type: str = "pil", + return_dict: bool = True, + ): + default_size = self._default_image_size() + height = int(height or default_size) + width = int(width or default_size) + if num_inference_steps is None: + num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1)) + self.check_inputs(height, width, num_inference_steps, output_type, sampler) + + device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device + dtype = next(self.generator.parameters()).dtype + pred_type = self.config.pred_type + class_labels_tensor = self._normalize_class_labels(class_labels, device=device) + batch_size = class_labels_tensor.shape[0] + latents = self.prepare_latents( + batch_size=batch_size, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + latents = self._run_denoising_loop( + latents, + class_labels_tensor, + batch_size, + device, + num_inference_steps, + sampler, + pred_type, + ) + image = self.decode_latents(latents, output_type=output_type) + self.maybe_free_model_hooks() + if not return_dict: + return (image,) + return ImagePipelineOutput(images=image) diff --git a/AFM-XL-2-56layer-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc b/AFM-XL-2-56layer-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b1e6a6c680d25e5d8b054f66e7342aa597359b3f Binary files /dev/null and b/AFM-XL-2-56layer-1NFE-guided/scheduler/__pycache__/scheduling_continuous_flow.cpython-312.pyc differ diff --git a/AFM-XL-2-56layer-1NFE-guided/scheduler/scheduler_config.json b/AFM-XL-2-56layer-1NFE-guided/scheduler/scheduler_config.json new file mode 100644 index 0000000000000000000000000000000000000000..0ab49028950992abf116c1fd1789a8de907dfc6d --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/scheduler/scheduler_config.json @@ -0,0 +1,5 @@ +{ + "_class_name": "ContinuousFlowMatchScheduler", + "_diffusers_version": "0.36.0", + "solver": "euler" +} diff --git a/AFM-XL-2-56layer-1NFE-guided/scheduler/scheduling_continuous_flow.py b/AFM-XL-2-56layer-1NFE-guided/scheduler/scheduling_continuous_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..83e6dfbb7db4484e8418d437eb794ec42428d663 --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/scheduler/scheduling_continuous_flow.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class ContinuousFlowMatchSchedulerOutput(BaseOutput): + prev_sample: torch.Tensor + + +class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin): + """Flow-matching scheduler for AFM with time in [1, 0].""" + + order = 2 + + @register_to_config + def __init__(self, solver: str = "euler"): + if solver not in {"euler", "heun"}: + raise ValueError("solver must be one of: 'euler', 'heun'.") + self.timesteps: Optional[torch.Tensor] = None + self.num_inference_steps: Optional[int] = None + self._step_index: Optional[int] = None + + @property + def init_noise_sigma(self) -> float: + return 1.0 + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + solver: Optional[str] = None, + ) -> None: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + self.num_inference_steps = num_inference_steps + if solver == "heun": + grid_size = (num_inference_steps // 2) + 1 + else: + grid_size = num_inference_steps + 1 + self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32) + self._step_index = 0 + if solver is not None: + self.register_to_config(solver=solver) + + def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + del timestep + return sample + + def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int: + if self._step_index is not None: + return self._step_index + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if timestep is None: + return 0 + t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0]) + matches = (self.timesteps - t_value).abs() < 1e-6 + if matches.any(): + return int(matches.nonzero(as_tuple=False)[0].item()) + return 0 + + def step( + self, + model_output: torch.Tensor, + timestep_src: Union[float, torch.Tensor], + timestep_tgt: Union[float, torch.Tensor], + sample: torch.Tensor, + model_output_next: Optional[torch.Tensor] = None, + prediction_type: str = "v", + return_dict: bool = True, + ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]: + if prediction_type == "x": + prev_sample = model_output + else: + t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype) + t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype) + while t_src.ndim < sample.ndim: + t_src = t_src.unsqueeze(-1) + t_tgt = t_tgt.unsqueeze(-1) + dt = t_src - t_tgt + if self.config.solver == "heun" and model_output_next is not None: + prev_sample = sample - dt * 0.5 * (model_output + model_output_next) + else: + prev_sample = sample - dt * model_output + + step_index = self._resolve_step_index(timestep_src) + self._step_index = step_index + 1 + + if not return_dict: + return (prev_sample,) + return ContinuousFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/AFM-XL-2-56layer-1NFE-guided/vae/config.json b/AFM-XL-2-56layer-1NFE-guided/vae/config.json new file mode 100644 index 0000000000000000000000000000000000000000..ae2ab37de2acb8bc348d3de809ae8385324be962 --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/vae/config.json @@ -0,0 +1,38 @@ +{ + "_class_name": "AutoencoderKL", + "_diffusers_version": "0.36.0", + "_name_or_path": "stabilityai/sd-vae-ft-mse", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 4, + "latents_mean": null, + "latents_std": null, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "scaling_factor": 0.18215, + "shift_factor": null, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/AFM-XL-2-56layer-1NFE-guided/vae/diffusion_pytorch_model.safetensors b/AFM-XL-2-56layer-1NFE-guided/vae/diffusion_pytorch_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90464d67ac7303d0ee4696334df13da130a948ea --- /dev/null +++ b/AFM-XL-2-56layer-1NFE-guided/vae/diffusion_pytorch_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815 +size 334643276 diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..edeb58589d17639c80576c050a0882b8943580d2 --- /dev/null +++ b/README.md @@ -0,0 +1,90 @@ +--- +license: mit +library_name: diffusers +pipeline_tag: text-to-image +tags: +- diffusers +- afm +- adversarial-flow-models +- class-conditional +- imagenet +inference: true +widget: +- output: + url: AFM-XL-2-56layer-1NFE-guided/demo.png +language: +- en +--- + +# BiliSakura/AFM-diffusers + +Self-contained [Adversarial Flow Models](https://arxiv.org/abs/2511.22475) checkpoints for Hugging Face diffusers. + +Converted from `ByteDance-Seed/Adversarial-Flow-Models` using `libs/AFM-diffusers/scripts/convert_afm_to_diffusers.py`. + +All models use LDM (Rombach et al., 2022) latent space with `sd-vae-ft-mse`. Guidance abbreviations: **CG** = classifier guidance (Dhariwal & Nichol, 2021), **DA** = data augmentation (Karras et al., 2020a). + +## Demo + +`AFM-XL-2-2NFE-noguide` — class **207** (*golden retriever*), seed **0**, 2 NFE: + +

+ AFM-XL-2-2NFE-noguide demo (class 207, seed 0) +

+ +Each variant folder includes `demo.png` generated with the same prompt settings. + +## Benchmark results (ImageNet 256×256) + +| Model | Params | Guidance | NFE | FID | sFID | IS | Prec. | Recall | Checkpoint | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | +| AFM-B/2 | 130M | None | 1 | 6.07 | 5.31 | 169.51 | 0.72 | 0.49 | `AFM-B-2-1NFE-noguide/` | +| AFM-M/2 | 306M | None | 1 | 5.21 | 5.60 | 178.48 | 0.75 | 0.54 | `AFM-M-2-1NFE-noguide/` | +| AFM-L/2 | 457M | None | 1 | 4.36 | 5.39 | 186.21 | 0.77 | 0.53 | `AFM-L-2-1NFE-noguide/` | +| AFM-XL/2 | 673M | None | 1 | 3.98 | 5.40 | 201.85 | 0.78 | 0.52 | `AFM-XL-2-1NFE-noguide/` | +| AFM-XL/2 | 673M | None | 2 | 2.36 | 4.35 | 235.77 | 0.81 | 0.52 | `AFM-XL-2-2NFE-noguide/` | +| AFM-B/2 | 130M | CG+DA | 1 | 3.05 | 5.32 | 269.18 | 0.81 | 0.51 | `AFM-B-2-1NFE-guided/` | +| AFM-M/2 | 306M | CG+DA | 1 | 2.82 | 5.20 | 279.12 | 0.81 | 0.50 | `AFM-M-2-1NFE-guided/` | +| AFM-L/2 | 457M | CG+DA | 1 | 2.63 | 5.10 | 277.96 | 0.81 | 0.52 | `AFM-L-2-1NFE-guided/` | +| AFM-XL/2 | 673M | CG+DA | 1 | 2.38 | 4.87 | 284.18 | 0.81 | 0.52 | `AFM-XL-2-1NFE-guided/` | +| AFM-XL/2 | 675M | CG+DA | 2 | 2.11 | 4.33 | 273.84 | 0.82 | 0.55 | `AFM-XL-2-2NFE-guided/` | +| AFM-XL/2 (2× deep, 56-layer) | 675M | CG+DA | 1 | 2.08 | 4.79 | 298.33 | 0.79 | 0.56 | `AFM-XL-2-56layer-1NFE-guided/` | +| AFM-XL/2 | 675M | CG+DA | 4 | 2.03 | 4.59 | 259.66 | 0.78 | 0.59 | `AFM-XL-2-4NFE-guided/` | +| AFM-XL/2 (4× deep, 112-layer) | 675M | CG+DA | 1 | 1.94 | 4.54 | 292.20 | 0.79 | 0.56 | `AFM-XL-2-112layer-1NFE-guided/` | + +## Available checkpoints + +| Variant | Model | Steps | Guidance | +| --- | --- | ---: | --- | +| `AFM-B-2-1NFE-guided/` | AFM-B/2 | 1 | guided | +| `AFM-B-2-1NFE-noguide/` | AFM-B/2 | 1 | noguide | +| `AFM-M-2-1NFE-guided/` | AFM-M/2 | 1 | guided | +| `AFM-M-2-1NFE-noguide/` | AFM-M/2 | 1 | noguide | +| `AFM-L-2-1NFE-guided/` | AFM-L/2 | 1 | guided | +| `AFM-L-2-1NFE-noguide/` | AFM-L/2 | 1 | noguide | +| `AFM-XL-2-1NFE-guided/` | AFM-XL/2 | 1 | guided | +| `AFM-XL-2-1NFE-noguide/` | AFM-XL/2 | 1 | noguide | +| `AFM-XL-2-2NFE-guided/` | AFM-XL/2 | 2 | guided | +| `AFM-XL-2-2NFE-noguide/` | AFM-XL/2 | 2 | noguide | +| `AFM-XL-2-4NFE-guided/` | AFM-XL/2 | 4 | guided | +| `AFM-XL-2-56layer-1NFE-guided/` | AFM-XL/2 | 1 | guided | +| `AFM-XL-2-112layer-1NFE-guided/` | AFM-XL/2 | 1 | guided | + +## Inference + +```python +from pathlib import Path +import torch +from diffusers import DiffusionPipeline + +model_dir = Path("./AFM-XL-2-1NFE-guided") +pipe = DiffusionPipeline.from_pretrained( + str(model_dir), + local_files_only=True, + custom_pipeline=str(model_dir / "pipeline.py"), + trust_remote_code=True, + torch_dtype=torch.bfloat16, +).to("cuda") + +image = pipe(class_labels="golden retriever", num_inference_steps=1).images[0] +```