Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+wget-log*
+*.pyc

PixNerd-XL-16-256/README.md

@@ -0,0 +1,41 @@
+# PixNerd-XL-16-256
+
+Self-contained PixNerd-XL/16 checkpoint inside [`BiliSakura/PixNerd-diffusers`](https://huggingface.co/BiliSakura/PixNerd-diffusers). Runtime dependencies: this folder + PyPI `diffusers`/`torch` only.
+
+## Hub path
+
+`BiliSakura/PixNerd-diffusers/PixNerd-XL-16-256`
+
+## Layout
+
+```text
+PixNerd-XL-16-256/
+├── pipeline.py
+├── model_index.json
+├── conversion_metadata.json
+├── transformer/
+└── scheduler/
+```
+
+## Load
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/PixNerd-diffusers/PixNerd-XL-16-256",
+    trust_remote_code=True,
+    torch_dtype=torch.float32,
+).to("cuda")
+
+images = pipe(
+    prompt=207,
+    height=256,
+    width=256,
+    num_inference_steps=25,
+    guidance_scale=4.0,
+    timeshift=3.0,
+    order=2,
+).images
+```

PixNerd-XL-16-256/model_index.json

@@ -1,12 +1,15 @@
 {
-  "_class_name": "PixNerdPipeline",
-  "_diffusers_version": "0.30.0",
+  "_class_name": [
+    "pipeline",
+    "PixNerdPipeline"
+  ],
+  "_diffusers_version": "0.36.0",
   "scheduler": [
-    "diffusers_modules.local.scheduling_pixnerd_flow_match",
+    "scheduling_pixnerd_flow_match",
     "PixNerdFlowMatchScheduler"
   ],
   "transformer": [
-    "diffusers_modules.local.modeling_pixnerd_transformer_2d",
+    "modeling_pixnerd_transformer_2d",
     "PixNerdTransformer2DModel"
   ]
 }

PixNerd-XL-16-256/pipeline.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import sys
 from dataclasses import dataclass
-from typing import List, Optional, Sequence, Union
+from pathlib import Path
+from typing import List, Literal, Optional, Sequence, Union
 
 import torch
 from diffusers import DiffusionPipeline
@@ -9,10 +11,8 @@ from diffusers.image_processor import VaeImageProcessor
 from diffusers.utils import BaseOutput
 from PIL import Image
 
-from .modeling_pixnerd_transformer_2d import PixNerdTransformer2DModel
-from .scheduling_pixnerd_flow_match import PixNerdFlowMatchScheduler
-
 ConditioningInput = Union[str, int, Sequence[Union[str, int]]]
+Language = Literal["en", "cn"]
 
 
 @dataclass
@@ -27,9 +27,11 @@ class PixNerdPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler: PixNerdFlowMatchScheduler,
+        scheduler,
         vae=None,
         conditioner=None,
+        id2label: Optional[dict[int, str]] = None,
+        id2label_cn: Optional[dict[int, str]] = None,
     ):
         super().__init__()
         if vae is None:
@@ -46,6 +48,170 @@ class PixNerdPipeline(DiffusionPipeline):
         )
         self.image_processor = VaeImageProcessor(vae_scale_factor=1)
 
+        if id2label is None and id2label_cn is None:
+            id2label, id2label_cn = self._load_repo_labels()
+        self._id2label = id2label or {}
+        self._id2label_cn = id2label_cn or {}
+        self.labels = self._build_label2id(self._id2label)
+        self.labels_cn = self._build_label2id(self._id2label_cn)
+        self._labels_loaded_from_path = bool(self._id2label or self._id2label_cn)
+
+    def _ensure_labels_loaded(self) -> None:
+        if self._labels_loaded_from_path:
+            return
+
+        path = getattr(getattr(self, "config", None), "_name_or_path", None) or getattr(self, "_name_or_path", None)
+        if not path:
+            return
+
+        id2label, id2label_cn = self._load_labels_for_path(path)
+        if id2label is None and id2label_cn is None:
+            self._labels_loaded_from_path = True
+            return
+
+        self._id2label = id2label or {}
+        self._id2label_cn = id2label_cn or {}
+        self.labels = self._build_label2id(self._id2label)
+        self.labels_cn = self._build_label2id(self._id2label_cn)
+        self._labels_loaded_from_path = True
+
+    @staticmethod
+    def _resolve_labels_dir(pretrained_model_name_or_path: Union[str, Path]) -> Optional[Path]:
+        path = Path(pretrained_model_name_or_path)
+        if not path.exists():
+            try:
+                from huggingface_hub import snapshot_download
+
+                path = Path(snapshot_download(pretrained_model_name_or_path))
+            except Exception:
+                return None
+
+        if (path / "model_index.json").exists():
+            labels_dir = path.parent / "labels"
+        else:
+            labels_dir = path / "labels"
+        return labels_dir if labels_dir.is_dir() else None
+
+    @classmethod
+    def _load_labels_for_path(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+    ) -> tuple[Optional[dict[int, str]], Optional[dict[int, str]]]:
+        labels_dir = cls._resolve_labels_dir(pretrained_model_name_or_path)
+        if labels_dir is None:
+            return None, None
+
+        labels_path = str(labels_dir)
+        inserted = False
+        if labels_path not in sys.path:
+            sys.path.insert(0, labels_path)
+            inserted = True
+        try:
+            from imagenet_labels import load_id2label
+
+            return (
+                load_id2label(labels_dir, lang="en"),
+                load_id2label(labels_dir, lang="cn"),
+            )
+        finally:
+            if inserted and labels_path in sys.path:
+                sys.path.remove(labels_path)
+
+    @staticmethod
+    def _load_repo_labels() -> tuple[Optional[dict[int, str]], Optional[dict[int, str]]]:
+        labels_dir = Path(__file__).resolve().parent.parent / "labels"
+        if not labels_dir.is_dir():
+            return None, None
+
+        labels_path = str(labels_dir)
+        inserted = False
+        if labels_path not in sys.path:
+            sys.path.insert(0, labels_path)
+            inserted = True
+        try:
+            from imagenet_labels import load_id2label
+
+            return (
+                load_id2label(labels_dir, lang="en"),
+                load_id2label(labels_dir, lang="cn"),
+            )
+        finally:
+            if inserted and labels_path in sys.path:
+                sys.path.remove(labels_path)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        pipe = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        id2label, id2label_cn = cls._load_labels_for_path(pretrained_model_name_or_path)
+        if id2label is not None or id2label_cn is not None:
+            pipe._id2label = id2label or {}
+            pipe._id2label_cn = id2label_cn or {}
+            pipe.labels = cls._build_label2id(pipe._id2label)
+            pipe.labels_cn = cls._build_label2id(pipe._id2label_cn)
+        return pipe
+
+    @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 id2label(self) -> dict[int, str]:
+        self._ensure_labels_loaded()
+        return self._id2label
+
+    @property
+    def id2label_cn(self) -> dict[int, str]:
+        self._ensure_labels_loaded()
+        return self._id2label_cn
+
+    def get_label_ids(
+        self,
+        labels: Union[str, List[str]],
+        *,
+        lang: Language = "en",
+    ) -> List[int]:
+        self._ensure_labels_loaded()
+        if isinstance(labels, str):
+            labels = [labels]
+
+        label2id = self.labels if lang == "en" else self.labels_cn
+        if not label2id:
+            raise ValueError(
+                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
+            )
+
+        missing = [label for label in labels if label not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(
+                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+            )
+        return [label2id[label] for label in labels]
+
+    def _resolve_prompt_item(self, value: Union[str, int]) -> int:
+        if isinstance(value, int):
+            return value
+        if value.isdigit():
+            return int(value)
+        if value in self.labels:
+            return self.labels[value]
+        if value in self.labels_cn:
+            return self.labels_cn[value]
+        raise ValueError(
+            f"Unknown class label {value!r}. Pass an ImageNet class id or a synonym from "
+            "`pipe.labels` / `pipe.labels_cn`."
+        )
+
+    def _resolve_prompts(self, prompts: List[Union[str, int]]) -> List[int]:
+        self._ensure_labels_loaded()
+        return [self._resolve_prompt_item(prompt) for prompt in prompts]
+
     @staticmethod
     def _fp_to_uint8(image: torch.Tensor) -> torch.Tensor:
         return torch.clip_((image + 1) * 127.5 + 0.5, 0, 255).to(torch.uint8)
@@ -71,10 +237,11 @@ class PixNerdPipeline(DiffusionPipeline):
         num_images_per_prompt: int,
     ):
         prompts = self._repeat(self._to_list(prompt), num_images_per_prompt)
+        resolved = self._resolve_prompts(prompts)
         metadata = {"device": self._execution_device}
         with torch.no_grad():
-            cond, uncond = self.conditioner(prompts, metadata)
-        return cond, uncond, prompts
+            cond, uncond = self.conditioner(resolved, metadata)
+        return cond, uncond, resolved
 
     def prepare_latents(
         self,
@@ -124,9 +291,10 @@ class PixNerdPipeline(DiffusionPipeline):
         cond, default_uncond, prompts = self.encode_prompt(prompt, num_images_per_prompt)
         if negative_prompt is not None:
             negative = self._repeat(self._to_list(negative_prompt), num_images_per_prompt)
+            resolved_negative = self._resolve_prompts(negative)
             metadata = {"device": self._execution_device}
             with torch.no_grad():
-                _, uncond = self.conditioner(negative, metadata)
+                _, uncond = self.conditioner(resolved_negative, metadata)
         else:
             uncond = default_uncond
         batch_size = len(prompts)
@@ -178,6 +346,7 @@ class PixNerdPipeline(DiffusionPipeline):
             return (output,)
         return PixNerdPipelineOutput(images=output)
 
+
 __all__ = [
     "PixNerdPipeline",
     "PixNerdPipelineOutput",

PixNerd-XL-16-256/scheduler/scheduling_pixnerd_flow_match.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, 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 PixNerdSchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+
+
+class PixNerdFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Diffusers-compatible scheduler wrapper for PixNerd's AdamLM flow-matching sampler.
+    """
+
+    config_name = "scheduler_config.json"
+    order = 1
+    init_noise_sigma = 1.0
+
+    @staticmethod
+    def _lagrange_coeffs(order: int, pre_ts: torch.Tensor, t_start: torch.Tensor, t_end: torch.Tensor) -> List[float]:
+        ts = [float(v) for v in pre_ts[-order:].tolist()]
+        a = float(t_start)
+        b = float(t_end)
+
+        if order == 1:
+            return [1.0]
+        if order == 2:
+            t1, t2 = ts
+            int1 = 0.5 / (t1 - t2) * ((b - t2) ** 2 - (a - t2) ** 2)
+            int2 = 0.5 / (t2 - t1) * ((b - t1) ** 2 - (a - t1) ** 2)
+            total = int1 + int2
+            return [int1 / total, int2 / total]
+        if order == 3:
+            t1, t2, t3 = ts
+            int1_denom = (t1 - t2) * (t1 - t3)
+            int1 = ((1 / 3) * b**3 - 0.5 * (t2 + t3) * b**2 + (t2 * t3) * b) - (
+                (1 / 3) * a**3 - 0.5 * (t2 + t3) * a**2 + (t2 * t3) * a
+            )
+            int1 = int1 / int1_denom
+            int2_denom = (t2 - t1) * (t2 - t3)
+            int2 = ((1 / 3) * b**3 - 0.5 * (t1 + t3) * b**2 + (t1 * t3) * b) - (
+                (1 / 3) * a**3 - 0.5 * (t1 + t3) * a**2 + (t1 * t3) * a
+            )
+            int2 = int2 / int2_denom
+            int3_denom = (t3 - t1) * (t3 - t2)
+            int3 = ((1 / 3) * b**3 - 0.5 * (t1 + t2) * b**2 + (t1 * t2) * b) - (
+                (1 / 3) * a**3 - 0.5 * (t1 + t2) * a**2 + (t1 * t2) * a
+            )
+            int3 = int3 / int3_denom
+            total = int1 + int2 + int3
+            return [int1 / total, int2 / total, int3 / total]
+        if order == 4:
+            t1, t2, t3, t4 = ts
+            int1_denom = (t1 - t2) * (t1 - t3) * (t1 - t4)
+            int1 = ((1 / 4) * b**4 - (1 / 3) * (t2 + t3 + t4) * b**3 + 0.5 * (t3 * t4 + t2 * t3 + t2 * t4) * b**2 - (t2 * t3 * t4) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t2 + t3 + t4) * a**3 + 0.5 * (t3 * t4 + t2 * t3 + t2 * t4) * a**2 - (t2 * t3 * t4) * a
+            )
+            int1 = int1 / int1_denom
+            int2_denom = (t2 - t1) * (t2 - t3) * (t2 - t4)
+            int2 = ((1 / 4) * b**4 - (1 / 3) * (t1 + t3 + t4) * b**3 + 0.5 * (t3 * t4 + t1 * t3 + t1 * t4) * b**2 - (t1 * t3 * t4) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t1 + t3 + t4) * a**3 + 0.5 * (t3 * t4 + t1 * t3 + t1 * t4) * a**2 - (t1 * t3 * t4) * a
+            )
+            int2 = int2 / int2_denom
+            int3_denom = (t3 - t1) * (t3 - t2) * (t3 - t4)
+            int3 = ((1 / 4) * b**4 - (1 / 3) * (t1 + t2 + t4) * b**3 + 0.5 * (t4 * t2 + t1 * t2 + t1 * t4) * b**2 - (t1 * t2 * t4) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t1 + t2 + t4) * a**3 + 0.5 * (t4 * t2 + t1 * t2 + t1 * t4) * a**2 - (t1 * t2 * t4) * a
+            )
+            int3 = int3 / int3_denom
+            int4_denom = (t4 - t1) * (t4 - t2) * (t4 - t3)
+            int4 = ((1 / 4) * b**4 - (1 / 3) * (t1 + t2 + t3) * b**3 + 0.5 * (t3 * t2 + t1 * t2 + t1 * t3) * b**2 - (t1 * t2 * t3) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t1 + t2 + t3) * a**3 + 0.5 * (t3 * t2 + t1 * t2 + t1 * t3) * a**2 - (t1 * t2 * t3) * a
+            )
+            int4 = int4 / int4_denom
+            total = int1 + int2 + int3 + int4
+            return [int1 / total, int2 / total, int3 / total, int4 / total]
+        raise ValueError(f"Unsupported solver order: {order}.")
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        num_inference_steps: int = 25,
+        guidance_scale: float = 4.0,
+        timeshift: float = 3.0,
+        order: int = 2,
+        guidance_interval_min: float = 0.0,
+        guidance_interval_max: float = 1.0,
+        last_step: Optional[float] = None,
+    ) -> None:
+        self.num_inference_steps = int(num_inference_steps)
+        self.guidance_scale = float(guidance_scale)
+        self.timeshift = float(timeshift)
+        self.order = int(order)
+        self.guidance_interval_min = float(guidance_interval_min)
+        self.guidance_interval_max = float(guidance_interval_max)
+        self.last_step = last_step
+        self._reset_state()
+
+    @classmethod
+    def from_sampler_spec(cls, sampler_spec: Dict[str, Any]) -> "PixNerdFlowMatchScheduler":
+        init_args = dict(sampler_spec.get("init_args", {}))
+        return cls(
+            num_inference_steps=int(init_args.get("num_steps", 25)),
+            guidance_scale=float(init_args.get("guidance", 4.0)),
+            timeshift=float(init_args.get("timeshift", 3.0)),
+            order=int(init_args.get("order", 2)),
+            guidance_interval_min=float(init_args.get("guidance_interval_min", 0.0)),
+            guidance_interval_max=float(init_args.get("guidance_interval_max", 1.0)),
+            last_step=init_args.get("last_step"),
+        )
+
+    def _reset_state(self) -> None:
+        self.timesteps: Optional[torch.Tensor] = None
+        self._timedeltas: Optional[torch.Tensor] = None
+        self._solver_coeffs = None
+        self._model_outputs = []
+        self._step_index = 0
+
+    @staticmethod
+    def _shift_respace_fn(t: torch.Tensor, shift: float = 3.0) -> torch.Tensor:
+        return t / (t + (1 - t) * shift)
+
+    def _build_solver_state(
+        self,
+        num_inference_steps: int,
+        timeshift: float,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, List[List[float]]]:
+        last_step = self.last_step
+        if last_step is None:
+            last_step = 1.0 / float(num_inference_steps)
+
+        endpoints = torch.linspace(0.0, 1 - float(last_step), int(num_inference_steps), dtype=torch.float32)
+        endpoints = torch.cat([endpoints, torch.tensor([1.0], dtype=torch.float32)], dim=0)
+        timesteps = self._shift_respace_fn(endpoints, timeshift).to(device=device)
+        timedeltas = (timesteps[1:] - timesteps[:-1]).to(device=device)
+
+        solver_coeffs: List[List[float]] = [[] for _ in range(int(num_inference_steps))]
+        for i in range(int(num_inference_steps)):
+            order = min(self.order, i + 1)
+            pre_ts = timesteps[: i + 1]
+            coeffs = self._lagrange_coeffs(order, pre_ts, pre_ts[i], timesteps[i + 1])
+            solver_coeffs[i] = coeffs
+        return timesteps[:-1], timedeltas, solver_coeffs
+
+    def set_timesteps(
+        self,
+        num_inference_steps: Optional[int] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        timeshift: Optional[float] = None,
+        guidance_scale: Optional[float] = None,
+        order: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        if num_inference_steps is not None:
+            self.num_inference_steps = int(num_inference_steps)
+        if timeshift is not None:
+            self.timeshift = float(timeshift)
+        if guidance_scale is not None:
+            self.guidance_scale = float(guidance_scale)
+        if order is not None:
+            self.order = int(order)
+
+        timesteps, timedeltas, solver_coeffs = self._build_solver_state(
+            self.num_inference_steps,
+            self.timeshift,
+            device=device,
+        )
+        self.timesteps = timesteps
+        self._timedeltas = timedeltas
+        self._solver_coeffs = solver_coeffs
+        self._model_outputs = []
+        self._step_index = 0
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return sample
+
+    def classifier_free_guidance(self, model_output: torch.Tensor) -> torch.Tensor:
+        if model_output.shape[0] % 2 != 0:
+            raise ValueError("Classifier-free guidance expects concatenated unconditional/conditional batches.")
+        uncond, cond = model_output.chunk(2, dim=0)
+        return uncond + self.guidance_scale * (cond - uncond)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+        **kwargs: Any,
+    ) -> Union[PixNerdSchedulerOutput, Tuple[torch.Tensor]]:
+        if self.timesteps is None or self._timedeltas is None or self._solver_coeffs is None:
+            raise RuntimeError("`set_timesteps` must be called before `step`.")
+        if self._step_index >= len(self._solver_coeffs):
+            raise RuntimeError("Scheduler step index exceeded configured timesteps.")
+
+        coeffs = self._solver_coeffs[self._step_index]
+        self._model_outputs.append(model_output)
+        order = len(coeffs)
+        pred = torch.zeros_like(model_output)
+        recent = self._model_outputs[-order:]
+        for coeff, output in zip(coeffs, recent):
+            pred = pred + coeff * output
+
+        prev_sample = sample + pred * self._timedeltas[self._step_index]
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return PixNerdSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        alpha = timesteps.view(-1, 1, 1, 1)
+        sigma = (1.0 - timesteps).view(-1, 1, 1, 1)
+        return alpha * original_samples + sigma * noise
+
+__all__ = [
+    "PixNerdFlowMatchScheduler",
+    "PixNerdSchedulerOutput",
+]

PixNerd-XL-16-256/transformer/config.json

@@ -3,13 +3,13 @@
   "_diffusers_version": "0.36.0",
   "compile_denoiser": false,
   "conditioner_spec": {
-    "class_path": "diffusers_modules.local.modeling_pixnerd_transformer_2d.LabelConditioner",
+    "class_path": "modeling_pixnerd_transformer_2d.LabelConditioner",
     "init_args": {
       "num_classes": 1000
     }
   },
   "denoiser_spec": {
-    "class_path": "diffusers_modules.local.modeling_pixnerd_transformer_2d.PixNerDiT",
+    "class_path": "modeling_pixnerd_transformer_2d.PixNerDiT",
     "init_args": {
       "hidden_size": 1152,
       "hidden_size_x": 64,
@@ -26,7 +26,7 @@
   "ema_decay": 0.9999,
   "use_ema": true,
   "vae_spec": {
-    "class_path": "diffusers_modules.local.modeling_pixnerd_transformer_2d.PixelAE",
+    "class_path": "modeling_pixnerd_transformer_2d.PixelAE",
     "init_args": {
       "scale": 1.0
     }

PixNerd-XL-16-256/transformer/modeling_pixnerd_transformer_2d.py

@@ -0,0 +1,749 @@
+from __future__ import annotations
+
+import copy
+import importlib
+import math
+import sys
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+from torch.nn.functional import scaled_dot_product_attention
+
+class BaseAE(torch.nn.Module):
+    def __init__(self, scale=1.0, shift=0.0):
+        super().__init__()
+        self.scale = scale
+        self.shift = shift
+
+    def encode(self, x):
+        return self._impl_encode(x) #.to(torch.bfloat16)
+
+    # @torch.autocast("cuda", dtype=torch.bfloat16)
+    def decode(self, x):
+        return self._impl_decode(x) #.to(torch.bfloat16)
+
+    def _impl_encode(self, x):
+        raise NotImplementedError
+
+    def _impl_decode(self, x):
+        raise NotImplementedError
+
+def uint82fp(x):
+    x = x.to(torch.float32)
+    x = (x - 127.5) / 127.5
+    return x
+
+def fp2uint8(x):
+    x = torch.clip_((x + 1) * 127.5 + 0.5, 0, 255).to(torch.uint8)
+    return x
+
+
+class PixelAE(BaseAE):
+    def __init__(self, scale=1.0, shift=0.0):
+        super().__init__(scale, shift)
+
+    def _impl_encode(self, x):
+        return x/self.scale+self.shift
+
+    def _impl_decode(self, x):
+        return (x-self.shift)*self.scale
+
+
+def resolve_conditioner_device(metadata: dict, fallback: torch.device | None = None) -> torch.device:
+    if metadata is None:
+        metadata = {}
+    if "device" in metadata and metadata["device"] is not None:
+        return torch.device(metadata["device"])
+    if fallback is not None:
+        return fallback
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class BaseConditioner(nn.Module):
+    def __init__(self):
+        super(BaseConditioner, self).__init__()
+
+    def _impl_condition(self, y, metadata)->torch.Tensor:
+        raise NotImplementedError()
+
+    def _impl_uncondition(self, y, metadata)->torch.Tensor:
+        raise NotImplementedError()
+
+    @torch.no_grad()
+    def __call__(self, y, metadata:dict={}):
+        condition = self._impl_condition(y, metadata)
+        uncondition = self._impl_uncondition(y, metadata)
+        if condition.dtype in [torch.float64, torch.float32, torch.float16]:
+            condition = condition.to(torch.bfloat16)
+        if uncondition.dtype in [torch.float64,torch.float32, torch.float16]:
+            uncondition = uncondition.to(torch.bfloat16)
+        return condition, uncondition
+
+
+class ComposeConditioner(BaseConditioner):
+    def __init__(self, conditioners:List[BaseConditioner]):
+        super().__init__()
+        self.conditioners = conditioners
+
+    def _impl_condition(self, y, metadata):
+        condition = []
+        for conditioner in self.conditioners:
+            condition.append(conditioner._impl_condition(y, metadata))
+        condition = torch.cat(condition, dim=1)
+        return condition
+
+    def _impl_uncondition(self, y, metadata):
+        uncondition = []
+        for conditioner in self.conditioners:
+            uncondition.append(conditioner._impl_uncondition(y, metadata))
+        uncondition = torch.cat(uncondition, dim=1)
+        return uncondition
+
+
+class LabelConditioner(BaseConditioner):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.null_condition = num_classes
+
+    def _impl_condition(self, y, metadata):
+        device = resolve_conditioner_device(metadata)
+        return torch.tensor(y, device=device).long()
+
+    def _impl_uncondition(self, y, metadata):
+        device = resolve_conditioner_device(metadata)
+        return torch.full((len(y),), self.null_condition, dtype=torch.long, device=device)
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale) + shift
+
+class Embed(nn.Module):
+    def __init__(
+            self,
+            in_chans: int = 3,
+            embed_dim: int = 768,
+            norm_layer = None,
+            bias: bool = True,
+    ):
+        super().__init__()
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.proj = nn.Linear(in_chans, embed_dim, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+    def forward(self, x):
+        x = self.proj(x)
+        x = self.norm(x)
+        return x
+
+class TimestepEmbedder(nn.Module):
+
+    def __init__(self, hidden_size, frequency_embedding_size=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, dim, max_period=10):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+        )
+        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):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes, hidden_size):
+        super().__init__()
+        self.embedding_table = nn.Embedding(num_classes, hidden_size)
+        self.num_classes = num_classes
+
+    def forward(self, labels,):
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.Linear(hidden_size, 2*hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+    ):
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+    def forward(self, x):
+        x =  self.w2(torch.nn.functional.silu(self.w1(x)) * self.w3(x))
+        return x
+
+def precompute_freqs_cis_2d(dim: int, height: int, width:int, theta: float = 10000.0, scale=16.0):
+    # assert  H * H == end
+    # flat_patch_pos = torch.linspace(-1, 1, end) # N = end
+    x_pos = torch.linspace(0, scale, width)
+    y_pos = torch.linspace(0, scale, height)
+    y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij")
+    y_pos = y_pos.reshape(-1)
+    x_pos = x_pos.reshape(-1)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) # Hc/4
+    x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
+    y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
+    x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
+    y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
+    freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) # N,Hc/4,2
+    freqs_cis = freqs_cis.reshape(height*width, -1)
+    return freqs_cis
+
+
+def apply_rotary_emb(
+        xq: torch.Tensor,
+        xk: torch.Tensor,
+        freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    freqs_cis = freqs_cis[None, :, None, :]
+    # xq : B N H Hc
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # B N H Hc/2
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) # B, N, H, Hc
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class RAttention(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = True,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = RMSNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor, pos, mask) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # B N H Hc
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q, k = apply_rotary_emb(q, k, freqs_cis=pos)
+        q = q.view(B, -1, self.num_heads, C // self.num_heads).transpose(1, 2)  # B, H, N, Hc
+        k = k.view(B, -1, self.num_heads, C // self.num_heads).transpose(1, 2).contiguous()  # B, H, N, Hc
+        v = v.view(B, -1, self.num_heads, C // self.num_heads).transpose(1, 2).contiguous()
+
+        x = scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+
+class FlattenDiTBlock(nn.Module):
+    def __init__(self, hidden_size, groups,  mlp_ratio=4.0, ):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
+        self.attn = RAttention(hidden_size, num_heads=groups, qkv_bias=False)
+        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp = FeedForward(hidden_size, mlp_hidden_dim)
+        self.adaLN_modulation = nn.Sequential(
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, x,  c, pos, mask=None):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
+        x = x + gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), pos, mask=mask)
+        x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+class NerfEmbedder(nn.Module):
+    def __init__(self, in_channels, hidden_size_input, max_freqs):
+        super().__init__()
+        self.max_freqs = max_freqs
+        self.hidden_size_input = hidden_size_input
+        self.embedder = nn.Sequential(
+            nn.Linear(in_channels+max_freqs**2, hidden_size_input, bias=True),
+        )
+
+    @lru_cache
+    def fetch_pos(self, patch_size, device, dtype):
+        pos_x = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y, pos_x = torch.meshgrid(pos_y, pos_x, indexing="ij")
+        pos_x = pos_x.reshape(-1, 1, 1)
+        pos_y = pos_y.reshape(-1, 1, 1)
+
+        freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device)
+        freqs_x = freqs[None, :, None]
+        freqs_y = freqs[None, None, :]
+        coeffs = (1 + freqs_x * freqs_y) ** -1
+        dct_x = torch.cos(pos_x * freqs_x * torch.pi)
+        dct_y = torch.cos(pos_y * freqs_y * torch.pi)
+        dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
+        return dct
+
+
+    def forward(self, inputs):
+        B, P2, C = inputs.shape
+        patch_size = int(P2 ** 0.5)
+        device = inputs.device
+        dtype = inputs.dtype
+        dct = self.fetch_pos(patch_size, device, dtype)
+        dct = dct.repeat(B, 1, 1)
+        inputs = torch.cat([inputs, dct], dim=-1)
+        inputs = self.embedder(inputs)
+        return inputs
+
+
+class NerfBlock(nn.Module):
+    def __init__(self, hidden_size_s, hidden_size_x, mlp_ratio=4):
+        super().__init__()
+        self.param_generator1 = nn.Sequential(
+            nn.Linear(hidden_size_s, 2*hidden_size_x**2*mlp_ratio, bias=True),
+        )
+        self.norm = RMSNorm(hidden_size_x, eps=1e-6)
+        self.mlp_ratio = mlp_ratio
+    def forward(self, x, s):
+        batch_size, num_x, hidden_size_x = x.shape
+        mlp_params1 = self.param_generator1(s)
+        fc1_param1, fc2_param1 = mlp_params1.chunk(2, dim=-1)
+        fc1_param1 = fc1_param1.view(batch_size, hidden_size_x, hidden_size_x*self.mlp_ratio)
+        fc2_param1 = fc2_param1.view(batch_size, hidden_size_x*self.mlp_ratio, hidden_size_x)
+
+        # normalize fc1
+        normalized_fc1_param1 = torch.nn.functional.normalize(fc1_param1, dim=-2)
+        # normalize fc2
+        normalized_fc2_param1 = torch.nn.functional.normalize(fc2_param1, dim=-2)
+        # mlp 1
+        res_x = x
+        x = self.norm(x)
+        x = torch.bmm(x, normalized_fc1_param1)
+        x = torch.nn.functional.silu(x)
+        x = torch.bmm(x, normalized_fc2_param1)
+        x = x + res_x
+        return x
+
+class NerfFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.linear(x)
+        return x
+
+class PixNerDiT(nn.Module):
+    def __init__(
+            self,
+            in_channels=4,
+            num_groups=12,
+            hidden_size=1152,
+            hidden_size_x=64,
+            nerf_mlpratio=4,
+            num_blocks=18,
+            num_cond_blocks=4,
+            patch_size=2,
+            num_classes=1000,
+            learn_sigma=True,
+            deep_supervision=0,
+            weight_path=None,
+            load_ema=False,
+    ):
+        super().__init__()
+        self.deep_supervision = deep_supervision
+        self.learn_sigma = learn_sigma
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.hidden_size = hidden_size
+        self.num_groups = num_groups
+        self.num_blocks = num_blocks
+        self.num_cond_blocks = num_cond_blocks
+        self.patch_size = patch_size
+        self.x_embedder = NerfEmbedder(in_channels, hidden_size_x, max_freqs=8)
+        self.s_embedder = Embed(in_channels*patch_size**2, hidden_size, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes+1, hidden_size)
+
+        self.final_layer = NerfFinalLayer(hidden_size_x, self.out_channels)
+
+        self.weight_path = weight_path
+
+        self.load_ema = load_ema
+        self.blocks = nn.ModuleList([
+            FlattenDiTBlock(self.hidden_size, self.num_groups) for _ in range(self.num_cond_blocks)
+        ])
+        self.blocks.extend([
+            NerfBlock(self.hidden_size, hidden_size_x, nerf_mlpratio) for _ in range(self.num_cond_blocks, self.num_blocks)
+        ])
+        self.initialize_weights()
+        self.precompute_pos = dict()
+
+    def fetch_pos(self, height, width, device):
+        if (height, width) in self.precompute_pos:
+            return self.precompute_pos[(height, width)].to(device)
+        else:
+            pos = precompute_freqs_cis_2d(self.hidden_size // self.num_groups, height, width).to(device)
+            self.precompute_pos[(height, width)] = pos
+            return pos
+
+    def initialize_weights(self):
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.s_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.s_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # zero init final layer
+        nn.init.zeros_(self.final_layer.linear.weight)
+        nn.init.zeros_(self.final_layer.linear.bias)
+
+
+    def forward(self, x, t, y, s=None, mask=None):
+        B, _, H, W = x.shape
+        pos = self.fetch_pos(H//self.patch_size, W//self.patch_size, x.device)
+        x = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2)
+        t = self.t_embedder(t.view(-1)).view(B, -1, self.hidden_size)
+        y = self.y_embedder(y).view(B, 1, self.hidden_size)
+        c = nn.functional.silu(t + y)
+        if s is None:
+            s = self.s_embedder(x)
+            for i in range(self.num_cond_blocks):
+                s = self.blocks[i](s, c, pos, mask)
+            s = nn.functional.silu(t + s)
+        batch_size, length, _ = s.shape
+        x = x.reshape(batch_size*length, self.in_channels, self.patch_size**2)
+        x = x.transpose(1, 2)
+        s = s.view(batch_size*length, self.hidden_size)
+        x = self.x_embedder(x)
+        for i in range(self.num_cond_blocks, self.num_blocks):
+            x = self.blocks[i](x, s)
+        x = self.final_layer(x)
+        x = x.transpose(1, 2)
+        x = x.reshape(batch_size, length, -1)
+        x = torch.nn.functional.fold(x.transpose(1, 2).contiguous(), (H, W), kernel_size=self.patch_size, stride=self.patch_size)
+        return x
+
+
+def to_container(config: Any) -> Any:
+    if hasattr(config, "items") and not isinstance(config, dict):
+        return {k: to_container(v) for k, v in config.items()}
+    if isinstance(config, list):
+        return [to_container(v) for v in config]
+    return config
+
+
+def load_symbol(path: str) -> Any:
+    module_path, name = path.rsplit(".", 1)
+    if module_path in {__name__, "modeling_pixnerd_transformer_2d"}:
+        return getattr(sys.modules[__name__], name)
+    module = importlib.import_module(module_path)
+    return getattr(module, name)
+
+
+def instantiate_from_spec(spec: Any) -> Any:
+    spec = to_container(spec)
+    if isinstance(spec, dict) and "class_path" in spec:
+        class_or_fn = load_symbol(spec["class_path"])
+        init_args = spec.get("init_args", {})
+        if isinstance(init_args, dict):
+            init_args = {k: instantiate_from_spec(v) for k, v in init_args.items()}
+        return class_or_fn(**init_args)
+    if isinstance(spec, dict):
+        return {k: instantiate_from_spec(v) for k, v in spec.items()}
+    if isinstance(spec, list):
+        return [instantiate_from_spec(v) for v in spec]
+    if isinstance(spec, str) and "." in spec:
+        try:
+            return load_symbol(spec)
+        except Exception:
+            return spec
+    return spec
+
+
+def clone_spec(spec: Dict[str, Any]) -> Dict[str, Any]:
+    return copy.deepcopy(to_container(spec))
+
+
+def load_prefixed_state_dict(
+    module: Optional[torch.nn.Module],
+    state_dict: Dict[str, torch.Tensor],
+    prefixes: Iterable[str],
+) -> bool:
+    if module is None:
+        return False
+    for prefix in prefixes:
+        subset = {
+            key[len(prefix) :]: value
+            for key, value in state_dict.items()
+            if key.startswith(prefix)
+        }
+        if subset:
+            module.load_state_dict(subset, strict=False)
+            return True
+    return False
+
+
+@dataclass
+class PixNerdTransformer2DModelOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+class PixNerdTransformer2DModel(ModelMixin, ConfigMixin):
+    config_name = "config.json"
+
+    @register_to_config
+    def __init__(
+        self,
+        denoiser_spec: Dict[str, Any],
+        conditioner_spec: Dict[str, Any],
+        vae_spec: Optional[Dict[str, Any]] = None,
+        diffusion_trainer_spec: Optional[Dict[str, Any]] = None,
+        use_ema: bool = True,
+        ema_decay: float = 0.9999,
+        compile_denoiser: bool = False,
+    ) -> None:
+        super().__init__()
+        self.denoiser = instantiate_from_spec(to_container(denoiser_spec))
+        self.conditioner = instantiate_from_spec(to_container(conditioner_spec))
+        self.vae = instantiate_from_spec(to_container(vae_spec)) if vae_spec is not None else None
+        self.diffusion_trainer = (
+            instantiate_from_spec(to_container(diffusion_trainer_spec))
+            if diffusion_trainer_spec is not None
+            else None
+        )
+
+        self.use_ema = bool(use_ema)
+        self.ema_decay = float(ema_decay)
+        self.ema_denoiser = copy.deepcopy(self.denoiser) if self.use_ema else None
+        if self.ema_denoiser is not None:
+            self.ema_denoiser.to(torch.float32)
+
+        if compile_denoiser and hasattr(self.denoiser, "compile"):
+            self.denoiser.compile()
+            if self.ema_denoiser is not None:
+                self.ema_denoiser.compile()
+
+        self._freeze_non_trainable_modules()
+        if self.ema_denoiser is not None:
+            self.sync_ema()
+
+    @property
+    def patch_size(self) -> int:
+        return int(getattr(self.denoiser, "patch_size", 1))
+
+    @property
+    def in_channels(self) -> int:
+        return int(getattr(self.denoiser, "in_channels", 3))
+
+    @classmethod
+    def from_project_config(
+        cls,
+        model_config: Dict[str, Any],
+        use_ema: bool = True,
+        compile_denoiser: bool = False,
+    ) -> "PixNerdTransformer2DModel":
+        model_config = to_container(model_config)
+        ema_decay = model_config.get("ema_tracker", {}).get("init_args", {}).get("decay", 0.9999)
+        return cls(
+            denoiser_spec=model_config["denoiser"],
+            conditioner_spec=model_config["conditioner"],
+            vae_spec=model_config.get("vae"),
+            diffusion_trainer_spec=model_config.get("diffusion_trainer"),
+            use_ema=use_ema,
+            ema_decay=ema_decay,
+            compile_denoiser=compile_denoiser,
+        )
+
+    @staticmethod
+    def _as_timestep_tensor(
+        timestep: Any,
+        batch_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if isinstance(timestep, torch.Tensor):
+            if timestep.ndim == 0:
+                return timestep.repeat(batch_size).to(device=device, dtype=torch.float32)
+            return timestep.to(device=device, dtype=torch.float32)
+        return torch.full((batch_size,), float(timestep), device=device, dtype=torch.float32)
+
+    def _freeze_module(self, module: Optional[torch.nn.Module]) -> None:
+        if module is None:
+            return
+        module.eval()
+        for parameter in module.parameters():
+            parameter.requires_grad = False
+
+    def _freeze_non_trainable_modules(self) -> None:
+        self._freeze_module(self.conditioner)
+        self._freeze_module(self.vae)
+        self._freeze_module(self.ema_denoiser)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Any,
+        encoder_hidden_states: torch.Tensor,
+        return_dict: bool = True,
+    ) -> PixNerdTransformer2DModelOutput | Tuple[torch.Tensor]:
+        t = self._as_timestep_tensor(timestep, sample.shape[0], sample.device)
+        out = self.denoiser(sample, t, encoder_hidden_states)
+        if not return_dict:
+            return (out,)
+        return PixNerdTransformer2DModelOutput(sample=out)
+
+    def predict_noise(
+        self,
+        sample: torch.Tensor,
+        timestep: Any,
+        encoder_hidden_states: torch.Tensor,
+        use_ema: bool = False,
+    ) -> torch.Tensor:
+        t = self._as_timestep_tensor(timestep, sample.shape[0], sample.device)
+        denoiser = self.get_inference_denoiser(use_ema=use_ema)
+        return denoiser(sample, t, encoder_hidden_states)
+
+    def get_inference_denoiser(self, use_ema: bool = True) -> torch.nn.Module:
+        if use_ema and self.ema_denoiser is not None:
+            return self.ema_denoiser
+        return self.denoiser
+
+    @torch.no_grad()
+    def get_conditioning(
+        self,
+        y: Iterable[Any],
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        metadata = {} if metadata is None else metadata
+        return self.conditioner(y, metadata)
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        if self.vae is None:
+            return x
+        return self.vae.encode(x)
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        if self.vae is None:
+            return latents
+        return self.vae.decode(latents)
+
+    @torch.no_grad()
+    def sync_ema(self) -> None:
+        if self.ema_denoiser is None:
+            return
+        self.ema_denoiser.load_state_dict(self.denoiser.state_dict(), strict=True)
+        self.ema_denoiser.to(torch.float32)
+
+    @torch.no_grad()
+    def ema_step(self, decay: Optional[float] = None) -> None:
+        if self.ema_denoiser is None:
+            return
+        decay = self.ema_decay if decay is None else float(decay)
+        for ema_param, param in zip(self.ema_denoiser.parameters(), self.denoiser.parameters()):
+            ema_param.mul_(decay).add_(param.detach().float(), alpha=1.0 - decay)
+
+    def compute_training_loss(
+        self,
+        x: torch.Tensor,
+        y: Iterable[Any],
+        scheduler: torch.nn.Module,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, torch.Tensor]:
+        if self.diffusion_trainer is None:
+            raise RuntimeError("diffusion_trainer is not configured.")
+        metadata = {} if metadata is None else metadata
+
+        with torch.no_grad():
+            x = self.encode(x)
+            condition, uncondition = self.get_conditioning(y, metadata)
+
+        return self.diffusion_trainer(
+            self.denoiser,
+            self.ema_denoiser if self.ema_denoiser is not None else self.denoiser,
+            scheduler,
+            x,
+            condition,
+            uncondition,
+            metadata,
+        )
+
+__all__ = [
+    "PixNerDiT",
+    "LabelConditioner",
+    "PixelAE",
+    "PixNerdTransformer2DModel",
+    "PixNerdTransformer2DModelOutput",
+]

PixNerd-XL-16-512/README.md

@@ -0,0 +1,41 @@
+# PixNerd-XL-16-512
+
+Self-contained PixNerd-XL/16 checkpoint inside [`BiliSakura/PixNerd-diffusers`](https://huggingface.co/BiliSakura/PixNerd-diffusers). Runtime dependencies: this folder + PyPI `diffusers`/`torch` only.
+
+## Hub path
+
+`BiliSakura/PixNerd-diffusers/PixNerd-XL-16-512`
+
+## Layout
+
+```text
+PixNerd-XL-16-512/
+├── pipeline.py
+├── model_index.json
+├── conversion_metadata.json
+├── transformer/
+└── scheduler/
+```
+
+## Load
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/PixNerd-diffusers/PixNerd-XL-16-512",
+    trust_remote_code=True,
+    torch_dtype=torch.float32,
+).to("cuda")
+
+images = pipe(
+    prompt=207,
+    height=512,
+    width=512,
+    num_inference_steps=25,
+    guidance_scale=4.0,
+    timeshift=3.0,
+    order=2,
+).images
+```

PixNerd-XL-16-512/model_index.json

@@ -1,12 +1,15 @@
 {
-  "_class_name": "PixNerdPipeline",
-  "_diffusers_version": "0.30.0",
+  "_class_name": [
+    "pipeline",
+    "PixNerdPipeline"
+  ],
+  "_diffusers_version": "0.36.0",
   "scheduler": [
-    "diffusers_modules.local.scheduling_pixnerd_flow_match",
+    "scheduling_pixnerd_flow_match",
     "PixNerdFlowMatchScheduler"
   ],
   "transformer": [
-    "diffusers_modules.local.modeling_pixnerd_transformer_2d",
+    "modeling_pixnerd_transformer_2d",
     "PixNerdTransformer2DModel"
   ]
 }

PixNerd-XL-16-512/pipeline.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import sys
 from dataclasses import dataclass
-from typing import List, Optional, Sequence, Union
+from pathlib import Path
+from typing import List, Literal, Optional, Sequence, Union
 
 import torch
 from diffusers import DiffusionPipeline
@@ -9,10 +11,8 @@ from diffusers.image_processor import VaeImageProcessor
 from diffusers.utils import BaseOutput
 from PIL import Image
 
-from .modeling_pixnerd_transformer_2d import PixNerdTransformer2DModel
-from .scheduling_pixnerd_flow_match import PixNerdFlowMatchScheduler
-
 ConditioningInput = Union[str, int, Sequence[Union[str, int]]]
+Language = Literal["en", "cn"]
 
 
 @dataclass
@@ -27,9 +27,11 @@ class PixNerdPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler: PixNerdFlowMatchScheduler,
+        scheduler,
         vae=None,
         conditioner=None,
+        id2label: Optional[dict[int, str]] = None,
+        id2label_cn: Optional[dict[int, str]] = None,
     ):
         super().__init__()
         if vae is None:
@@ -46,6 +48,170 @@ class PixNerdPipeline(DiffusionPipeline):
         )
         self.image_processor = VaeImageProcessor(vae_scale_factor=1)
 
+        if id2label is None and id2label_cn is None:
+            id2label, id2label_cn = self._load_repo_labels()
+        self._id2label = id2label or {}
+        self._id2label_cn = id2label_cn or {}
+        self.labels = self._build_label2id(self._id2label)
+        self.labels_cn = self._build_label2id(self._id2label_cn)
+        self._labels_loaded_from_path = bool(self._id2label or self._id2label_cn)
+
+    def _ensure_labels_loaded(self) -> None:
+        if self._labels_loaded_from_path:
+            return
+
+        path = getattr(getattr(self, "config", None), "_name_or_path", None) or getattr(self, "_name_or_path", None)
+        if not path:
+            return
+
+        id2label, id2label_cn = self._load_labels_for_path(path)
+        if id2label is None and id2label_cn is None:
+            self._labels_loaded_from_path = True
+            return
+
+        self._id2label = id2label or {}
+        self._id2label_cn = id2label_cn or {}
+        self.labels = self._build_label2id(self._id2label)
+        self.labels_cn = self._build_label2id(self._id2label_cn)
+        self._labels_loaded_from_path = True
+
+    @staticmethod
+    def _resolve_labels_dir(pretrained_model_name_or_path: Union[str, Path]) -> Optional[Path]:
+        path = Path(pretrained_model_name_or_path)
+        if not path.exists():
+            try:
+                from huggingface_hub import snapshot_download
+
+                path = Path(snapshot_download(pretrained_model_name_or_path))
+            except Exception:
+                return None
+
+        if (path / "model_index.json").exists():
+            labels_dir = path.parent / "labels"
+        else:
+            labels_dir = path / "labels"
+        return labels_dir if labels_dir.is_dir() else None
+
+    @classmethod
+    def _load_labels_for_path(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+    ) -> tuple[Optional[dict[int, str]], Optional[dict[int, str]]]:
+        labels_dir = cls._resolve_labels_dir(pretrained_model_name_or_path)
+        if labels_dir is None:
+            return None, None
+
+        labels_path = str(labels_dir)
+        inserted = False
+        if labels_path not in sys.path:
+            sys.path.insert(0, labels_path)
+            inserted = True
+        try:
+            from imagenet_labels import load_id2label
+
+            return (
+                load_id2label(labels_dir, lang="en"),
+                load_id2label(labels_dir, lang="cn"),
+            )
+        finally:
+            if inserted and labels_path in sys.path:
+                sys.path.remove(labels_path)
+
+    @staticmethod
+    def _load_repo_labels() -> tuple[Optional[dict[int, str]], Optional[dict[int, str]]]:
+        labels_dir = Path(__file__).resolve().parent.parent / "labels"
+        if not labels_dir.is_dir():
+            return None, None
+
+        labels_path = str(labels_dir)
+        inserted = False
+        if labels_path not in sys.path:
+            sys.path.insert(0, labels_path)
+            inserted = True
+        try:
+            from imagenet_labels import load_id2label
+
+            return (
+                load_id2label(labels_dir, lang="en"),
+                load_id2label(labels_dir, lang="cn"),
+            )
+        finally:
+            if inserted and labels_path in sys.path:
+                sys.path.remove(labels_path)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        pipe = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        id2label, id2label_cn = cls._load_labels_for_path(pretrained_model_name_or_path)
+        if id2label is not None or id2label_cn is not None:
+            pipe._id2label = id2label or {}
+            pipe._id2label_cn = id2label_cn or {}
+            pipe.labels = cls._build_label2id(pipe._id2label)
+            pipe.labels_cn = cls._build_label2id(pipe._id2label_cn)
+        return pipe
+
+    @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 id2label(self) -> dict[int, str]:
+        self._ensure_labels_loaded()
+        return self._id2label
+
+    @property
+    def id2label_cn(self) -> dict[int, str]:
+        self._ensure_labels_loaded()
+        return self._id2label_cn
+
+    def get_label_ids(
+        self,
+        labels: Union[str, List[str]],
+        *,
+        lang: Language = "en",
+    ) -> List[int]:
+        self._ensure_labels_loaded()
+        if isinstance(labels, str):
+            labels = [labels]
+
+        label2id = self.labels if lang == "en" else self.labels_cn
+        if not label2id:
+            raise ValueError(
+                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
+            )
+
+        missing = [label for label in labels if label not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(
+                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+            )
+        return [label2id[label] for label in labels]
+
+    def _resolve_prompt_item(self, value: Union[str, int]) -> int:
+        if isinstance(value, int):
+            return value
+        if value.isdigit():
+            return int(value)
+        if value in self.labels:
+            return self.labels[value]
+        if value in self.labels_cn:
+            return self.labels_cn[value]
+        raise ValueError(
+            f"Unknown class label {value!r}. Pass an ImageNet class id or a synonym from "
+            "`pipe.labels` / `pipe.labels_cn`."
+        )
+
+    def _resolve_prompts(self, prompts: List[Union[str, int]]) -> List[int]:
+        self._ensure_labels_loaded()
+        return [self._resolve_prompt_item(prompt) for prompt in prompts]
+
     @staticmethod
     def _fp_to_uint8(image: torch.Tensor) -> torch.Tensor:
         return torch.clip_((image + 1) * 127.5 + 0.5, 0, 255).to(torch.uint8)
@@ -71,10 +237,11 @@ class PixNerdPipeline(DiffusionPipeline):
         num_images_per_prompt: int,
     ):
         prompts = self._repeat(self._to_list(prompt), num_images_per_prompt)
+        resolved = self._resolve_prompts(prompts)
         metadata = {"device": self._execution_device}
         with torch.no_grad():
-            cond, uncond = self.conditioner(prompts, metadata)
-        return cond, uncond, prompts
+            cond, uncond = self.conditioner(resolved, metadata)
+        return cond, uncond, resolved
 
     def prepare_latents(
         self,
@@ -124,9 +291,10 @@ class PixNerdPipeline(DiffusionPipeline):
         cond, default_uncond, prompts = self.encode_prompt(prompt, num_images_per_prompt)
         if negative_prompt is not None:
             negative = self._repeat(self._to_list(negative_prompt), num_images_per_prompt)
+            resolved_negative = self._resolve_prompts(negative)
             metadata = {"device": self._execution_device}
             with torch.no_grad():
-                _, uncond = self.conditioner(negative, metadata)
+                _, uncond = self.conditioner(resolved_negative, metadata)
         else:
             uncond = default_uncond
         batch_size = len(prompts)
@@ -178,6 +346,7 @@ class PixNerdPipeline(DiffusionPipeline):
             return (output,)
         return PixNerdPipelineOutput(images=output)
 
+
 __all__ = [
     "PixNerdPipeline",
     "PixNerdPipelineOutput",

PixNerd-XL-16-512/scheduler/scheduling_pixnerd_flow_match.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, 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 PixNerdSchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+
+
+class PixNerdFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Diffusers-compatible scheduler wrapper for PixNerd's AdamLM flow-matching sampler.
+    """
+
+    config_name = "scheduler_config.json"
+    order = 1
+    init_noise_sigma = 1.0
+
+    @staticmethod
+    def _lagrange_coeffs(order: int, pre_ts: torch.Tensor, t_start: torch.Tensor, t_end: torch.Tensor) -> List[float]:
+        ts = [float(v) for v in pre_ts[-order:].tolist()]
+        a = float(t_start)
+        b = float(t_end)
+
+        if order == 1:
+            return [1.0]
+        if order == 2:
+            t1, t2 = ts
+            int1 = 0.5 / (t1 - t2) * ((b - t2) ** 2 - (a - t2) ** 2)
+            int2 = 0.5 / (t2 - t1) * ((b - t1) ** 2 - (a - t1) ** 2)
+            total = int1 + int2
+            return [int1 / total, int2 / total]
+        if order == 3:
+            t1, t2, t3 = ts
+            int1_denom = (t1 - t2) * (t1 - t3)
+            int1 = ((1 / 3) * b**3 - 0.5 * (t2 + t3) * b**2 + (t2 * t3) * b) - (
+                (1 / 3) * a**3 - 0.5 * (t2 + t3) * a**2 + (t2 * t3) * a
+            )
+            int1 = int1 / int1_denom
+            int2_denom = (t2 - t1) * (t2 - t3)
+            int2 = ((1 / 3) * b**3 - 0.5 * (t1 + t3) * b**2 + (t1 * t3) * b) - (
+                (1 / 3) * a**3 - 0.5 * (t1 + t3) * a**2 + (t1 * t3) * a
+            )
+            int2 = int2 / int2_denom
+            int3_denom = (t3 - t1) * (t3 - t2)
+            int3 = ((1 / 3) * b**3 - 0.5 * (t1 + t2) * b**2 + (t1 * t2) * b) - (
+                (1 / 3) * a**3 - 0.5 * (t1 + t2) * a**2 + (t1 * t2) * a
+            )
+            int3 = int3 / int3_denom
+            total = int1 + int2 + int3
+            return [int1 / total, int2 / total, int3 / total]
+        if order == 4:
+            t1, t2, t3, t4 = ts
+            int1_denom = (t1 - t2) * (t1 - t3) * (t1 - t4)
+            int1 = ((1 / 4) * b**4 - (1 / 3) * (t2 + t3 + t4) * b**3 + 0.5 * (t3 * t4 + t2 * t3 + t2 * t4) * b**2 - (t2 * t3 * t4) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t2 + t3 + t4) * a**3 + 0.5 * (t3 * t4 + t2 * t3 + t2 * t4) * a**2 - (t2 * t3 * t4) * a
+            )
+            int1 = int1 / int1_denom
+            int2_denom = (t2 - t1) * (t2 - t3) * (t2 - t4)
+            int2 = ((1 / 4) * b**4 - (1 / 3) * (t1 + t3 + t4) * b**3 + 0.5 * (t3 * t4 + t1 * t3 + t1 * t4) * b**2 - (t1 * t3 * t4) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t1 + t3 + t4) * a**3 + 0.5 * (t3 * t4 + t1 * t3 + t1 * t4) * a**2 - (t1 * t3 * t4) * a
+            )
+            int2 = int2 / int2_denom
+            int3_denom = (t3 - t1) * (t3 - t2) * (t3 - t4)
+            int3 = ((1 / 4) * b**4 - (1 / 3) * (t1 + t2 + t4) * b**3 + 0.5 * (t4 * t2 + t1 * t2 + t1 * t4) * b**2 - (t1 * t2 * t4) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t1 + t2 + t4) * a**3 + 0.5 * (t4 * t2 + t1 * t2 + t1 * t4) * a**2 - (t1 * t2 * t4) * a
+            )
+            int3 = int3 / int3_denom
+            int4_denom = (t4 - t1) * (t4 - t2) * (t4 - t3)
+            int4 = ((1 / 4) * b**4 - (1 / 3) * (t1 + t2 + t3) * b**3 + 0.5 * (t3 * t2 + t1 * t2 + t1 * t3) * b**2 - (t1 * t2 * t3) * b) - (
+                (1 / 4) * a**4 - (1 / 3) * (t1 + t2 + t3) * a**3 + 0.5 * (t3 * t2 + t1 * t2 + t1 * t3) * a**2 - (t1 * t2 * t3) * a
+            )
+            int4 = int4 / int4_denom
+            total = int1 + int2 + int3 + int4
+            return [int1 / total, int2 / total, int3 / total, int4 / total]
+        raise ValueError(f"Unsupported solver order: {order}.")
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        num_inference_steps: int = 25,
+        guidance_scale: float = 4.0,
+        timeshift: float = 3.0,
+        order: int = 2,
+        guidance_interval_min: float = 0.0,
+        guidance_interval_max: float = 1.0,
+        last_step: Optional[float] = None,
+    ) -> None:
+        self.num_inference_steps = int(num_inference_steps)
+        self.guidance_scale = float(guidance_scale)
+        self.timeshift = float(timeshift)
+        self.order = int(order)
+        self.guidance_interval_min = float(guidance_interval_min)
+        self.guidance_interval_max = float(guidance_interval_max)
+        self.last_step = last_step
+        self._reset_state()
+
+    @classmethod
+    def from_sampler_spec(cls, sampler_spec: Dict[str, Any]) -> "PixNerdFlowMatchScheduler":
+        init_args = dict(sampler_spec.get("init_args", {}))
+        return cls(
+            num_inference_steps=int(init_args.get("num_steps", 25)),
+            guidance_scale=float(init_args.get("guidance", 4.0)),
+            timeshift=float(init_args.get("timeshift", 3.0)),
+            order=int(init_args.get("order", 2)),
+            guidance_interval_min=float(init_args.get("guidance_interval_min", 0.0)),
+            guidance_interval_max=float(init_args.get("guidance_interval_max", 1.0)),
+            last_step=init_args.get("last_step"),
+        )
+
+    def _reset_state(self) -> None:
+        self.timesteps: Optional[torch.Tensor] = None
+        self._timedeltas: Optional[torch.Tensor] = None
+        self._solver_coeffs = None
+        self._model_outputs = []
+        self._step_index = 0
+
+    @staticmethod
+    def _shift_respace_fn(t: torch.Tensor, shift: float = 3.0) -> torch.Tensor:
+        return t / (t + (1 - t) * shift)
+
+    def _build_solver_state(
+        self,
+        num_inference_steps: int,
+        timeshift: float,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, List[List[float]]]:
+        last_step = self.last_step
+        if last_step is None:
+            last_step = 1.0 / float(num_inference_steps)
+
+        endpoints = torch.linspace(0.0, 1 - float(last_step), int(num_inference_steps), dtype=torch.float32)
+        endpoints = torch.cat([endpoints, torch.tensor([1.0], dtype=torch.float32)], dim=0)
+        timesteps = self._shift_respace_fn(endpoints, timeshift).to(device=device)
+        timedeltas = (timesteps[1:] - timesteps[:-1]).to(device=device)
+
+        solver_coeffs: List[List[float]] = [[] for _ in range(int(num_inference_steps))]
+        for i in range(int(num_inference_steps)):
+            order = min(self.order, i + 1)
+            pre_ts = timesteps[: i + 1]
+            coeffs = self._lagrange_coeffs(order, pre_ts, pre_ts[i], timesteps[i + 1])
+            solver_coeffs[i] = coeffs
+        return timesteps[:-1], timedeltas, solver_coeffs
+
+    def set_timesteps(
+        self,
+        num_inference_steps: Optional[int] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        timeshift: Optional[float] = None,
+        guidance_scale: Optional[float] = None,
+        order: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        if num_inference_steps is not None:
+            self.num_inference_steps = int(num_inference_steps)
+        if timeshift is not None:
+            self.timeshift = float(timeshift)
+        if guidance_scale is not None:
+            self.guidance_scale = float(guidance_scale)
+        if order is not None:
+            self.order = int(order)
+
+        timesteps, timedeltas, solver_coeffs = self._build_solver_state(
+            self.num_inference_steps,
+            self.timeshift,
+            device=device,
+        )
+        self.timesteps = timesteps
+        self._timedeltas = timedeltas
+        self._solver_coeffs = solver_coeffs
+        self._model_outputs = []
+        self._step_index = 0
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return sample
+
+    def classifier_free_guidance(self, model_output: torch.Tensor) -> torch.Tensor:
+        if model_output.shape[0] % 2 != 0:
+            raise ValueError("Classifier-free guidance expects concatenated unconditional/conditional batches.")
+        uncond, cond = model_output.chunk(2, dim=0)
+        return uncond + self.guidance_scale * (cond - uncond)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+        **kwargs: Any,
+    ) -> Union[PixNerdSchedulerOutput, Tuple[torch.Tensor]]:
+        if self.timesteps is None or self._timedeltas is None or self._solver_coeffs is None:
+            raise RuntimeError("`set_timesteps` must be called before `step`.")
+        if self._step_index >= len(self._solver_coeffs):
+            raise RuntimeError("Scheduler step index exceeded configured timesteps.")
+
+        coeffs = self._solver_coeffs[self._step_index]
+        self._model_outputs.append(model_output)
+        order = len(coeffs)
+        pred = torch.zeros_like(model_output)
+        recent = self._model_outputs[-order:]
+        for coeff, output in zip(coeffs, recent):
+            pred = pred + coeff * output
+
+        prev_sample = sample + pred * self._timedeltas[self._step_index]
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return PixNerdSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        alpha = timesteps.view(-1, 1, 1, 1)
+        sigma = (1.0 - timesteps).view(-1, 1, 1, 1)
+        return alpha * original_samples + sigma * noise
+
+__all__ = [
+    "PixNerdFlowMatchScheduler",
+    "PixNerdSchedulerOutput",
+]

PixNerd-XL-16-512/transformer/config.json

@@ -3,13 +3,13 @@
   "_diffusers_version": "0.36.0",
   "compile_denoiser": false,
   "conditioner_spec": {
-    "class_path": "diffusers_modules.local.modeling_pixnerd_transformer_2d.LabelConditioner",
+    "class_path": "modeling_pixnerd_transformer_2d.LabelConditioner",
     "init_args": {
       "num_classes": 1000
     }
   },
   "denoiser_spec": {
-    "class_path": "diffusers_modules.local.modeling_pixnerd_transformer_2d.PixNerDiT",
+    "class_path": "modeling_pixnerd_transformer_2d.PixNerDiT",
     "init_args": {
       "hidden_size": 1152,
       "hidden_size_x": 64,
@@ -26,7 +26,7 @@
   "ema_decay": 0.9999,
   "use_ema": true,
   "vae_spec": {
-    "class_path": "diffusers_modules.local.modeling_pixnerd_transformer_2d.PixelAE",
+    "class_path": "modeling_pixnerd_transformer_2d.PixelAE",
     "init_args": {
       "scale": 1.0
     }

PixNerd-XL-16-512/transformer/modeling_pixnerd_transformer_2d.py

@@ -0,0 +1,749 @@
+from __future__ import annotations
+
+import copy
+import importlib
+import math
+import sys
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+from torch.nn.functional import scaled_dot_product_attention
+
+class BaseAE(torch.nn.Module):
+    def __init__(self, scale=1.0, shift=0.0):
+        super().__init__()
+        self.scale = scale
+        self.shift = shift
+
+    def encode(self, x):
+        return self._impl_encode(x) #.to(torch.bfloat16)
+
+    # @torch.autocast("cuda", dtype=torch.bfloat16)
+    def decode(self, x):
+        return self._impl_decode(x) #.to(torch.bfloat16)
+
+    def _impl_encode(self, x):
+        raise NotImplementedError
+
+    def _impl_decode(self, x):
+        raise NotImplementedError
+
+def uint82fp(x):
+    x = x.to(torch.float32)
+    x = (x - 127.5) / 127.5
+    return x
+
+def fp2uint8(x):
+    x = torch.clip_((x + 1) * 127.5 + 0.5, 0, 255).to(torch.uint8)
+    return x
+
+
+class PixelAE(BaseAE):
+    def __init__(self, scale=1.0, shift=0.0):
+        super().__init__(scale, shift)
+
+    def _impl_encode(self, x):
+        return x/self.scale+self.shift
+
+    def _impl_decode(self, x):
+        return (x-self.shift)*self.scale
+
+
+def resolve_conditioner_device(metadata: dict, fallback: torch.device | None = None) -> torch.device:
+    if metadata is None:
+        metadata = {}
+    if "device" in metadata and metadata["device"] is not None:
+        return torch.device(metadata["device"])
+    if fallback is not None:
+        return fallback
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class BaseConditioner(nn.Module):
+    def __init__(self):
+        super(BaseConditioner, self).__init__()
+
+    def _impl_condition(self, y, metadata)->torch.Tensor:
+        raise NotImplementedError()
+
+    def _impl_uncondition(self, y, metadata)->torch.Tensor:
+        raise NotImplementedError()
+
+    @torch.no_grad()
+    def __call__(self, y, metadata:dict={}):
+        condition = self._impl_condition(y, metadata)
+        uncondition = self._impl_uncondition(y, metadata)
+        if condition.dtype in [torch.float64, torch.float32, torch.float16]:
+            condition = condition.to(torch.bfloat16)
+        if uncondition.dtype in [torch.float64,torch.float32, torch.float16]:
+            uncondition = uncondition.to(torch.bfloat16)
+        return condition, uncondition
+
+
+class ComposeConditioner(BaseConditioner):
+    def __init__(self, conditioners:List[BaseConditioner]):
+        super().__init__()
+        self.conditioners = conditioners
+
+    def _impl_condition(self, y, metadata):
+        condition = []
+        for conditioner in self.conditioners:
+            condition.append(conditioner._impl_condition(y, metadata))
+        condition = torch.cat(condition, dim=1)
+        return condition
+
+    def _impl_uncondition(self, y, metadata):
+        uncondition = []
+        for conditioner in self.conditioners:
+            uncondition.append(conditioner._impl_uncondition(y, metadata))
+        uncondition = torch.cat(uncondition, dim=1)
+        return uncondition
+
+
+class LabelConditioner(BaseConditioner):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.null_condition = num_classes
+
+    def _impl_condition(self, y, metadata):
+        device = resolve_conditioner_device(metadata)
+        return torch.tensor(y, device=device).long()
+
+    def _impl_uncondition(self, y, metadata):
+        device = resolve_conditioner_device(metadata)
+        return torch.full((len(y),), self.null_condition, dtype=torch.long, device=device)
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale) + shift
+
+class Embed(nn.Module):
+    def __init__(
+            self,
+            in_chans: int = 3,
+            embed_dim: int = 768,
+            norm_layer = None,
+            bias: bool = True,
+    ):
+        super().__init__()
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.proj = nn.Linear(in_chans, embed_dim, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+    def forward(self, x):
+        x = self.proj(x)
+        x = self.norm(x)
+        return x
+
+class TimestepEmbedder(nn.Module):
+
+    def __init__(self, hidden_size, frequency_embedding_size=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, dim, max_period=10):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+        )
+        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):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes, hidden_size):
+        super().__init__()
+        self.embedding_table = nn.Embedding(num_classes, hidden_size)
+        self.num_classes = num_classes
+
+    def forward(self, labels,):
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.Linear(hidden_size, 2*hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+    ):
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+    def forward(self, x):
+        x =  self.w2(torch.nn.functional.silu(self.w1(x)) * self.w3(x))
+        return x
+
+def precompute_freqs_cis_2d(dim: int, height: int, width:int, theta: float = 10000.0, scale=16.0):
+    # assert  H * H == end
+    # flat_patch_pos = torch.linspace(-1, 1, end) # N = end
+    x_pos = torch.linspace(0, scale, width)
+    y_pos = torch.linspace(0, scale, height)
+    y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij")
+    y_pos = y_pos.reshape(-1)
+    x_pos = x_pos.reshape(-1)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) # Hc/4
+    x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
+    y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
+    x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
+    y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
+    freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) # N,Hc/4,2
+    freqs_cis = freqs_cis.reshape(height*width, -1)
+    return freqs_cis
+
+
+def apply_rotary_emb(
+        xq: torch.Tensor,
+        xk: torch.Tensor,
+        freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    freqs_cis = freqs_cis[None, :, None, :]
+    # xq : B N H Hc
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # B N H Hc/2
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) # B, N, H, Hc
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class RAttention(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = True,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = RMSNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor, pos, mask) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # B N H Hc
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q, k = apply_rotary_emb(q, k, freqs_cis=pos)
+        q = q.view(B, -1, self.num_heads, C // self.num_heads).transpose(1, 2)  # B, H, N, Hc
+        k = k.view(B, -1, self.num_heads, C // self.num_heads).transpose(1, 2).contiguous()  # B, H, N, Hc
+        v = v.view(B, -1, self.num_heads, C // self.num_heads).transpose(1, 2).contiguous()
+
+        x = scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+
+class FlattenDiTBlock(nn.Module):
+    def __init__(self, hidden_size, groups,  mlp_ratio=4.0, ):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
+        self.attn = RAttention(hidden_size, num_heads=groups, qkv_bias=False)
+        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp = FeedForward(hidden_size, mlp_hidden_dim)
+        self.adaLN_modulation = nn.Sequential(
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, x,  c, pos, mask=None):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
+        x = x + gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), pos, mask=mask)
+        x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+class NerfEmbedder(nn.Module):
+    def __init__(self, in_channels, hidden_size_input, max_freqs):
+        super().__init__()
+        self.max_freqs = max_freqs
+        self.hidden_size_input = hidden_size_input
+        self.embedder = nn.Sequential(
+            nn.Linear(in_channels+max_freqs**2, hidden_size_input, bias=True),
+        )
+
+    @lru_cache
+    def fetch_pos(self, patch_size, device, dtype):
+        pos_x = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y, pos_x = torch.meshgrid(pos_y, pos_x, indexing="ij")
+        pos_x = pos_x.reshape(-1, 1, 1)
+        pos_y = pos_y.reshape(-1, 1, 1)
+
+        freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device)
+        freqs_x = freqs[None, :, None]
+        freqs_y = freqs[None, None, :]
+        coeffs = (1 + freqs_x * freqs_y) ** -1
+        dct_x = torch.cos(pos_x * freqs_x * torch.pi)
+        dct_y = torch.cos(pos_y * freqs_y * torch.pi)
+        dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
+        return dct
+
+
+    def forward(self, inputs):
+        B, P2, C = inputs.shape
+        patch_size = int(P2 ** 0.5)
+        device = inputs.device
+        dtype = inputs.dtype
+        dct = self.fetch_pos(patch_size, device, dtype)
+        dct = dct.repeat(B, 1, 1)
+        inputs = torch.cat([inputs, dct], dim=-1)
+        inputs = self.embedder(inputs)
+        return inputs
+
+
+class NerfBlock(nn.Module):
+    def __init__(self, hidden_size_s, hidden_size_x, mlp_ratio=4):
+        super().__init__()
+        self.param_generator1 = nn.Sequential(
+            nn.Linear(hidden_size_s, 2*hidden_size_x**2*mlp_ratio, bias=True),
+        )
+        self.norm = RMSNorm(hidden_size_x, eps=1e-6)
+        self.mlp_ratio = mlp_ratio
+    def forward(self, x, s):
+        batch_size, num_x, hidden_size_x = x.shape
+        mlp_params1 = self.param_generator1(s)
+        fc1_param1, fc2_param1 = mlp_params1.chunk(2, dim=-1)
+        fc1_param1 = fc1_param1.view(batch_size, hidden_size_x, hidden_size_x*self.mlp_ratio)
+        fc2_param1 = fc2_param1.view(batch_size, hidden_size_x*self.mlp_ratio, hidden_size_x)
+
+        # normalize fc1
+        normalized_fc1_param1 = torch.nn.functional.normalize(fc1_param1, dim=-2)
+        # normalize fc2
+        normalized_fc2_param1 = torch.nn.functional.normalize(fc2_param1, dim=-2)
+        # mlp 1
+        res_x = x
+        x = self.norm(x)
+        x = torch.bmm(x, normalized_fc1_param1)
+        x = torch.nn.functional.silu(x)
+        x = torch.bmm(x, normalized_fc2_param1)
+        x = x + res_x
+        return x
+
+class NerfFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.linear(x)
+        return x
+
+class PixNerDiT(nn.Module):
+    def __init__(
+            self,
+            in_channels=4,
+            num_groups=12,
+            hidden_size=1152,
+            hidden_size_x=64,
+            nerf_mlpratio=4,
+            num_blocks=18,
+            num_cond_blocks=4,
+            patch_size=2,
+            num_classes=1000,
+            learn_sigma=True,
+            deep_supervision=0,
+            weight_path=None,
+            load_ema=False,
+    ):
+        super().__init__()
+        self.deep_supervision = deep_supervision
+        self.learn_sigma = learn_sigma
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.hidden_size = hidden_size
+        self.num_groups = num_groups
+        self.num_blocks = num_blocks
+        self.num_cond_blocks = num_cond_blocks
+        self.patch_size = patch_size
+        self.x_embedder = NerfEmbedder(in_channels, hidden_size_x, max_freqs=8)
+        self.s_embedder = Embed(in_channels*patch_size**2, hidden_size, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes+1, hidden_size)
+
+        self.final_layer = NerfFinalLayer(hidden_size_x, self.out_channels)
+
+        self.weight_path = weight_path
+
+        self.load_ema = load_ema
+        self.blocks = nn.ModuleList([
+            FlattenDiTBlock(self.hidden_size, self.num_groups) for _ in range(self.num_cond_blocks)
+        ])
+        self.blocks.extend([
+            NerfBlock(self.hidden_size, hidden_size_x, nerf_mlpratio) for _ in range(self.num_cond_blocks, self.num_blocks)
+        ])
+        self.initialize_weights()
+        self.precompute_pos = dict()
+
+    def fetch_pos(self, height, width, device):
+        if (height, width) in self.precompute_pos:
+            return self.precompute_pos[(height, width)].to(device)
+        else:
+            pos = precompute_freqs_cis_2d(self.hidden_size // self.num_groups, height, width).to(device)
+            self.precompute_pos[(height, width)] = pos
+            return pos
+
+    def initialize_weights(self):
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.s_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.s_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # zero init final layer
+        nn.init.zeros_(self.final_layer.linear.weight)
+        nn.init.zeros_(self.final_layer.linear.bias)
+
+
+    def forward(self, x, t, y, s=None, mask=None):
+        B, _, H, W = x.shape
+        pos = self.fetch_pos(H//self.patch_size, W//self.patch_size, x.device)
+        x = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2)
+        t = self.t_embedder(t.view(-1)).view(B, -1, self.hidden_size)
+        y = self.y_embedder(y).view(B, 1, self.hidden_size)
+        c = nn.functional.silu(t + y)
+        if s is None:
+            s = self.s_embedder(x)
+            for i in range(self.num_cond_blocks):
+                s = self.blocks[i](s, c, pos, mask)
+            s = nn.functional.silu(t + s)
+        batch_size, length, _ = s.shape
+        x = x.reshape(batch_size*length, self.in_channels, self.patch_size**2)
+        x = x.transpose(1, 2)
+        s = s.view(batch_size*length, self.hidden_size)
+        x = self.x_embedder(x)
+        for i in range(self.num_cond_blocks, self.num_blocks):
+            x = self.blocks[i](x, s)
+        x = self.final_layer(x)
+        x = x.transpose(1, 2)
+        x = x.reshape(batch_size, length, -1)
+        x = torch.nn.functional.fold(x.transpose(1, 2).contiguous(), (H, W), kernel_size=self.patch_size, stride=self.patch_size)
+        return x
+
+
+def to_container(config: Any) -> Any:
+    if hasattr(config, "items") and not isinstance(config, dict):
+        return {k: to_container(v) for k, v in config.items()}
+    if isinstance(config, list):
+        return [to_container(v) for v in config]
+    return config
+
+
+def load_symbol(path: str) -> Any:
+    module_path, name = path.rsplit(".", 1)
+    if module_path in {__name__, "modeling_pixnerd_transformer_2d"}:
+        return getattr(sys.modules[__name__], name)
+    module = importlib.import_module(module_path)
+    return getattr(module, name)
+
+
+def instantiate_from_spec(spec: Any) -> Any:
+    spec = to_container(spec)
+    if isinstance(spec, dict) and "class_path" in spec:
+        class_or_fn = load_symbol(spec["class_path"])
+        init_args = spec.get("init_args", {})
+        if isinstance(init_args, dict):
+            init_args = {k: instantiate_from_spec(v) for k, v in init_args.items()}
+        return class_or_fn(**init_args)
+    if isinstance(spec, dict):
+        return {k: instantiate_from_spec(v) for k, v in spec.items()}
+    if isinstance(spec, list):
+        return [instantiate_from_spec(v) for v in spec]
+    if isinstance(spec, str) and "." in spec:
+        try:
+            return load_symbol(spec)
+        except Exception:
+            return spec
+    return spec
+
+
+def clone_spec(spec: Dict[str, Any]) -> Dict[str, Any]:
+    return copy.deepcopy(to_container(spec))
+
+
+def load_prefixed_state_dict(
+    module: Optional[torch.nn.Module],
+    state_dict: Dict[str, torch.Tensor],
+    prefixes: Iterable[str],
+) -> bool:
+    if module is None:
+        return False
+    for prefix in prefixes:
+        subset = {
+            key[len(prefix) :]: value
+            for key, value in state_dict.items()
+            if key.startswith(prefix)
+        }
+        if subset:
+            module.load_state_dict(subset, strict=False)
+            return True
+    return False
+
+
+@dataclass
+class PixNerdTransformer2DModelOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+class PixNerdTransformer2DModel(ModelMixin, ConfigMixin):
+    config_name = "config.json"
+
+    @register_to_config
+    def __init__(
+        self,
+        denoiser_spec: Dict[str, Any],
+        conditioner_spec: Dict[str, Any],
+        vae_spec: Optional[Dict[str, Any]] = None,
+        diffusion_trainer_spec: Optional[Dict[str, Any]] = None,
+        use_ema: bool = True,
+        ema_decay: float = 0.9999,
+        compile_denoiser: bool = False,
+    ) -> None:
+        super().__init__()
+        self.denoiser = instantiate_from_spec(to_container(denoiser_spec))
+        self.conditioner = instantiate_from_spec(to_container(conditioner_spec))
+        self.vae = instantiate_from_spec(to_container(vae_spec)) if vae_spec is not None else None
+        self.diffusion_trainer = (
+            instantiate_from_spec(to_container(diffusion_trainer_spec))
+            if diffusion_trainer_spec is not None
+            else None
+        )
+
+        self.use_ema = bool(use_ema)
+        self.ema_decay = float(ema_decay)
+        self.ema_denoiser = copy.deepcopy(self.denoiser) if self.use_ema else None
+        if self.ema_denoiser is not None:
+            self.ema_denoiser.to(torch.float32)
+
+        if compile_denoiser and hasattr(self.denoiser, "compile"):
+            self.denoiser.compile()
+            if self.ema_denoiser is not None:
+                self.ema_denoiser.compile()
+
+        self._freeze_non_trainable_modules()
+        if self.ema_denoiser is not None:
+            self.sync_ema()
+
+    @property
+    def patch_size(self) -> int:
+        return int(getattr(self.denoiser, "patch_size", 1))
+
+    @property
+    def in_channels(self) -> int:
+        return int(getattr(self.denoiser, "in_channels", 3))
+
+    @classmethod
+    def from_project_config(
+        cls,
+        model_config: Dict[str, Any],
+        use_ema: bool = True,
+        compile_denoiser: bool = False,
+    ) -> "PixNerdTransformer2DModel":
+        model_config = to_container(model_config)
+        ema_decay = model_config.get("ema_tracker", {}).get("init_args", {}).get("decay", 0.9999)
+        return cls(
+            denoiser_spec=model_config["denoiser"],
+            conditioner_spec=model_config["conditioner"],
+            vae_spec=model_config.get("vae"),
+            diffusion_trainer_spec=model_config.get("diffusion_trainer"),
+            use_ema=use_ema,
+            ema_decay=ema_decay,
+            compile_denoiser=compile_denoiser,
+        )
+
+    @staticmethod
+    def _as_timestep_tensor(
+        timestep: Any,
+        batch_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if isinstance(timestep, torch.Tensor):
+            if timestep.ndim == 0:
+                return timestep.repeat(batch_size).to(device=device, dtype=torch.float32)
+            return timestep.to(device=device, dtype=torch.float32)
+        return torch.full((batch_size,), float(timestep), device=device, dtype=torch.float32)
+
+    def _freeze_module(self, module: Optional[torch.nn.Module]) -> None:
+        if module is None:
+            return
+        module.eval()
+        for parameter in module.parameters():
+            parameter.requires_grad = False
+
+    def _freeze_non_trainable_modules(self) -> None:
+        self._freeze_module(self.conditioner)
+        self._freeze_module(self.vae)
+        self._freeze_module(self.ema_denoiser)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Any,
+        encoder_hidden_states: torch.Tensor,
+        return_dict: bool = True,
+    ) -> PixNerdTransformer2DModelOutput | Tuple[torch.Tensor]:
+        t = self._as_timestep_tensor(timestep, sample.shape[0], sample.device)
+        out = self.denoiser(sample, t, encoder_hidden_states)
+        if not return_dict:
+            return (out,)
+        return PixNerdTransformer2DModelOutput(sample=out)
+
+    def predict_noise(
+        self,
+        sample: torch.Tensor,
+        timestep: Any,
+        encoder_hidden_states: torch.Tensor,
+        use_ema: bool = False,
+    ) -> torch.Tensor:
+        t = self._as_timestep_tensor(timestep, sample.shape[0], sample.device)
+        denoiser = self.get_inference_denoiser(use_ema=use_ema)
+        return denoiser(sample, t, encoder_hidden_states)
+
+    def get_inference_denoiser(self, use_ema: bool = True) -> torch.nn.Module:
+        if use_ema and self.ema_denoiser is not None:
+            return self.ema_denoiser
+        return self.denoiser
+
+    @torch.no_grad()
+    def get_conditioning(
+        self,
+        y: Iterable[Any],
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        metadata = {} if metadata is None else metadata
+        return self.conditioner(y, metadata)
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        if self.vae is None:
+            return x
+        return self.vae.encode(x)
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        if self.vae is None:
+            return latents
+        return self.vae.decode(latents)
+
+    @torch.no_grad()
+    def sync_ema(self) -> None:
+        if self.ema_denoiser is None:
+            return
+        self.ema_denoiser.load_state_dict(self.denoiser.state_dict(), strict=True)
+        self.ema_denoiser.to(torch.float32)
+
+    @torch.no_grad()
+    def ema_step(self, decay: Optional[float] = None) -> None:
+        if self.ema_denoiser is None:
+            return
+        decay = self.ema_decay if decay is None else float(decay)
+        for ema_param, param in zip(self.ema_denoiser.parameters(), self.denoiser.parameters()):
+            ema_param.mul_(decay).add_(param.detach().float(), alpha=1.0 - decay)
+
+    def compute_training_loss(
+        self,
+        x: torch.Tensor,
+        y: Iterable[Any],
+        scheduler: torch.nn.Module,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, torch.Tensor]:
+        if self.diffusion_trainer is None:
+            raise RuntimeError("diffusion_trainer is not configured.")
+        metadata = {} if metadata is None else metadata
+
+        with torch.no_grad():
+            x = self.encode(x)
+            condition, uncondition = self.get_conditioning(y, metadata)
+
+        return self.diffusion_trainer(
+            self.denoiser,
+            self.ema_denoiser if self.ema_denoiser is not None else self.denoiser,
+            scheduler,
+            x,
+            condition,
+            uncondition,
+            metadata,
+        )
+
+__all__ = [
+    "PixNerDiT",
+    "LabelConditioner",
+    "PixelAE",
+    "PixNerdTransformer2DModel",
+    "PixNerdTransformer2DModelOutput",
+]

README.md

@@ -11,71 +11,129 @@ language:
   - en
 ---
 
-# PixNerd-XL-16 Diffusers Checkpoints
+# BiliSakura/PixNerd-diffusers
 
-Production-ready Diffusers export of PixNerd-XL/16 class-conditional ImageNet checkpoints.
+Self-contained PixNerd-XL/16 checkpoints for Hugging Face diffusers. **No external code repo is required** — each subfolder ships its own `pipeline.py`, component modules, and weights.
 
-## Available Checkpoints
+This repo is derived from the development bundle in [Visual-Generative-Foundation-Model-Collection](https://github.com/Bili-Sakura/Visual-Generative-Foundation-Model-Collection), but inference only needs:
 
-- `PixNerd-XL-16-256`
-  - source: `epoch%3D319-step%3D1600000_emainit.ckpt`
-  - target resolution: `256x256`
-- `PixNerd-XL-16-512`
-  - source: `res512_ft200k_epoch%3D325-step%3D1800000_emainit.ckpt`
-  - target resolution: `512x512`
+- This model repo (`BiliSakura/PixNerd-diffusers`)
+- PyPI `diffusers`, `torch`, `huggingface_hub`
 
-Both checkpoints are packaged with:
+This Hugging Face repo hosts **multiple self-contained checkpoints as subfolders**. Each subfolder includes its own `pipeline.py`, `model_index.json`, weights, and component code (`transformer/`, `scheduler/`).
 
-- `pipeline.py`
-- `modeling_pixnerd_transformer_2d.py`
-- `scheduling_pixnerd_flow_match.py`
-- `transformer/` weights + config
-- `scheduler/` config
+## Available checkpoints
 
-## Requirements
+| Subfolder | Resolution | Source checkpoint |
+| --- | --- | --- |
+| [`PixNerd-XL-16-256/`](PixNerd-XL-16-256/) | 256×256 | `epoch%3D319-step%3D1600000_emainit.ckpt` |
+| [`PixNerd-XL-16-512/`](PixNerd-XL-16-512/) | 512×512 | `res512_ft200k_epoch%3D325-step%3D1800000_emainit.ckpt` |
 
-```bash
-pip install torch diffusers
-```
+Both checkpoints are ImageNet class-conditional PixNerd-XL/16 exports with flow-matching sampling.
+
+## ImageNet class labels
+
+ImageNet-1k labels live in shared [`labels/`](labels/) at the repo root (not duplicated per variant). Format follows Hugging Face / DiT convention:
+
+| File | Direction | Value format |
+| --- | --- | --- |
+| `labels/id2label_en.json` | id → English | comma-separated synonyms, e.g. `"207": "golden retriever"` |
+| `labels/id2label_cn.json` | id → Chinese | comma-separated synonyms, e.g. `"207": "金毛猎犬"` |
+
+After `PixNerdPipeline.from_pretrained(...)`, the pipeline exposes:
 
-## Inference (Python)
+- `pipe.id2label` / `pipe.id2label_cn` — inspect id → label correspondence
+- `pipe.labels` / `pipe.labels_cn` — reverse maps (synonym → id), sorted for browsing
+- `pipe.get_label_ids("golden retriever")` or `pipe.get_label_ids("金毛猎犬", lang="cn")`
+- `pipe(prompt="golden retriever", ...)` — string labels resolved automatically
+
+Why JSON at repo root instead of a Python dict in each variant?
+
+1. **Explicit correspondence** — users can open `id2label_en.json` and see every id without running code.
+2. **Hub-compatible** — same shape as `facebook/DiT-XL-2-256` and other vision checkpoints.
+3. **Shared across variants** — both PixNerd checkpoints use the same 1000 ImageNet classes.
+4. **Bilingual without duplication** — English and Chinese are separate files; the pipeline loads both.
+
+## Load from Hugging Face
 
 ```python
 import torch
 from diffusers import DiffusionPipeline
 
-model_dir = "PixNerd-XL-16-256"  # or PixNerd-XL-16-512
+variant = "PixNerd-XL-16-256"  # or PixNerd-XL-16-512
+resolution = 256 if variant.endswith("256") else 512
+
 pipe = DiffusionPipeline.from_pretrained(
-    model_dir,
-    custom_pipeline=f"{model_dir}/pipeline.py",
+    f"BiliSakura/PixNerd-diffusers/{variant}",
+    trust_remote_code=True,
     torch_dtype=torch.float32,
-).to("cpu")  # use "cuda" if available
+).to("cuda")
 
-# Class-conditional generation: class label 207 (golden retriever)
 images = pipe(
-    prompt=[207],
-    num_images_per_prompt=1,
-    height=256,
-    width=256,
+    prompt=207,
+    height=resolution,
+    width=resolution,
     num_inference_steps=25,
     guidance_scale=4.0,
     timeshift=3.0,
     order=2,
 ).images
 
-images[0].save("sample.png")
+print(pipe.id2label[207])          # "golden retriever"
+print(pipe.id2label_cn[207])       # "金毛猎犬"
+pipe.get_label_ids("golden retriever")  # [207]
+images = pipe(prompt="golden retriever", height=resolution, width=resolution).images
 ```
 
-## Interface Notes
+## Load from a local clone
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+repo = "models/BiliSakura/PixNerd-diffusers"
+variant = "PixNerd-XL-16-256"
+
+pipe = DiffusionPipeline.from_pretrained(
+    f"{repo}/{variant}",
+    trust_remote_code=True,
+    torch_dtype=torch.float32,
+).to("cuda")
 
-- The pipeline uses `prompt` for conditioning input.
-- For class-conditional generation, pass integer labels, e.g. `prompt=[207]`.
-- `height` and `width` should match checkpoint intent (256 or 512), but custom sizes work if divisible by patch size.
+images = pipe(prompt=207, height=256, width=256).images
+```
+
+## Repo layout
+
+```text
+BiliSakura/PixNerd-diffusers/
+├── README.md
+├── labels/
+│   ├── id2label_en.json      # ImageNet id -> English synonyms
+│   ├── id2label_cn.json      # ImageNet id -> Chinese synonyms
+│   └── imagenet_labels.py    # loader helpers
+├── PixNerd-XL-16-256/
+│   ├── README.md
+│   ├── pipeline.py
+│   ├── model_index.json
+│   ├── conversion_metadata.json
+│   ├── transformer/
+│   └── scheduler/
+└── PixNerd-XL-16-512/
+    ├── README.md
+    ├── pipeline.py
+    ├── model_index.json
+    ├── conversion_metadata.json
+    ├── transformer/
+    └── scheduler/
+```
 
-## Reproducibility Metadata
+## Interface notes
 
+- The pipeline uses `prompt` for class conditioning input.
+- Pass integer ImageNet ids (`prompt=207`) or human-readable synonyms (`prompt="golden retriever"`).
+- `height` and `width` should match checkpoint intent (256 or 512), but custom sizes work if divisible by patch size (16).
 - Architecture and conversion provenance are recorded in each checkpoint's `conversion_metadata.json`.
-- Transformer and scheduler runtime classes are defined in repository-local Python modules shipped with each checkpoint.
 
 ## Limitations
 

labels/id2label_cn.json

@@ -0,0 +1,1002 @@
+{
+  "0": "丁鲷",
+  "1": "金鱼",
+  "2": "大白鲨",
+  "3": "虎鲨",
+  "4": "锤头鲨",
+  "5": "电鳐",
+  "6": "黄貂鱼",
+  "7": "公鸡",
+  "8": "母鸡",
+  "9": "鸵鸟",
+  "10": "燕雀",
+  "11": "金翅雀",
+  "12": "家朱雀",
+  "13": "灯芯草雀",
+  "14": "靛蓝雀,靛蓝鸟",
+  "15": "蓝鹀",
+  "16": "夜莺",
+  "17": "松鸦",
+  "18": "喜鹊",
+  "19": "山雀",
+  "20": "河鸟",
+  "21": "鸢（猛禽）",
+  "22": "秃头鹰",
+  "23": "秃鹫",
+  "24": "大灰猫头鹰",
+  "25": "欧洲火蝾螈",
+  "26": "普通蝾螈",
+  "27": "水蜥",
+  "28": "斑点蝾螈",
+  "29": "蝾螈,泥狗",
+  "30": "牛蛙",
+  "31": "树蛙",
+  "32": "尾蛙,铃蟾蜍,肋蟾蜍,尾蟾蜍",
+  "33": "红海龟",
+  "34": "皮革龟",
+  "35": "泥龟",
+  "36": "淡水龟",
+  "37": "箱龟",
+  "38": "带状壁虎",
+  "39": "普通鬣蜥",
+  "40": "美国变色龙",
+  "41": "鞭尾蜥蜴",
+  "42": "飞龙科蜥蜴",
+  "43": "褶边蜥蜴",
+  "44": "鳄鱼蜥蜴",
+  "45": "毒蜥",
+  "46": "绿蜥蜴",
+  "47": "非洲变色龙",
+  "48": "科莫多蜥蜴",
+  "49": "非洲鳄,尼罗河鳄鱼",
+  "50": "美国鳄鱼,鳄鱼",
+  "51": "三角龙",
+  "52": "雷蛇,蠕虫蛇",
+  "53": "环蛇,环颈蛇",
+  "54": "希腊蛇",
+  "55": "绿蛇,草蛇",
+  "56": "国王蛇",
+  "57": "袜带蛇,草蛇",
+  "58": "水蛇",
+  "59": "藤蛇",
+  "60": "夜蛇",
+  "61": "大蟒蛇",
+  "62": "岩石蟒蛇,岩蛇,蟒蛇",
+  "63": "印度眼镜蛇",
+  "64": "绿曼巴",
+  "65": "海蛇",
+  "66": "角腹蛇",
+  "67": "菱纹响尾蛇",
+  "68": "角响尾蛇",
+  "69": "三叶虫",
+  "70": "盲蜘蛛",
+  "71": "蝎子",
+  "72": "黑金花园蜘蛛",
+  "73": "谷仓蜘蛛",
+  "74": "花园蜘蛛",
+  "75": "黑寡妇蜘蛛",
+  "76": "狼蛛",
+  "77": "狼蜘蛛,狩猎蜘蛛",
+  "78": "壁虱",
+  "79": "蜈蚣",
+  "80": "黑松鸡",
+  "81": "松鸡,雷鸟",
+  "82": "披肩鸡,披肩榛鸡",
+  "83": "草原鸡,草原松鸡",
+  "84": "孔雀",
+  "85": "鹌鹑",
+  "86": "鹧鸪",
+  "87": "非洲灰鹦鹉",
+  "88": "金刚鹦鹉",
+  "89": "硫冠鹦鹉",
+  "90": "短尾鹦鹉",
+  "91": "褐翅鸦鹃",
+  "92": "蜜蜂",
+  "93": "犀鸟",
+  "94": "蜂鸟",
+  "95": "鹟䴕",
+  "96": "犀鸟",
+  "97": "野鸭",
+  "98": "红胸秋沙鸭",
+  "99": "鹅",
+  "100": "黑天鹅",
+  "101": "大象",
+  "102": "针鼹鼠",
+  "103": "鸭嘴兽",
+  "104": "沙袋鼠",
+  "105": "考拉,考拉熊",
+  "106": "袋熊",
+  "107": "水母",
+  "108": "海葵",
+  "109": "脑珊瑚",
+  "110": "扁形虫扁虫",
+  "111": "线虫,蛔虫",
+  "112": "海螺",
+  "113": "蜗牛",
+  "114": "鼻涕虫",
+  "115": "海参",
+  "116": "石鳖",
+  "117": "鹦鹉螺",
+  "118": "珍宝蟹",
+  "119": "石蟹",
+  "120": "招潮蟹",
+  "121": "帝王蟹,阿拉斯加蟹,阿拉斯加帝王蟹",
+  "122": "美国龙虾,缅因州龙虾",
+  "123": "大螯虾",
+  "124": "小龙虾",
+  "125": "寄居蟹",
+  "126": "等足目动物(明虾和螃蟹近亲)",
+  "127": "白鹳",
+  "128": "黑鹳",
+  "129": "鹭",
+  "130": "火烈鸟",
+  "131": "小蓝鹭",
+  "132": "美国鹭,大白鹭",
+  "133": "麻鸦",
+  "134": "鹤",
+  "135": "秧鹤",
+  "136": "欧洲水鸡,紫水鸡",
+  "137": "沼泽泥母鸡,水母鸡",
+  "138": "鸨",
+  "139": "红翻石鹬",
+  "140": "红背鹬,黑腹滨鹬",
+  "141": "红脚鹬",
+  "142": "半蹼鹬",
+  "143": "蛎鹬",
+  "144": "鹈鹕",
+  "145": "国王企鹅",
+  "146": "信天翁,大海鸟",
+  "147": "灰鲸",
+  "148": "杀人鲸,逆戟鲸,虎鲸",
+  "149": "海牛",
+  "150": "海狮",
+  "151": "奇瓦瓦",
+  "152": "日本猎犬",
+  "153": "马尔济斯犬",
+  "154": "狮子狗",
+  "155": "西施犬",
+  "156": "布莱尼姆猎犬",
+  "157": "巴比狗",
+  "158": "玩具犬",
+  "159": "罗得西亚长背猎狗",
+  "160": "阿富汗猎犬",
+  "161": "猎犬",
+  "162": "比格犬,猎兔犬",
+  "163": "侦探犬",
+  "164": "蓝色快狗",
+  "165": "黑褐猎浣熊犬",
+  "166": "沃克猎犬",
+  "167": "英国猎狐犬",
+  "168": "美洲赤狗",
+  "169": "俄罗斯猎狼犬",
+  "170": "爱尔兰猎狼犬",
+  "171": "意大利灰狗",
+  "172": "惠比特犬",
+  "173": "依比沙猎犬",
+  "174": "挪威猎犬",
+  "175": "奥达猎犬,水獭猎犬",
+  "176": "沙克犬,瞪羚猎犬",
+  "177": "苏格兰猎鹿犬,猎鹿犬",
+  "178": "威玛猎犬",
+  "179": "斯塔福德郡牛头梗,斯塔福德郡斗牛梗",
+  "180": "美国斯塔福德郡梗,美国比特斗牛梗,斗牛梗",
+  "181": "贝德灵顿梗",
+  "182": "边境梗",
+  "183": "凯丽蓝梗",
+  "184": "爱尔兰梗",
+  "185": "诺福克梗",
+  "186": "诺维奇梗",
+  "187": "约克郡梗",
+  "188": "刚毛猎狐梗",
+  "189": "莱克兰梗",
+  "190": "锡利哈姆梗",
+  "191": "艾尔谷犬",
+  "192": "凯恩梗",
+  "193": "澳大利亚梗",
+  "194": "丹迪丁蒙梗",
+  "195": "波士顿梗",
+  "196": "迷你雪纳瑞犬",
+  "197": "巨型雪纳瑞犬",
+  "198": "标准雪纳瑞犬",
+  "199": "苏格兰梗",
+  "200": "西藏梗,菊花狗",
+  "201": "丝毛梗",
+  "202": "软毛麦色梗",
+  "203": "西高地白梗",
+  "204": "拉萨阿普索犬",
+  "205": "平毛寻回犬",
+  "206": "卷毛寻回犬",
+  "207": "金毛猎犬",
+  "208": "拉布拉多猎犬",
+  "209": "乞沙比克猎犬",
+  "210": "德国短毛猎犬",
+  "211": "维兹拉犬",
+  "212": "英国谍犬",
+  "213": "爱尔兰雪达犬,红色猎犬",
+  "214": "戈登雪达犬",
+  "215": "布列塔尼犬猎犬",
+  "216": "黄毛,黄毛猎犬",
+  "217": "英国史宾格犬",
+  "218": "威尔士史宾格犬",
+  "219": "可卡犬,英国可卡犬",
+  "220": "萨塞克斯猎犬",
+  "221": "爱尔兰水猎犬",
+  "222": "哥威斯犬",
+  "223": "舒柏奇犬",
+  "224": "比利时牧羊犬",
+  "225": "马里努阿犬",
+  "226": "伯瑞犬",
+  "227": "凯尔皮犬",
+  "228": "匈牙利牧羊犬",
+  "229": "老英国牧羊犬",
+  "230": "喜乐蒂牧羊犬",
+  "231": "牧羊犬",
+  "232": "边境牧羊犬",
+  "233": "法兰德斯牧牛狗",
+  "234": "罗特韦尔犬",
+  "235": "德国牧羊犬,德国警犬,阿尔萨斯",
+  "236": "多伯曼犬,杜宾犬",
+  "237": "迷你杜宾犬",
+  "238": "大瑞士山地犬",
+  "239": "伯恩山犬",
+  "240": "Appenzeller狗",
+  "241": "EntleBucher狗",
+  "242": "拳师狗",
+  "243": "斗牛獒",
+  "244": "藏獒",
+  "245": "法国斗牛犬",
+  "246": "大丹犬",
+  "247": "圣伯纳德狗",
+  "248": "爱斯基摩犬,哈士奇",
+  "249": "雪橇犬,阿拉斯加爱斯基摩狗",
+  "250": "哈士奇",
+  "251": "达尔马提亚,教练车狗",
+  "252": "狮毛狗",
+  "253": "巴辛吉狗",
+  "254": "哈巴狗,狮子狗",
+  "255": "莱昂贝格狗",
+  "256": "纽芬兰岛狗",
+  "257": "大白熊犬",
+  "258": "萨摩耶犬",
+  "259": "博美犬",
+  "260": "松狮,松狮",
+  "261": "荷兰卷尾狮毛狗",
+  "262": "布鲁塞尔格林芬犬",
+  "263": "彭布洛克威尔士科基犬",
+  "264": "威尔士柯基犬",
+  "265": "玩具贵宾犬",
+  "266": "迷你贵宾犬",
+  "267": "标准贵宾犬",
+  "268": "墨西哥无毛犬",
+  "269": "灰狼",
+  "270": "白狼,北极狼",
+  "271": "红太狼,鬃狼,犬犬鲁弗斯",
+  "272": "狼,草原狼,刷狼,郊狼",
+  "273": "澳洲野狗,澳大利亚野犬",
+  "274": "豺",
+  "275": "非洲猎犬,土狼犬",
+  "276": "鬣狗",
+  "277": "红狐狸",
+  "278": "沙狐",
+  "279": "北极狐狸,白狐狸",
+  "280": "灰狐狸",
+  "281": "虎斑猫",
+  "282": "山猫,虎猫",
+  "283": "波斯猫",
+  "284": "暹罗暹罗猫,",
+  "285": "埃及猫",
+  "286": "美洲狮,美洲豹",
+  "287": "猞猁,山猫",
+  "288": "豹子",
+  "289": "雪豹",
+  "290": "美洲虎",
+  "291": "狮子",
+  "292": "老虎",
+  "293": "猎豹",
+  "294": "棕熊",
+  "295": "美洲黑熊",
+  "296": "冰熊,北极熊",
+  "297": "懒熊",
+  "298": "猫鼬",
+  "299": "猫鼬,海猫",
+  "300": "虎甲虫",
+  "301": "瓢虫",
+  "302": "土鳖虫",
+  "303": "天牛",
+  "304": "龟甲虫",
+  "305": "粪甲虫",
+  "306": "犀牛甲虫",
+  "307": "象甲",
+  "308": "苍蝇",
+  "309": "蜜蜂",
+  "310": "蚂蚁",
+  "311": "蚱蜢",
+  "312": "蟋蟀",
+  "313": "竹节虫",
+  "314": "蟑螂",
+  "315": "螳螂",
+  "316": "蝉",
+  "317": "叶蝉",
+  "318": "草蜻蛉",
+  "319": "蜻蜓",
+  "320": "豆娘,蜻蛉",
+  "321": "优红蛱蝶",
+  "322": "小环蝴蝶",
+  "323": "君主蝴蝶,大斑蝶",
+  "324": "菜粉蝶",
+  "325": "白蝴蝶",
+  "326": "灰蝶",
+  "327": "海星",
+  "328": "海胆",
+  "329": "海参,海黄瓜",
+  "330": "野兔",
+  "331": "兔",
+  "332": "安哥拉兔",
+  "333": "仓鼠",
+  "334": "刺猬,豪猪,",
+  "335": "黑松鼠",
+  "336": "土拨鼠",
+  "337": "海狸",
+  "338": "豚鼠,豚鼠",
+  "339": "栗色马",
+  "340": "斑马",
+  "341": "猪",
+  "342": "野猪",
+  "343": "疣猪",
+  "344": "河马",
+  "345": "牛",
+  "346": "水牛,亚洲水牛",
+  "347": "野牛",
+  "348": "公羊",
+  "349": "大角羊,洛矶山大角羊",
+  "350": "山羊",
+  "351": "狷羚",
+  "352": "黑斑羚",
+  "353": "瞪羚",
+  "354": "阿拉伯单峰骆驼,骆驼",
+  "355": "羊驼",
+  "356": "黄鼠狼",
+  "357": "水貂",
+  "358": "臭猫",
+  "359": "黑足鼬",
+  "360": "水獭",
+  "361": "臭鼬,木猫",
+  "362": "獾",
+  "363": "犰狳",
+  "364": "树懒",
+  "365": "猩猩,婆罗洲猩猩",
+  "366": "大猩猩",
+  "367": "黑猩猩",
+  "368": "长臂猿",
+  "369": "合趾猿长臂猿,合趾猿",
+  "370": "长尾猴",
+  "371": "赤猴",
+  "372": "狒狒",
+  "373": "恒河猴,猕猴",
+  "374": "白头叶猴",
+  "375": "疣猴",
+  "376": "长鼻猴",
+  "377": "狨（美洲产小型长尾猴）",
+  "378": "卷尾猴",
+  "379": "吼猴",
+  "380": "伶猴",
+  "381": "蜘蛛猴",
+  "382": "松鼠猴",
+  "383": "马达加斯加环尾狐猴,鼠狐猴",
+  "384": "大狐猴,马达加斯加大狐猴",
+  "385": "印度大象,亚洲象",
+  "386": "非洲象,非洲象",
+  "387": "小熊猫",
+  "388": "大熊猫",
+  "389": "杖鱼",
+  "390": "鳗鱼",
+  "391": "银鲑,银鲑���",
+  "392": "三色刺蝶鱼",
+  "393": "海葵鱼",
+  "394": "鲟鱼",
+  "395": "雀鳝",
+  "396": "狮子鱼",
+  "397": "河豚",
+  "398": "算盘",
+  "399": "长袍",
+  "400": "学位袍",
+  "401": "手风琴",
+  "402": "原声吉他",
+  "403": "航空母舰",
+  "404": "客机",
+  "405": "飞艇",
+  "406": "祭坛",
+  "407": "救护车",
+  "408": "水陆两用车",
+  "409": "模拟时钟",
+  "410": "蜂房",
+  "411": "围裙",
+  "412": "垃圾桶",
+  "413": "攻击步枪,枪",
+  "414": "背包",
+  "415": "面包店,面包铺,",
+  "416": "平衡木",
+  "417": "热气球",
+  "418": "圆珠笔",
+  "419": "创可贴",
+  "420": "班卓琴",
+  "421": "栏杆,楼梯扶手",
+  "422": "杠铃",
+  "423": "理发师的椅子",
+  "424": "理发店",
+  "425": "牲口棚",
+  "426": "晴雨表",
+  "427": "圆筒",
+  "428": "园地小车,手推车",
+  "429": "棒球",
+  "430": "篮球",
+  "431": "婴儿床",
+  "432": "巴松管,低音管",
+  "433": "游泳帽",
+  "434": "沐浴毛巾",
+  "435": "浴缸,澡盆",
+  "436": "沙滩车,旅行车",
+  "437": "灯塔",
+  "438": "高脚杯",
+  "439": "熊皮高帽",
+  "440": "啤酒瓶",
+  "441": "啤酒杯",
+  "442": "钟塔",
+  "443": "（小儿用的）围嘴",
+  "444": "串联自行车,",
+  "445": "比基尼",
+  "446": "装订册",
+  "447": "双筒望远镜",
+  "448": "鸟舍",
+  "449": "船库",
+  "450": "雪橇",
+  "451": "饰扣式领带",
+  "452": "阔边女帽",
+  "453": "书橱",
+  "454": "书店,书摊",
+  "455": "瓶盖",
+  "456": "弓箭",
+  "457": "蝴蝶结领结",
+  "458": "铜制牌位",
+  "459": "奶罩",
+  "460": "防波堤,海堤",
+  "461": "铠甲",
+  "462": "扫帚",
+  "463": "桶",
+  "464": "扣环",
+  "465": "防弹背心",
+  "466": "动车,子弹头列车",
+  "467": "肉铺,肉菜市场",
+  "468": "出租车",
+  "469": "大锅",
+  "470": "蜡烛",
+  "471": "大炮",
+  "472": "独木舟",
+  "473": "开瓶器,开罐器",
+  "474": "开衫",
+  "475": "车镜",
+  "476": "旋转木马",
+  "477": "木匠的工具包,工具包",
+  "478": "纸箱",
+  "479": "车轮",
+  "480": "取款机,自动取款机",
+  "481": "盒式录音带",
+  "482": "卡带播放器",
+  "483": "城堡",
+  "484": "双体船",
+  "485": "CD播放器",
+  "486": "大提琴",
+  "487": "移动电话,手机",
+  "488": "铁链",
+  "489": "围栏",
+  "490": "链甲",
+  "491": "电锯,油锯",
+  "492": "箱子",
+  "493": "衣柜,洗脸台",
+  "494": "编钟,钟,锣",
+  "495": "中国橱柜",
+  "496": "圣诞袜",
+  "497": "教堂,教堂建筑",
+  "498": "电影院,剧场",
+  "499": "切肉刀,菜刀",
+  "500": "悬崖屋",
+  "501": "斗篷",
+  "502": "木屐,木鞋",
+  "503": "鸡尾酒调酒器",
+  "504": "咖啡杯",
+  "505": "咖啡壶",
+  "506": "螺旋结构（楼梯）",
+  "507": "组合锁",
+  "508": "电脑键盘,键盘",
+  "509": "糖果,糖果店",
+  "510": "集装箱船",
+  "511": "敞篷车",
+  "512": "开瓶器,瓶螺杆",
+  "513": "短号,喇叭",
+  "514": "牛仔靴",
+  "515": "牛仔帽",
+  "516": "摇篮",
+  "517": "起重机",
+  "518": "头盔",
+  "519": "板条箱",
+  "520": "小儿床",
+  "521": "砂锅",
+  "522": "槌球",
+  "523": "拐杖",
+  "524": "胸甲",
+  "525": "大坝,堤防",
+  "526": "书桌",
+  "527": "台式电脑",
+  "528": "有线电话",
+  "529": "尿布湿",
+  "530": "数字时钟",
+  "531": "数字手表",
+  "532": "餐桌板",
+  "533": "抹布",
+  "534": "洗碗机,洗碟机",
+  "535": "盘式制动器",
+  "536": "码头,船坞,码头设施",
+  "537": "狗拉雪橇",
+  "538": "圆顶",
+  "539": "门垫,垫子",
+  "540": "钻井平台,海上钻井",
+  "541": "鼓,乐器,鼓膜",
+  "542": "鼓槌",
+  "543": "哑铃",
+  "544": "荷兰烤箱",
+  "545": "电风扇,鼓风机",
+  "546": "电吉他",
+  "547": "电力机车",
+  "548": "电视,电视柜",
+  "549": "信封",
+  "550": "浓缩咖啡机",
+  "551": "扑面粉",
+  "552": "女用长围巾",
+  "553": "文件,文件柜,档案柜",
+  "554": "消防船",
+  "555": "消防车",
+  "556": "火炉栏",
+  "557": "旗杆",
+  "558": "长笛",
+  "559": "折叠椅",
+  "560": "橄榄球头盔",
+  "561": "叉车",
+  "562": "喷泉",
+  "563": "钢笔",
+  "564": "有四根帷柱的床",
+  "565": "运货车厢",
+  "566": "圆号,喇叭",
+  "567": "煎锅",
+  "568": "裘皮大衣",
+  "569": "垃圾车",
+  "570": "防毒面具,呼吸器",
+  "571": "汽油泵",
+  "572": "高脚杯",
+  "573": "卡丁车",
+  "574": "高尔夫球",
+  "575": "高尔夫球车",
+  "576": "狭长小船",
+  "577": "锣",
+  "578": "礼服",
+  "579": "钢琴",
+  "580": "温室,苗圃",
+  "581": "散热器格栅",
+  "582": "杂货店,食品市场",
+  "583": "断头台",
+  "584": "小发夹",
+  "585": "头发喷雾",
+  "586": "半履带装甲车",
+  "587": "锤子",
+  "588": "大篮子",
+  "589": "手摇鼓风机,吹风机",
+  "590": "手提电脑",
+  "591": "手帕",
+  "592": "硬盘",
+  "593": "口琴,口风琴",
+  "594": "竖琴",
+  "595": "收割机",
+  "596": "斧头",
+  "597": "手枪皮套",
+  "598": "家庭影院",
+  "599": "蜂窝",
+  "600": "钩爪",
+  "601": "衬裙",
+  "602": "单杠",
+  "603": "马车",
+  "604": "沙漏",
+  "605": "手机，iPad",
+  "606": "熨斗",
+  "607": "南瓜灯笼",
+  "608": "牛仔裤,蓝色牛仔裤",
+  "609": "吉普车",
+  "610": "运动衫,T恤",
+  "611": "拼图",
+  "612": "人力车",
+  "613": "操纵杆",
+  "614": "和服",
+  "615": "护膝",
+  "616": "蝴蝶结",
+  "617": "大褂,实验室外套",
+  "618": "长柄勺",
+  "619": "灯罩",
+  "620": "笔记本电脑",
+  "621": "割草机",
+  "622": "镜头盖",
+  "623": "开信刀,裁纸刀",
+  "624": "图书馆",
+  "625": "救生艇",
+  "626": "点火器,打火机",
+  "627": "豪华轿车",
+  "628": "远洋班轮",
+  "629": "唇膏,口红",
+  "630": "平底便鞋",
+  "631": "洗剂",
+  "632": "扬声器",
+  "633": "放大镜",
+  "634": "锯木厂",
+  "635": "磁罗盘",
+  "636": "邮袋",
+  "637": "信箱",
+  "638": "女游泳衣",
+  "639": "有肩带浴衣",
+  "640": "窨井盖",
+  "641": "沙球（一种打击乐器）",
+  "642": "马林巴木琴",
+  "643": "面膜",
+  "644": "火柴",
+  "645": "花柱",
+  "646": "迷宫",
+  "647": "量杯",
+  "648": "药箱",
+  "649": "巨石,巨石结构",
+  "650": "麦克风",
+  "651": "微波炉",
+  "652": "军装",
+  "653": "奶桶",
+  "654": "迷你巴士",
+  "655": "迷你裙",
+  "656": "面包车",
+  "657": "导弹",
+  "658": "连指手套",
+  "659": "搅拌钵",
+  "660": "活动房屋（由汽车拖拉的）",
+  "661": "T型发动机小汽车",
+  "662": "调制解调器",
+  "663": "修道院",
+  "664": "显示器",
+  "665": "电瓶车",
+  "666": "砂浆",
+  "667": "学士",
+  "668": "清真寺",
+  "669": "蚊帐",
+  "670": "摩托车",
+  "671": "山地自行车",
+  "672": "登山帐",
+  "673": "鼠标,电脑鼠标",
+  "674": "捕鼠器",
+  "675": "搬家车",
+  "676": "口套",
+  "677": "钉子",
+  "678": "颈托",
+  "679": "项链",
+  "680": "乳头（瓶）",
+  "681": "笔记本,笔记本电脑",
+  "682": "方尖碑",
+  "683": "双簧管",
+  "684": "陶笛,卵形笛",
+  "685": "里程表",
+  "686": "滤油器",
+  "687": "风琴,管风琴",
+  "688": "示波器",
+  "689": "罩裙",
+  "690": "牛车",
+  "691": "氧气面罩",
+  "692": "包装",
+  "693": "船桨",
+  "694": "明轮,桨轮",
+  "695": "挂锁,扣锁",
+  "696": "画笔",
+  "697": "睡衣",
+  "698": "宫殿",
+  "699": "排箫,鸣管",
+  "700": "纸巾",
+  "701": "降落伞",
+  "702": "双杠",
+  "703": "公园长椅",
+  "704": "停车收费表,停车计时器",
+  "705": "客车,教练车",
+  "706": "露台,阳台",
+  "707": "付费电话",
+  "708": "基座,基脚",
+  "709": "铅笔盒",
+  "710": "卷笔刀",
+  "711": "香水（瓶）",
+  "712": "培养皿",
+  "713": "复印机",
+  "714": "拨弦片,拨子",
+  "715": "尖顶头盔",
+  "716": "栅栏,栅栏",
+  "717": "皮卡,皮卡车",
+  "718": "桥墩",
+  "719": "存钱罐",
+  "720": "药瓶",
+  "721": "枕头",
+  "722": "乒乓球",
+  "723": "风车",
+  "724": "海盗船",
+  "725": "水罐",
+  "726": "木工刨",
+  "727": "天文馆",
+  "728": "塑料袋",
+  "729": "板架",
+  "730": "犁型铲雪机",
+  "731": "手压皮碗泵",
+  "732": "宝丽来相机",
+  "733": "电线杆",
+  "734": "警车,巡逻车",
+  "735": "雨披",
+  "736": "台球桌",
+  "737": "充气饮料瓶",
+  "738": "花盆",
+  "739": "陶工旋盘",
+  "740": "电钻",
+  "741": "祈祷垫,地毯",
+  "742": "打印机",
+  "743": "监狱",
+  "744": "炮弹,导弹",
+  "745": "投影仪",
+  "746": "冰球",
+  "747": "沙包,吊球",
+  "748": "钱包",
+  "749": "羽管笔",
+  "750": "被子",
+  "751": "赛车",
+  "752": "球拍",
+  "753": "散热器",
+  "754": "收音机",
+  "755": "射电望远镜,无线电反射器",
+  "756": "雨桶",
+  "757": "休闲车,房车",
+  "758": "卷轴,卷筒",
+  "759": "反射式照相机",
+  "760": "冰箱,冰柜",
+  "761": "遥控器",
+  "762": "餐厅,饮食店,食堂",
+  "763": "左轮手枪",
+  "764": "步枪",
+  "765": "摇椅",
+  "766": "电转烤肉架",
+  "767": "橡皮",
+  "768": "橄榄球",
+  "769": "直尺",
+  "770": "跑步鞋",
+  "771": "保险柜",
+  "772": "安全别针",
+  "773": "盐瓶（调味用）",
+  "774": "凉鞋",
+  "775": "纱笼,围裙",
+  "776": "萨克斯管",
+  "777": "剑鞘",
+  "778": "秤,称重机",
+  "779": "校车",
+  "780": "帆船",
+  "781": "记分牌",
+  "782": "屏幕",
+  "783": "螺丝",
+  "784": "螺丝刀",
+  "785": "安全带",
+  "786": "缝纫机",
+  "787": "盾牌,盾牌",
+  "788": "皮鞋店,鞋店",
+  "789": "障子",
+  "790": "购物篮",
+  "791": "购物车",
+  "792": "铁锹",
+  "793": "浴帽",
+  "794": "浴帘",
+  "795": "滑雪板",
+  "796": "滑雪面罩",
+  "797": "睡袋",
+  "798": "滑尺",
+  "799": "滑动门",
+  "800": "角子老虎机",
+  "801": "潜水通气管",
+  "802": "雪橇",
+  "803": "扫雪机,扫雪机",
+  "804": "皂液器",
+  "805": "足球",
+  "806": "袜子",
+  "807": "碟式太阳能,太阳能集热器,太阳能炉",
+  "808": "宽边帽",
+  "809": "汤碗",
+  "810": "空格键",
+  "811": "空间加热器",
+  "812": "航天飞机",
+  "813": "铲（搅拌或涂敷用的）",
+  "814": "快艇",
+  "815": "蜘蛛网",
+  "816": "纺锤,纱锭",
+  "817": "跑车",
+  "818": "聚光灯",
+  "819": "舞台",
+  "820": "蒸汽机车",
+  "821": "钢拱桥",
+  "822": "钢滚筒",
+  "823": "听诊器",
+  "824": "女用披肩",
+  "825": "石头墙",
+  "826": "秒表",
+  "827": "火炉",
+  "828": "过滤器",
+  "829": "有轨电车,电车",
+  "830": "担架",
+  "831": "沙发床",
+  "832": "佛塔",
+  "833": "潜艇,潜水艇",
+  "834": "套装,衣服",
+  "835": "日晷",
+  "836": "太阳镜",
+  "837": "太阳镜,墨镜",
+  "838": "防晒霜,防晒剂",
+  "839": "悬索桥",
+  "840": "拖把",
+  "841": "运动衫",
+  "842": "游泳裤",
+  "843": "秋千",
+  "844": "开关,电器开关",
+  "845": "注射器",
+  "846": "台灯",
+  "847": "坦克,装甲战车,装甲战斗车辆",
+  "848": "磁带播放器",
+  "849": "茶壶",
+  "850": "泰迪,泰迪熊",
+  "851": "电视",
+  "852": "网球",
+  "853": "茅草,茅草屋顶",
+  "854": "幕布,剧院的帷幕",
+  "855": "顶针",
+  "856": "脱粒机",
+  "857": "宝座",
+  "858": "瓦屋顶",
+  "859": "烤面包机",
+  "860": "烟草店,烟草",
+  "861": "马桶",
+  "862": "火炬",
+  "863": "图腾柱",
+  "864": "拖车,牵引车,清障车",
+  "865": "玩具店",
+  "866": "拖拉机",
+  "867": "拖车,铰接式卡车",
+  "868": "托盘",
+  "869": "风衣",
+  "870": "三轮车",
+  "871": "三体船",
+  "872": "三脚架",
+  "873": "凯旋门",
+  "874": "无轨电车",
+  "875": "长号",
+  "876": "浴盆,浴缸",
+  "877": "旋转式栅门",
+  "878": "打字机键盘",
+  "879": "伞",
+  "880": "独轮车",
+  "881": "直立式钢琴",
+  "882": "真空吸尘器",
+  "883": "花瓶",
+  "884": "拱顶",
+  "885": "天鹅绒",
+  "886": "自动售货机",
+  "887": "祭服",
+  "888": "高架桥",
+  "889": "小提琴,小提琴",
+  "890": "排球",
+  "891": "松饼机",
+  "892": "挂钟",
+  "893": "钱包,皮夹",
+  "894": "衣柜,壁橱",
+  "895": "军用飞机",
+  "896": "洗脸盆,洗手盆",
+  "897": "洗衣机,自动洗衣机",
+  "898": "水瓶",
+  "899": "水壶",
+  "900": "水塔",
+  "901": "威士忌壶",
+  "902": "哨子",
+  "903": "假发",
+  "904": "纱窗",
+  "905": "百叶窗",
+  "906": "温莎领带",
+  "907": "葡萄酒瓶",
+  "908": "飞机翅膀,飞机",
+  "909": "炒菜锅",
+  "910": "木制的勺子",
+  "911": "毛织品,羊绒",
+  "912": "栅栏,围栏",
+  "913": "沉船",
+  "914": "双桅船",
+  "915": "蒙古包",
+  "916": "网站,互联网网站",
+  "917": "漫画",
+  "918": "纵横字谜",
+  "919": "路标",
+  "920": "交通信号灯",
+  "921": "防尘罩,书皮",
+  "922": "菜单",
+  "923": "盘子",
+  "924": "鳄梨酱",
+  "925": "清汤",
+  "926": "罐焖土豆烧肉",
+  "927": "蛋糕",
+  "928": "冰淇淋",
+  "929": "雪糕,冰棍,冰棒",
+  "930": "法式面包",
+  "931": "百吉饼",
+  "932": "椒盐脆饼",
+  "933": "芝士汉堡",
+  "934": "热狗",
+  "935": "土豆泥",
+  "936": "结球甘蓝",
+  "937": "西兰花",
+  "938": "菜花",
+  "939": "绿皮密生西葫芦",
+  "940": "西葫芦",
+  "941": "小青南瓜",
+  "942": "南瓜",
+  "943": "黄瓜",
+  "944": "朝鲜蓟",
+  "945": "甜椒",
+  "946": "刺棘蓟",
+  "947": "蘑菇",
+  "948": "绿苹果",
+  "949": "草莓",
+  "950": "橘子",
+  "951": "柠檬",
+  "952": "无花果",
+  "953": "菠萝",
+  "954": "香蕉",
+  "955": "菠萝蜜",
+  "956": "蛋奶冻苹果",
+  "957": "石榴",
+  "958": "干草",
+  "959": "烤面条加干酪沙司",
+  "960": "巧克力酱,巧克力糖浆",
+  "961": "面团",
+  "962": "瑞士肉包,肉饼",
+  "963": "披萨,披萨饼",
+  "964": "馅饼",
+  "965": "卷饼",
+  "966": "红葡萄酒",
+  "967": "意大利浓咖啡",
+  "968": "杯子",
+  "969": "蛋酒",
+  "970": "高山",
+  "971": "泡泡",
+  "972": "悬崖",
+  "973": "珊瑚礁",
+  "974": "间歇泉",
+  "975": "湖边,湖岸",
+  "976": "海角",
+  "977": "沙洲,沙坝",
+  "978": "海滨,海岸",
+  "979": "峡谷",
+  "980": "火山",
+  "981": "棒球,棒球运动员",
+  "982": "新郎",
+  "983": "潜水员",
+  "984": "油菜",
+  "985": "雏菊",
+  "986": "杓兰",
+  "987": "玉米",
+  "988": "橡子",
+  "989": "玫瑰果",
+  "990": "七叶树果实",
+  "991": "珊瑚菌",
+  "992": "木耳",
+  "993": "鹿花菌",
+  "994": "鬼笔菌",
+  "995": "地星（菌类）",
+  "996": "多叶奇果菌",
+  "997": "牛肝菌",
+  "998": "玉米穗",
+  "999": "卫生纸"
+}

labels/id2label_en.json

@@ -0,0 +1,1002 @@
+{
+  "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"
+}

labels/imagenet_labels.py

@@ -0,0 +1,61 @@
+"""ImageNet-1k class labels for ADM class-conditional generation.
+
+Labels are stored as Hugging Face-style ``id2label`` JSON maps (string keys ``"0"``–``"999"``).
+Each value is a comma-separated list of synonyms for that class id.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Literal
+
+Language = Literal["en", "cn"]
+
+_LABELS_DIR = Path(__file__).resolve().parent
+
+
+def load_id2label(
+    labels_dir: Path | str | None = None,
+    lang: Language = "en",
+) -> dict[int, str]:
+    """Load ``id2label`` from ``id2label_en.json`` or ``id2label_cn.json``."""
+    root = Path(labels_dir) if labels_dir is not None else _LABELS_DIR
+    filename = "id2label_en.json" if lang == "en" else "id2label_cn.json"
+    path = root / filename
+    if not path.exists():
+        raise FileNotFoundError(f"ImageNet label file not found: {path}")
+
+    raw = json.loads(path.read_text(encoding="utf-8"))
+    return {int(key): value for key, value in raw.items()}
+
+
+def build_label2id(id2label: dict[int, str]) -> dict[str, int]:
+    """Build a synonym -> class id map from an ``id2label`` dict (DiT-style)."""
+    labels: dict[str, int] = {}
+    for class_id, value in id2label.items():
+        for synonym in value.split(","):
+            synonym = synonym.strip()
+            if synonym:
+                labels[synonym] = int(class_id)
+    return dict(sorted(labels.items()))
+
+
+def resolve_label_ids(
+    labels: str | list[str],
+    label2id: dict[str, int],
+    *,
+    lang: Language = "en",
+) -> list[int]:
+    """Map one or more label strings to ImageNet class ids."""
+    if isinstance(labels, str):
+        labels = [labels]
+
+    missing = [label for label in labels if label not in label2id]
+    if missing:
+        preview = ", ".join(list(label2id.keys())[:8])
+        raise ValueError(
+            f"Unknown label(s) for lang={lang!r}: {missing}. "
+            f"Example valid labels: {preview}, ..."
+        )
+    return [label2id[label] for label in labels]