Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+PixelDiT-T2I-1024/demo.png filter=lfs diff=lfs merge=lfs -text
+PixelDiT-T2I-1024/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+PixelDiT-XL-16-256/demo.png filter=lfs diff=lfs merge=lfs -text
+PixelDiT-XL-16-512/demo.png filter=lfs diff=lfs merge=lfs -text

PixelDiT-T2I-1024/model_index.json

@@ -0,0 +1,23 @@
+{
+  "_class_name": [
+    "pipeline",
+    "PixelDiTT2IPipeline"
+  ],
+  "_diffusers_version": "0.35.1",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "transformer": [
+    "transformer_pixeldit_t2i",
+    "PixelDiTT2ITransformer2DModel"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Gemma2Model"
+  ],
+  "tokenizer": [
+    "transformers",
+    "GemmaTokenizerFast"
+  ]
+}

PixelDiT-T2I-1024/pipeline.py

@@ -0,0 +1,401 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+import importlib.util
+
+def _load_flow_dpm_solver():
+    diffusers_src = os.environ.get("PIXELDIT_DIFFUSERS_SRC")
+    if diffusers_src:
+        diffusers_src = str(Path(diffusers_src).resolve())
+        if diffusers_src not in sys.path:
+            sys.path.insert(0, diffusers_src)
+    else:
+        for parent in Path(__file__).resolve().parents:
+            candidate_root = parent / "libs/diffusers/src"
+            if (candidate_root / "diffusers/schedulers/flow_dpm.py").is_file():
+                if str(candidate_root) not in sys.path:
+                    sys.path.insert(0, str(candidate_root))
+                break
+
+    from diffusers.schedulers.flow_dpm import create_flow_dpm_solver
+
+    return create_flow_dpm_solver
+
+
+DEFAULT_TEXT_ENCODER_REPO = "google/gemma-2-2b-it"
+DEFAULT_NEGATIVE_PROMPT = "low quality, worst quality, over-saturated, blurry, deformed, watermark"
+DEFAULT_CHI_PROMPT = "\n".join(
+    [
+        'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:',
+        "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
+        "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
+        "Here are examples of how to transform or refine prompts:",
+        "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.",
+        "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.",
+        "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
+        "User Prompt: ",
+    ]
+)
+
+
+def _default_t2i_scheduler(**kwargs) -> FlowMatchEulerDiscreteScheduler:
+    defaults = {
+        "num_train_timesteps": 1000,
+        "shift": 4.0,
+        "stochastic_sampling": False,
+    }
+    defaults.update(kwargs)
+    return FlowMatchEulerDiscreteScheduler(**defaults)
+
+
+class PixelDiTT2IPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-image generation using PixelDiT (Pixel Diffusion Transformer).
+
+    Parameters:
+        transformer ([`PixelDiTT2ITransformer2DModel`]):
+            Text-conditioned PixelDiT transformer that predicts flow-matching velocity in pixel space.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            Configures flow shift for PixelDiT T2I flow DPM sampling.
+        text_encoder (`transformers.PreTrainedModel`, *optional*):
+            Gemma decoder used to embed prompts.
+        tokenizer (`transformers.PreTrainedTokenizer`, *optional*):
+            Tokenizer paired with the text encoder.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer"
+    _optional_components = ["text_encoder", "tokenizer"]
+
+    def __init__(
+        self,
+        transformer,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        text_encoder=None,
+        tokenizer=None,
+        model_max_length: int = 300,
+        default_negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+    ):
+        super().__init__()
+        scheduler = scheduler or _default_t2i_scheduler()
+        self.register_modules(
+            transformer=transformer,
+            scheduler=scheduler,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+        )
+        self.model_max_length = int(model_max_length)
+        self.default_negative_prompt = default_negative_prompt
+
+    def _load_text_stack(
+        self,
+        base_path: Path,
+        text_encoder_subfolder: Optional[str],
+        tokenizer_subfolder: Optional[str],
+        text_encoder_repo: str,
+        model_kwargs: dict,
+    ) -> None:
+        if self.text_encoder is not None and self.tokenizer is not None:
+            return
+
+        text_encoder = self.text_encoder
+        tokenizer = self.tokenizer
+        if text_encoder_subfolder is not None and text_encoder is None:
+            text_encoder = AutoModelForCausalLM.from_pretrained(
+                str(base_path / text_encoder_subfolder), **model_kwargs
+            ).get_decoder()
+        if tokenizer_subfolder is not None and tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(str(base_path / tokenizer_subfolder))
+            tokenizer.padding_side = "right"
+
+        if text_encoder is None or tokenizer is None:
+            tokenizer = tokenizer or AutoTokenizer.from_pretrained(text_encoder_repo)
+            tokenizer.padding_side = "right"
+            text_encoder = text_encoder or AutoModelForCausalLM.from_pretrained(
+                text_encoder_repo, **model_kwargs
+            ).get_decoder()
+
+        self.register_modules(text_encoder=text_encoder, tokenizer=tokenizer)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        model_kwargs = dict(kwargs)
+        transformer_subfolder = model_kwargs.pop("transformer_subfolder", None)
+        scheduler_subfolder = model_kwargs.pop("scheduler_subfolder", None)
+        text_encoder_subfolder = model_kwargs.pop("text_encoder_subfolder", None)
+        tokenizer_subfolder = model_kwargs.pop("tokenizer_subfolder", None)
+        scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
+        text_encoder_repo = model_kwargs.pop("text_encoder_repo", DEFAULT_TEXT_ENCODER_REPO)
+        base_path = Path(pretrained_model_name_or_path)
+
+        if transformer_subfolder is None and (base_path / "transformer").exists():
+            transformer_subfolder = "transformer"
+        if scheduler_subfolder is None and (base_path / "scheduler").exists():
+            scheduler_subfolder = "scheduler"
+        if text_encoder_subfolder is None and (base_path / "text_encoder").exists():
+            text_encoder_subfolder = "text_encoder"
+        if tokenizer_subfolder is None and (base_path / "tokenizer").exists():
+            tokenizer_subfolder = "tokenizer"
+
+        try:
+            pipe = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+        except Exception:
+            if transformer_subfolder is not None:
+                transformer_path = str(base_path / transformer_subfolder)
+            else:
+                transformer_path = pretrained_model_name_or_path
+
+            transformer_module_path = Path(transformer_path) / 'transformer_pixeldit_t2i.py'
+            spec = importlib.util.spec_from_file_location('transformer_pixeldit_t2i', transformer_module_path)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            transformer = module.PixelDiTT2ITransformer2DModel.from_pretrained(transformer_path, **model_kwargs)
+            try:
+                scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+                    pretrained_model_name_or_path,
+                    subfolder=scheduler_subfolder,
+                    **scheduler_kwargs,
+                )
+            except Exception:
+                scheduler = _default_t2i_scheduler(**scheduler_kwargs)
+            pipe = cls(transformer=transformer, scheduler=scheduler)
+
+        pipe._load_text_stack(
+            base_path=base_path,
+            text_encoder_subfolder=text_encoder_subfolder,
+            tokenizer_subfolder=tokenizer_subfolder,
+            text_encoder_repo=text_encoder_repo,
+            model_kwargs=model_kwargs,
+        )
+        return pipe
+
+    @property
+    def _select_index(self) -> List[int]:
+        return [0] + list(range(-self.model_max_length + 1, 0))
+
+    def _ensure_text_stack_loaded(self) -> None:
+        if self.text_encoder is not None and self.tokenizer is not None:
+            return
+        variant_path = getattr(self.config, "_name_or_path", None)
+        base_path = Path(variant_path) if variant_path else Path(".")
+        text_encoder_subfolder = "text_encoder" if (base_path / "text_encoder").exists() else None
+        tokenizer_subfolder = "tokenizer" if (base_path / "tokenizer").exists() else None
+        self._load_text_stack(
+            base_path=base_path,
+            text_encoder_subfolder=text_encoder_subfolder,
+            tokenizer_subfolder=tokenizer_subfolder,
+            text_encoder_repo=DEFAULT_TEXT_ENCODER_REPO,
+            model_kwargs={"torch_dtype": self.transformer.dtype},
+        )
+
+    @torch.inference_mode()
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        do_classifier_free_guidance: bool = True,
+        use_chi_prompt: bool = True,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        self._ensure_text_stack_loaded()
+        if self.text_encoder is None or self.tokenizer is None:
+            raise ValueError("Text-to-image generation requires `text_encoder` and `tokenizer`.")
+
+        device = device or self._execution_device
+        dtype = dtype or self.transformer.dtype
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        batch_size = len(prompt)
+
+        if use_chi_prompt and DEFAULT_CHI_PROMPT:
+            prompts_all = [DEFAULT_CHI_PROMPT + item for item in prompt]
+            chi_prompt_tokens = len(self.tokenizer.encode(DEFAULT_CHI_PROMPT))
+            max_length_all = chi_prompt_tokens + self.model_max_length - 2
+        else:
+            prompts_all = prompt
+            max_length_all = self.model_max_length
+
+        tokenized = self.tokenizer(
+            prompts_all,
+            max_length=max_length_all,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        input_ids = tokenized.input_ids.to(device)
+        attention_mask = tokenized.attention_mask.to(device)
+        prompt_embeds = self.text_encoder(input_ids, attention_mask=attention_mask)[0]
+        select_index = self._select_index
+        prompt_embeds = prompt_embeds[:, select_index]
+        prompt_attention_mask = attention_mask[:, select_index]
+        prompt_embeds = prompt_embeds.to(dtype=dtype)
+
+        if not do_classifier_free_guidance:
+            return prompt_embeds, prompt_attention_mask
+
+        if negative_prompt is None:
+            negative_prompt = self.default_negative_prompt
+        if isinstance(negative_prompt, str):
+            negative_prompt = [negative_prompt] * batch_size
+
+        null_tokenized = self.tokenizer(
+            negative_prompt,
+            max_length=self.model_max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        null_input_ids = null_tokenized.input_ids.to(device)
+        null_attention_mask = null_tokenized.attention_mask.to(device)
+        negative_embeds = self.text_encoder(null_input_ids, attention_mask=null_attention_mask)[0].to(dtype=dtype)
+        negative_attention_mask = null_attention_mask
+
+        prompt_embeds = torch.cat([negative_embeds, prompt_embeds], dim=0)
+        prompt_attention_mask = torch.cat([negative_attention_mask, prompt_attention_mask], dim=0)
+        return prompt_embeds, prompt_attention_mask
+
+    @staticmethod
+    def _resolve_flow_shift(scheduler, image_size: int) -> float:
+        flow_shift = getattr(scheduler.config, "flow_shift", None)
+        if flow_shift is not None:
+            return float(flow_shift)
+        shift = getattr(scheduler.config, "shift", None)
+        if shift is not None:
+            return float(shift)
+        return 4.0 if image_size >= 1024 else 1.0
+
+    def _make_dpm_transformer_fn(self):
+        transformer = self.transformer
+
+        def forward(x, t_input, y, **model_kwargs):
+            if y.dim() == 4:
+                y = y.squeeze(1)
+            return transformer(
+                x,
+                timestep=t_input.to(dtype=transformer.dtype),
+                encoder_hidden_states=y.to(dtype=transformer.dtype),
+                encoder_attention_mask=None,
+            ).sample
+
+        return forward
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        guidance_scale: Optional[float] = None,
+        guidance_interval_min: float = 0.0,
+        guidance_interval_max: float = 1.0,
+        use_chi_prompt: bool = True,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        num_inference_steps: int = 50,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        if num_inference_steps < 1:
+            raise ValueError("num_inference_steps must be >= 1.")
+        if output_type not in {"pil", "np", "pt"}:
+            raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        batch_size = len(prompt)
+
+        image_size = int(getattr(self.transformer.config, "sample_size", 1024))
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
+        channels = int(self.transformer.config.in_channels)
+
+        if guidance_scale is None:
+            guidance_scale = 2.75
+
+        prompt_embeds, _ = self.encode_prompt(
+            prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=True,
+            use_chi_prompt=use_chi_prompt,
+        )
+
+        latents = randn_tensor(
+            shape=(batch_size, channels, height, width),
+            generator=generator,
+            device=self._execution_device,
+            dtype=self.transformer.dtype,
+        )
+
+        flow_shift = self._resolve_flow_shift(self.scheduler, image_size)
+        solver_order = 2
+        skip_type = "time_uniform_flow"
+        method = "multistep"
+        condition = prompt_embeds[batch_size:].unsqueeze(1)
+        uncondition = prompt_embeds[:batch_size].unsqueeze(1)
+        create_flow_dpm_solver = _load_flow_dpm_solver()
+        dpm_solver = create_flow_dpm_solver(
+            self._make_dpm_transformer_fn(),
+            condition=condition,
+            uncondition=uncondition,
+            cfg_scale=float(guidance_scale),
+            model_kwargs={},
+            interval_guidance=[guidance_interval_min, guidance_interval_max],
+        )
+        os.environ["DPM_TQDM"] = "True"
+        latents = dpm_solver.sample(
+            latents,
+            steps=num_inference_steps,
+            order=solver_order,
+            skip_type=skip_type,
+            method=method,
+            flow_shift=flow_shift,
+        )
+
+        images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
+        if output_type == "pt":
+            images = images_pt
+        elif output_type == "np":
+            images = images_pt.permute(0, 2, 3, 1).numpy()
+        else:
+            images = self.numpy_to_pil(images_pt.permute(0, 2, 3, 1).numpy())
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)
+
+
+PixelDiTT2IPipelineOutput = ImagePipelineOutput

PixelDiT-T2I-1024/scheduler/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.36.0",
+  "num_train_timesteps": 1000,
+  "shift": 4.0,
+  "stochastic_sampling": false
+}

PixelDiT-T2I-1024/text_encoder/config.json

@@ -0,0 +1,36 @@
+{
+  "_name_or_path": "google/gemma-2-2b-it",
+  "architectures": [
+    "Gemma2ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_logit_softcapping": 50.0,
+  "bos_token_id": 2,
+  "cache_implementation": "hybrid",
+  "eos_token_id": [
+    1,
+    107
+  ],
+  "final_logit_softcapping": 30.0,
+  "head_dim": 256,
+  "hidden_act": "gelu_pytorch_tanh",
+  "hidden_activation": "gelu_pytorch_tanh",
+  "hidden_size": 2304,
+  "initializer_range": 0.02,
+  "intermediate_size": 9216,
+  "max_position_embeddings": 8192,
+  "model_type": "gemma2",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 26,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "query_pre_attn_scalar": 256,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000.0,
+  "sliding_window": 4096,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.46.0",
+  "use_cache": true,
+  "vocab_size": 256000
+}

PixelDiT-T2I-1024/text_encoder/configuration.json

@@ -0,0 +1 @@
+{"framework": "pytorch", "task": "text-generation", "allow_remote": true}

PixelDiT-T2I-1024/text_encoder/generation_config.json

@@ -0,0 +1,11 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 2,
+  "cache_implementation": "hybrid",
+  "eos_token_id": [
+    1,
+    107
+  ],
+  "pad_token_id": 0,
+  "transformers_version": "4.46.0"
+}

PixelDiT-T2I-1024/text_encoder/model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:532d792c9178805064170a3ec485b7dedbfccc6fd297b92c31a6091b6c7e41bf
+size 4988025760

PixelDiT-T2I-1024/text_encoder/model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d6d9ce84db398fb6e0191f91542e5da0a73da2cb695e172a24edc2146dc8d20
+size 240691728

PixelDiT-T2I-1024/text_encoder/model.safetensors.index.json

@@ -0,0 +1,295 @@
+{
+  "metadata": {
+    "total_size": 5228683776
+  },
+  "weight_map": {
+    "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.post_feedforward_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.pre_feedforward_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.post_feedforward_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.pre_feedforward_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.post_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.pre_feedforward_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.norm.weight": "model-00002-of-00002.safetensors"
+  }
+}

PixelDiT-T2I-1024/tokenizer/special_tokens_map.json

@@ -0,0 +1,34 @@
+{
+  "additional_special_tokens": [
+    "<start_of_turn>",
+    "<end_of_turn>"
+  ],
+  "bos_token": {
+    "content": "<bos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<eos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

PixelDiT-T2I-1024/tokenizer/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f289bc05132635a8bc7aca7aa21255efd5e18f3710f43e3cdb96bcd41be4922
+size 17525357

PixelDiT-T2I-1024/tokenizer/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2
+size 4241003

PixelDiT-T2I-1024/tokenizer/tokenizer_config.json

@@ -0,0 +1,2013 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<eos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "5": {
+      "content": "<2mass>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "6": {
+      "content": "[@BOS@]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "7": {
+      "content": "<unused0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "8": {
+      "content": "<unused1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "9": {
+      "content": "<unused2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "10": {
+      "content": "<unused3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "11": {
+      "content": "<unused4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "12": {
+      "content": "<unused5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "13": {
+      "content": "<unused6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "14": {
+      "content": "<unused7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "15": {
+      "content": "<unused8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "16": {
+      "content": "<unused9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "17": {
+      "content": "<unused10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "18": {
+      "content": "<unused11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "19": {
+      "content": "<unused12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "20": {
+      "content": "<unused13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "21": {
+      "content": "<unused14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "22": {
+      "content": "<unused15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "23": {
+      "content": "<unused16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "24": {
+      "content": "<unused17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "25": {
+      "content": "<unused18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "26": {
+      "content": "<unused19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "27": {
+      "content": "<unused20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "28": {
+      "content": "<unused21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "29": {
+      "content": "<unused22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "30": {
+      "content": "<unused23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "31": {
+      "content": "<unused24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "32": {
+      "content": "<unused25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "33": {
+      "content": "<unused26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "34": {
+      "content": "<unused27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "35": {
+      "content": "<unused28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "36": {
+      "content": "<unused29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "37": {
+      "content": "<unused30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "38": {
+      "content": "<unused31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "39": {
+      "content": "<unused32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "40": {
+      "content": "<unused33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "41": {
+      "content": "<unused34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "42": {
+      "content": "<unused35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "43": {
+      "content": "<unused36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "44": {
+      "content": "<unused37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "45": {
+      "content": "<unused38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "46": {
+      "content": "<unused39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "47": {
+      "content": "<unused40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "48": {
+      "content": "<unused41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "49": {
+      "content": "<unused42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50": {
+      "content": "<unused43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "51": {
+      "content": "<unused44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "52": {
+      "content": "<unused45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "53": {
+      "content": "<unused46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "54": {
+      "content": "<unused47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "55": {
+      "content": "<unused48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "56": {
+      "content": "<unused49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "57": {
+      "content": "<unused50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "58": {
+      "content": "<unused51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "59": {
+      "content": "<unused52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "60": {
+      "content": "<unused53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "61": {
+      "content": "<unused54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "62": {
+      "content": "<unused55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "63": {
+      "content": "<unused56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "64": {
+      "content": "<unused57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "65": {
+      "content": "<unused58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "66": {
+      "content": "<unused59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "67": {
+      "content": "<unused60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "68": {
+      "content": "<unused61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "69": {
+      "content": "<unused62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "70": {
+      "content": "<unused63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "71": {
+      "content": "<unused64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "72": {
+      "content": "<unused65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "73": {
+      "content": "<unused66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "74": {
+      "content": "<unused67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "75": {
+      "content": "<unused68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "76": {
+      "content": "<unused69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "77": {
+      "content": "<unused70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "78": {
+      "content": "<unused71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "79": {
+      "content": "<unused72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "80": {
+      "content": "<unused73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "81": {
+      "content": "<unused74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "82": {
+      "content": "<unused75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "83": {
+      "content": "<unused76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "84": {
+      "content": "<unused77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "85": {
+      "content": "<unused78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "86": {
+      "content": "<unused79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "87": {
+      "content": "<unused80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "88": {
+      "content": "<unused81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "89": {
+      "content": "<unused82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "90": {
+      "content": "<unused83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "91": {
+      "content": "<unused84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "92": {
+      "content": "<unused85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "93": {
+      "content": "<unused86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "94": {
+      "content": "<unused87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "95": {
+      "content": "<unused88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "96": {
+      "content": "<unused89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "97": {
+      "content": "<unused90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "98": {
+      "content": "<unused91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "99": {
+      "content": "<unused92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100": {
+      "content": "<unused93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "101": {
+      "content": "<unused94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "102": {
+      "content": "<unused95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "103": {
+      "content": "<unused96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "104": {
+      "content": "<unused97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "105": {
+      "content": "<unused98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "106": {
+      "content": "<start_of_turn>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "107": {
+      "content": "<end_of_turn>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "108": {
+      "content": "\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "109": {
+      "content": "\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "110": {
+      "content": "\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "111": {
+      "content": "\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "112": {
+      "content": "\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "113": {
+      "content": "\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "114": {
+      "content": "\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "115": {
+      "content": "\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "116": {
+      "content": "\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "117": {
+      "content": "\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "118": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "119": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "120": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "121": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "122": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "123": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "124": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "125": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "126": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "127": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "128": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "129": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "130": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "131": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "132": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "133": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "134": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "135": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "136": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "137": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "138": {
+      "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "139": {
+      "content": "▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "140": {
+      "content": "▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "141": {
+      "content": "▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "142": {
+      "content": "▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "143": {
+      "content": "▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "144": {
+      "content": "▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "145": {
+      "content": "▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "146": {
+      "content": "▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "147": {
+      "content": "▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "148": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "149": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "150": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "152": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "153": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "154": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "155": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "156": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "157": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "158": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "159": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "160": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "161": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "162": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "163": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "164": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "165": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "166": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "167": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "168": {
+      "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "169": {
+      "content": "<table>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "170": {
+      "content": "<caption>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "171": {
+      "content": "<thead>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "172": {
+      "content": "<tbody>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "173": {
+      "content": "<tfoot>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "174": {
+      "content": "<tr>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "175": {
+      "content": "<th>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "176": {
+      "content": "<td>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "177": {
+      "content": "</table>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "178": {
+      "content": "</caption>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "179": {
+      "content": "</thead>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "180": {
+      "content": "</tbody>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "181": {
+      "content": "</tfoot>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "182": {
+      "content": "</tr>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "183": {
+      "content": "</th>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "184": {
+      "content": "</td>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "185": {
+      "content": "<h1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "186": {
+      "content": "<h2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "187": {
+      "content": "<h3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "188": {
+      "content": "<h4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "189": {
+      "content": "<h5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "190": {
+      "content": "<h6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "191": {
+      "content": "<blockquote>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "192": {
+      "content": "</h1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "193": {
+      "content": "</h2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "194": {
+      "content": "</h3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "195": {
+      "content": "</h4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "196": {
+      "content": "</h5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "197": {
+      "content": "</h6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "198": {
+      "content": "</blockquote>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "199": {
+      "content": "<strong>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "200": {
+      "content": "<em>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "201": {
+      "content": "<b>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "202": {
+      "content": "<i>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "203": {
+      "content": "<u>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "204": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "205": {
+      "content": "<sub>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "206": {
+      "content": "<sup>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "207": {
+      "content": "<code>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "208": {
+      "content": "</strong>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "209": {
+      "content": "</em>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "210": {
+      "content": "</b>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "211": {
+      "content": "</i>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "212": {
+      "content": "</u>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "213": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "214": {
+      "content": "</sub>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "215": {
+      "content": "</sup>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "216": {
+      "content": "</code>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255968": {
+      "content": "[toxicity=0]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255969": {
+      "content": "\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255970": {
+      "content": "\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255971": {
+      "content": "\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255972": {
+      "content": "\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255973": {
+      "content": "\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255974": {
+      "content": "\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255975": {
+      "content": "\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255976": {
+      "content": "\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255977": {
+      "content": "\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255978": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255979": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255980": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255981": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255982": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255983": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255984": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255985": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255986": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255987": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255988": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255989": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255990": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255991": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255992": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255993": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255994": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255995": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255996": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255997": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255998": {
+      "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "255999": {
+      "content": "<unused99>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<start_of_turn>",
+    "<end_of_turn>"
+  ],
+  "bos_token": "<bos>",
+  "chat_template": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<eos>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

PixelDiT-T2I-1024/transformer/config.json

@@ -0,0 +1,21 @@
+{
+  "_class_name": "PixelDiTT2ITransformer2DModel",
+  "hidden_size": 1536,
+  "in_channels": 3,
+  "model_type": "pixeldit",
+  "num_groups": 24,
+  "num_text_blocks": 4,
+  "patch_depth": 14,
+  "patch_size": 16,
+  "pixel_attn_hidden_size": 1152,
+  "pixel_depth": 2,
+  "pixel_hidden_size": 16,
+  "pixel_num_groups": 16,
+  "repa_encoder_index": 6,
+  "sample_size": 1024,
+  "text_rope_theta": 10000.0,
+  "txt_embed_dim": 2304,
+  "txt_max_length": 300,
+  "use_pixel_abs_pos": true,
+  "use_text_rope": true
+}

PixelDiT-T2I-1024/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c06618b1540ff926b6bb165d0dfbe14c1b228e66b28e72459dcef1c92eb7768c
+size 5209924596

PixelDiT-T2I-1024/transformer/transformer_pixeldit.py

@@ -0,0 +1,639 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from typing import Dict, Literal, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import RMSNorm
+
+
+PIXELDIT_PRESET_CONFIGS: Dict[str, Dict[str, object]] = {
+    "pixeldit-xl": {
+        "sample_size": 256,
+        "num_groups": 16,
+        "hidden_size": 1152,
+        "pixel_hidden_size": 16,
+        "patch_depth": 26,
+        "pixel_depth": 4,
+        "patch_size": 16,
+    },
+}
+
+
+def remap_legacy_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Map wrapper/backbone keys from legacy checkpoints to native PixelDiTTransformer2DModel keys."""
+    remapped: Dict[str, torch.Tensor] = {}
+    prefixes = ("transformer.", "model.", "module.", "denoiser.", "net.")
+    for key, value in state_dict.items():
+        new_key = key
+        for prefix in prefixes:
+            if new_key.startswith(prefix):
+                new_key = new_key[len(prefix) :]
+                break
+        remapped[new_key] = value
+    return remapped
+
+
+def config_from_legacy(config: Dict[str, object]) -> Dict[str, object]:
+    """Build native config kwargs from a legacy config.json dict."""
+    model_type = config.get("model_type") or config.get("model_name") or config.get("model_size")
+    if model_type not in PIXELDIT_PRESET_CONFIGS:
+        raise ValueError(f"Unknown PixelDiT preset '{model_type}'. Known: {list(PIXELDIT_PRESET_CONFIGS)}")
+
+    preset = dict(PIXELDIT_PRESET_CONFIGS[model_type])
+    preset["num_classes"] = int(config.get("num_classes") or config.get("num_class_embeds") or 1000)
+    preset["in_channels"] = int(config.get("in_channels", 3))
+    preset["use_pixel_abs_pos"] = bool(config.get("use_pixel_abs_pos", True))
+    preset["model_type"] = model_type
+
+    for key in ("sample_size", "num_groups", "hidden_size", "pixel_hidden_size", "patch_depth", "pixel_depth", "patch_size"):
+        if config.get(key) is not None:
+            preset[key] = config[key]
+
+    return preset
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int) -> np.ndarray:
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0)
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    return get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
+    if embed_dim % 2 != 0:
+        raise ValueError("Embedding dimension must be even for 2D sin/cos positional embeddings.")
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    if embed_dim % 2 != 0:
+        raise ValueError("Embedding dimension must be even for 1D sin/cos positional embeddings.")
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+def apply_adaln(hidden_states: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return hidden_states * (1 + scale) + shift
+
+
+def precompute_freqs_cis_2d(dim: int, height: int, width: int, theta: float = 10000.0, scale: float = 16.0):
+    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))
+    x_freqs = torch.outer(x_pos, freqs).float()
+    y_freqs = torch.outer(y_pos, freqs).float()
+    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)
+    return freqs_cis.reshape(height * width, -1)
+
+
+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_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 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)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class TimestepConditioner(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(timesteps: torch.Tensor, dim: int, max_period: int = 10):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half
+        )
+        args = timesteps[..., 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, timesteps: torch.Tensor) -> torch.Tensor:
+        timestep_freq = self.timestep_embedding(timesteps, self.frequency_embedding_size)
+        mlp_dtype = next(self.mlp.parameters()).dtype
+        if timestep_freq.dtype != mlp_dtype:
+            timestep_freq = timestep_freq.to(mlp_dtype)
+        return self.mlp(timestep_freq)
+
+
+class ClassEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int):
+        super().__init__()
+        self.embedding_table = nn.Embedding(num_classes, hidden_size)
+        self.num_classes = num_classes
+
+    def forward(self, labels: torch.Tensor) -> torch.Tensor:
+        return self.embedding_table(labels)
+
+
+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, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(hidden_states)) * self.w3(hidden_states))
+
+
+class RotaryAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError("dim should be divisible by num_heads")
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, eps=eps) 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, hidden_states: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        batch_size, length, channels = hidden_states.shape
+        qkv = (
+            self.qkv(hidden_states)
+            .reshape(batch_size, length, 3, self.num_heads, channels // self.num_heads)
+            .permute(2, 0, 1, 3, 4)
+        )
+        query, key, value = qkv[0], qkv[1], qkv[2]
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        query, key = apply_rotary_emb(query, key, freqs_cis=pos)
+        query = query.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2)
+        key = key.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2).contiguous()
+        value = value.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2).contiguous()
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, length, channels)
+        hidden_states = self.proj(hidden_states)
+        return self.proj_drop(hidden_states)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int, mlp_ratio: float = 4.0, drop: float = 0.0):
+        super().__init__()
+        hidden_dim = int(dim * mlp_ratio)
+        self.fc1 = nn.Linear(dim, hidden_dim)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_dim, dim)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.drop(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return self.drop(hidden_states)
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, out_channels: int, eps: float = 1e-6):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size, eps=eps)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(hidden_states)
+        return self.linear(hidden_states)
+
+
+class PatchTokenEmbedder(nn.Module):
+    def __init__(self, in_chans: int, embed_dim: int, 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, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.proj(hidden_states)
+        return self.norm(hidden_states)
+
+
+class PixelTokenEmbedder(nn.Module):
+    def __init__(self, in_channels: int, hidden_size_output: int, use_pixel_abs_pos: bool = True):
+        super().__init__()
+        self.in_channels = int(in_channels)
+        self.hidden_size_output = int(hidden_size_output)
+        self.use_pixel_abs_pos = bool(use_pixel_abs_pos)
+        self.proj = nn.Linear(self.in_channels, self.hidden_size_output, bias=True)
+        self._pos_cache: Dict[tuple[str, int, int], torch.Tensor] = {}
+
+    def _fetch_pixel_pos_image(self, height: int, width: int, device: torch.device, dtype: torch.dtype):
+        key = ("image", height, width)
+        if key in self._pos_cache:
+            return self._pos_cache[key].to(device=device, dtype=dtype)
+        if height == width:
+            pos_np = get_2d_sincos_pos_embed(self.hidden_size_output, height)
+        else:
+            grid_h = np.arange(height, dtype=np.float32)
+            grid_w = np.arange(width, dtype=np.float32)
+            grid = np.meshgrid(grid_w, grid_h)
+            grid = np.stack(grid, axis=0).reshape(2, 1, height, width)
+            pos_np = get_2d_sincos_pos_embed_from_grid(self.hidden_size_output, grid)
+        pos = torch.from_numpy(pos_np).to(device=device, dtype=dtype)
+        self._pos_cache[key] = pos
+        return pos
+
+    def forward(self, inputs: torch.Tensor, img_height: int, img_width: int, patch_size: int):
+        if inputs.dim() != 4:
+            raise ValueError("PixelTokenEmbedder expects inputs of shape [B,C,H,W]")
+        batch_size, channels, height, width = inputs.shape
+        if height != img_height or width != img_width:
+            raise ValueError("Input resolution does not match img_height/img_width.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+        h_tokens, w_tokens = height // patch_size, width // patch_size
+        patch_area = patch_size * patch_size
+        hidden_states = inputs.permute(0, 2, 3, 1).contiguous()
+        hidden_states = self.proj(hidden_states)
+        if self.use_pixel_abs_pos:
+            pos_full = self._fetch_pixel_pos_image(height, width, inputs.device, inputs.dtype)
+            hidden_states = hidden_states + pos_full.view(height, width, self.hidden_size_output).unsqueeze(0)
+        hidden_states = hidden_states.view(batch_size, h_tokens, patch_size, w_tokens, patch_size, self.hidden_size_output)
+        hidden_states = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous()
+        return hidden_states.view(batch_size * h_tokens * w_tokens, patch_area, self.hidden_size_output)
+
+
+class AugmentedDiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, groups: int, mlp_ratio: float = 4.0, adaLN_modulation=None, eps: float = 1e-6):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size, eps=eps)
+        self.attn = RotaryAttention(hidden_size, num_heads=groups, qkv_bias=False, eps=eps)
+        self.norm2 = RMSNorm(hidden_size, eps=eps)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp = FeedForward(hidden_size, mlp_hidden_dim)
+        self.adaLN_modulation = adaLN_modulation if adaLN_modulation is not None else nn.Sequential(
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, hidden_states: torch.Tensor, conditioning: torch.Tensor, pos: torch.Tensor):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(conditioning).chunk(
+            6, dim=-1
+        )
+        hidden_states = hidden_states + gate_msa * self.attn(
+            apply_adaln(self.norm1(hidden_states), shift_msa, scale_msa), pos
+        )
+        hidden_states = hidden_states + gate_mlp * self.mlp(
+            apply_adaln(self.norm2(hidden_states), shift_mlp, scale_mlp)
+        )
+        return hidden_states
+
+
+class PiTBlock(nn.Module):
+    def __init__(
+        self,
+        pixel_hidden_size: int,
+        patch_hidden_size: int,
+        patch_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_hidden_size: Optional[int] = None,
+        attn_num_heads: Optional[int] = None,
+        rope_fn=None,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.pixel_dim = int(pixel_hidden_size)
+        self.context_dim = int(patch_hidden_size)
+        self.patch_size = int(patch_size)
+        self.attn_dim = int(attn_hidden_size) if attn_hidden_size is not None else self.context_dim
+        self.num_heads = int(attn_num_heads) if attn_num_heads is not None else int(num_heads)
+        if self.attn_dim % self.num_heads != 0:
+            raise ValueError("pixel attention hidden size must be divisible by pixel num_heads")
+        patch_area = self.patch_size * self.patch_size
+        self.compress_to_attn = nn.Linear(patch_area * self.pixel_dim, self.attn_dim, bias=True)
+        self.expand_from_attn = nn.Linear(self.attn_dim, patch_area * self.pixel_dim, bias=True)
+        self.norm1 = RMSNorm(self.pixel_dim, eps=eps)
+        self.attn = RotaryAttention(self.attn_dim, num_heads=self.num_heads, qkv_bias=False, eps=eps)
+        self.norm2 = RMSNorm(self.pixel_dim, eps=eps)
+        self.mlp = MLP(self.pixel_dim, mlp_ratio=mlp_ratio, drop=0.0)
+        self.adaLN_modulation = nn.Sequential(nn.Linear(self.context_dim, 6 * self.pixel_dim * patch_area, bias=True))
+        self._pos_cache: Dict[tuple[int, int], torch.Tensor] = {}
+        self._rope_fn = rope_fn if rope_fn is not None else precompute_freqs_cis_2d
+
+    def _fetch_pos(self, height: int, width: int, device: torch.device):
+        key = (height, width)
+        if key in self._pos_cache:
+            return self._pos_cache[key].to(device)
+        pos = self._rope_fn(self.attn_dim // self.num_heads, height, width).to(device)
+        self._pos_cache[key] = pos
+        return pos
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        conditioning: torch.Tensor,
+        image_height: int,
+        image_width: int,
+        patch_size: int,
+    ) -> torch.Tensor:
+        batch_tokens, patch_area, channels = hidden_states.shape
+        if channels != self.pixel_dim:
+            raise ValueError(f"PiTBlock expected pixel_dim={self.pixel_dim}, got {channels}")
+        if image_height % patch_size != 0 or image_width % patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+        h_tokens, w_tokens = image_height // patch_size, image_width // patch_size
+        length = h_tokens * w_tokens
+        batch_size = batch_tokens // length
+        cond_params = self.adaLN_modulation(conditioning).view(batch_tokens, patch_area, 6 * self.pixel_dim)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(cond_params, 6, dim=-1)
+        hidden_norm = apply_adaln(self.norm1(hidden_states), shift_msa, scale_msa)
+        hidden_flat = hidden_norm.view(batch_tokens, patch_area * self.pixel_dim)
+        hidden_comp = self.compress_to_attn(hidden_flat).view(batch_size, length, self.attn_dim)
+        pos_comp = self._fetch_pos(h_tokens, w_tokens, hidden_states.device)
+        attn_out = self.attn(hidden_comp, pos_comp)
+        attn_flat = self.expand_from_attn(attn_out.view(batch_size * length, self.attn_dim))
+        attn_exp = attn_flat.view(batch_tokens, patch_area, self.pixel_dim)
+        hidden_states = hidden_states + gate_msa * attn_exp
+        mlp_out = self.mlp(apply_adaln(self.norm2(hidden_states), shift_mlp, scale_mlp))
+        hidden_states = hidden_states + gate_mlp * mlp_out
+        return hidden_states
+
+
+class PixelDiTTransformer2DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+    _skip_layerwise_casting_patterns = ["pos", "_pos_cache"]
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 256,
+        in_channels: int = 3,
+        num_groups: int = 16,
+        hidden_size: int = 1152,
+        pixel_hidden_size: int = 16,
+        patch_depth: int = 26,
+        pixel_depth: int = 4,
+        patch_size: int = 16,
+        num_classes: int = 1000,
+        use_pixel_abs_pos: bool = True,
+        norm_eps: float = 1e-6,
+        model_type: str | None = None,
+        num_class_embeds: int | None = None,
+    ):
+        super().__init__()
+        if num_class_embeds is not None:
+            num_classes = int(num_class_embeds)
+        if model_type in PIXELDIT_PRESET_CONFIGS:
+            preset = PIXELDIT_PRESET_CONFIGS[model_type]
+            sample_size = int(preset["sample_size"])
+            num_groups = int(preset["num_groups"])
+            hidden_size = int(preset["hidden_size"])
+            pixel_hidden_size = int(preset["pixel_hidden_size"])
+            patch_depth = int(preset["patch_depth"])
+            pixel_depth = int(preset["pixel_depth"])
+            patch_size = int(preset["patch_size"])
+
+        self.sample_size = int(sample_size)
+        self.in_channels = int(in_channels)
+        self.out_channels = int(in_channels)
+        self.hidden_size = int(hidden_size)
+        self.num_groups = int(num_groups)
+        self.patch_depth = int(patch_depth)
+        self.pixel_depth = int(pixel_depth)
+        self.patch_size = int(patch_size)
+        self.pixel_hidden_size = int(pixel_hidden_size)
+        self.num_classes = int(num_classes)
+        self.use_pixel_abs_pos = bool(use_pixel_abs_pos)
+        self.norm_eps = float(norm_eps)
+        self.gradient_checkpointing = False
+
+        if self.pixel_depth <= 0:
+            raise ValueError("PixelDiT expects pixel_depth > 0 to preserve the dual-level pipeline")
+
+        self.pixel_embedder = PixelTokenEmbedder(
+            self.in_channels, self.pixel_hidden_size, use_pixel_abs_pos=self.use_pixel_abs_pos
+        )
+        self.s_embedder = PatchTokenEmbedder(self.in_channels * self.patch_size**2, self.hidden_size, bias=True)
+        self.t_embedder = TimestepConditioner(self.hidden_size)
+        self.y_embedder = ClassEmbedder(self.num_classes + 1, self.hidden_size)
+
+        self.final_layer = FinalLayer(self.pixel_hidden_size, self.out_channels, eps=self.norm_eps)
+        self.patch_blocks = nn.ModuleList(
+            [AugmentedDiTBlock(self.hidden_size, self.num_groups, eps=self.norm_eps) for _ in range(self.patch_depth)]
+        )
+        self.pixel_blocks = nn.ModuleList(
+            [
+                PiTBlock(
+                    self.pixel_hidden_size,
+                    self.hidden_size,
+                    patch_size=self.patch_size,
+                    num_heads=self.num_groups,
+                    mlp_ratio=4.0,
+                    eps=self.norm_eps,
+                )
+                for _ in range(self.pixel_depth)
+            ]
+        )
+        self._precompute_pos: Dict[tuple[int, int], torch.Tensor] = {}
+        self._initialize_weights()
+
+    def _fetch_pos(self, height: int, width: int, device: torch.device):
+        key = (height, width)
+        if key in self._precompute_pos:
+            return self._precompute_pos[key].to(device)
+        pos = precompute_freqs_cis_2d(self.hidden_size // self.num_groups, height, width).to(device)
+        self._precompute_pos[key] = pos
+        return pos
+
+    def _initialize_weights(self) -> None:
+        weight = self.s_embedder.proj.weight.data
+        nn.init.xavier_uniform_(weight.view([weight.shape[0], -1]))
+        nn.init.constant_(self.s_embedder.proj.bias, 0)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        nn.init.zeros_(self.final_layer.linear.weight)
+        nn.init.zeros_(self.final_layer.linear.bias)
+        for block in self.patch_blocks:
+            nn.init.zeros_(block.adaLN_modulation[0].weight)
+            nn.init.zeros_(block.adaLN_modulation[0].bias)
+        for block in self.pixel_blocks:
+            nn.init.zeros_(block.adaLN_modulation[0].weight)
+            nn.init.zeros_(block.adaLN_modulation[0].bias)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float],
+        class_labels: Union[torch.Tensor, int],
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]:
+        if sample.dim() != 4:
+            raise ValueError("PixelDiTTransformer2DModel expects sample of shape [B,C,H,W]")
+        batch_size, _, height, width = sample.shape
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+
+        timestep = torch.as_tensor(timestep, device=sample.device)
+        if timestep.ndim == 0:
+            timestep = timestep.repeat(batch_size)
+        else:
+            timestep = timestep.reshape(-1)
+            if timestep.shape[0] == 1 and batch_size > 1:
+                timestep = timestep.repeat(batch_size)
+
+        if not torch.is_tensor(class_labels):
+            class_labels = torch.tensor(class_labels, device=sample.device, dtype=torch.long)
+        class_labels = class_labels.to(device=sample.device, dtype=torch.long).reshape(-1)
+        if class_labels.shape[0] == 1 and batch_size > 1:
+            class_labels = class_labels.repeat(batch_size)
+
+        pos = self._fetch_pos(height // self.patch_size, width // self.patch_size, sample.device)
+        x_patches = F.unfold(sample, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2)
+        t_emb = self.t_embedder(timestep.view(-1)).view(batch_size, -1, self.hidden_size)
+        y_emb = self.y_embedder(class_labels).view(batch_size, 1, self.hidden_size)
+        conditioning = F.silu(t_emb + y_emb)
+
+        patch_states = self.s_embedder(x_patches)
+        for block in self.patch_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(hidden_states, cond, position):
+                    return block(hidden_states, cond, position)
+
+                patch_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward, patch_states, conditioning, pos, use_reentrant=False
+                )
+            else:
+                patch_states = block(patch_states, conditioning, pos)
+        patch_states = F.silu(t_emb + patch_states)
+
+        length = patch_states.shape[1]
+        conditioning_states = patch_states.view(batch_size * length, self.hidden_size)
+        pixel_states = self.pixel_embedder(
+            sample, img_height=height, img_width=width, patch_size=self.patch_size
+        )
+        for block in self.pixel_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(hidden_states, cond):
+                    return block(hidden_states, cond, height, width, self.patch_size)
+
+                pixel_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward, pixel_states, conditioning_states, use_reentrant=False
+                )
+            else:
+                pixel_states = block(pixel_states, conditioning_states, height, width, self.patch_size)
+        pixel_states = self.final_layer(pixel_states)
+
+        patch_area = self.patch_size * self.patch_size
+        pixel_states = pixel_states.view(batch_size, length, patch_area, self.out_channels).permute(0, 3, 2, 1)
+        pixel_states = pixel_states.contiguous().view(batch_size, self.out_channels * patch_area, length)
+        output = F.fold(pixel_states, (height, width), kernel_size=self.patch_size, stride=self.patch_size)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    @classmethod
+    def from_pixeldit_checkpoint(
+        cls,
+        checkpoint_path: str,
+        model_type: Literal["pixeldit-xl"] = "pixeldit-xl",
+        map_location: str = "cpu",
+        strict: bool = True,
+    ) -> Tuple["PixelDiTTransformer2DModel", Dict[str, object]]:
+        if model_type not in PIXELDIT_PRESET_CONFIGS:
+            raise ValueError(f"Unknown PixelDiT preset '{model_type}'.")
+
+        if checkpoint_path.endswith(".safetensors"):
+            try:
+                from safetensors.torch import load_file
+            except ImportError as error:
+                raise ImportError("Install safetensors to load .safetensors checkpoints.") from error
+            state_dict = load_file(checkpoint_path, device=map_location)
+        else:
+            loaded = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
+            if isinstance(loaded, Mapping):
+                state_dict = loaded
+                for key in ("state_dict", "model", "module", "denoiser"):
+                    if key in state_dict and isinstance(state_dict[key], dict):
+                        state_dict = state_dict[key]
+                        break
+            else:
+                raise ValueError("Unsupported checkpoint format.")
+
+        config = dict(PIXELDIT_PRESET_CONFIGS[model_type])
+        config["model_type"] = model_type
+        model = cls(**config)
+        model.load_state_dict(remap_legacy_state_dict(state_dict), strict=strict)
+
+        metadata = {
+            "checkpoint_path": checkpoint_path,
+            "model_type": model_type,
+        }
+        return model, metadata
+
+    def to_pixeldit_checkpoint(self, prefix: str = "") -> Dict[str, torch.Tensor]:
+        checkpoint: Dict[str, torch.Tensor] = {}
+        for key, value in self.state_dict().items():
+            checkpoint[f"{prefix}{key}"] = value.detach().cpu()
+        return checkpoint
+
+
+PixelDiTDiffusersModel = PixelDiTTransformer2DModel

PixelDiT-T2I-1024/transformer/transformer_pixeldit_t2i.py

@@ -0,0 +1,537 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Dict, Literal, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import RMSNorm
+
+import sys
+from pathlib import Path as _Path
+_MODULE_DIR = _Path(__file__).resolve().parent
+if str(_MODULE_DIR) not in sys.path:
+    sys.path.insert(0, str(_MODULE_DIR))
+from transformer_pixeldit import (
+    FinalLayer,
+    FeedForward,
+    PatchTokenEmbedder,
+    PiTBlock,
+    PixelTokenEmbedder,
+    TimestepConditioner,
+    apply_adaln,
+    apply_rotary_emb,
+    precompute_freqs_cis_2d,
+)
+
+
+PIXELDIT_T2I_PRESET_CONFIGS: Dict[str, Dict[str, object]] = {
+    "pixeldit-t2i-1300m": {
+        "sample_size": 1024,
+        "num_groups": 24,
+        "hidden_size": 1536,
+        "pixel_hidden_size": 16,
+        "pixel_attn_hidden_size": 1152,
+        "pixel_num_groups": 16,
+        "patch_depth": 14,
+        "pixel_depth": 2,
+        "num_text_blocks": 4,
+        "patch_size": 16,
+        "txt_embed_dim": 2304,
+        "txt_max_length": 300,
+        "use_text_rope": True,
+        "text_rope_theta": 10000.0,
+        "repa_encoder_index": 6,
+    },
+}
+
+
+def remap_t2i_legacy_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Map legacy T2I checkpoint keys to native PixelDiTT2ITransformer2DModel keys."""
+    remapped: Dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        if key.startswith("_repa_projector."):
+            continue
+        new_key = key[5:] if key.startswith("core.") else key
+        remapped[new_key] = value
+    return remapped
+
+
+def config_from_legacy_t2i(config: Dict[str, object]) -> Dict[str, object]:
+    """Build native T2I config kwargs from a legacy config.json dict."""
+    model_type = config.get("model_type")
+    if model_type == "pixeldit" and config.get("architectures") == ["PixDiT_T2I"]:
+        model_type = "pixeldit-t2i-1300m"
+    if model_type not in PIXELDIT_T2I_PRESET_CONFIGS:
+        raise ValueError(
+            f"Unknown PixelDiT T2I preset '{model_type}'. Known: {list(PIXELDIT_T2I_PRESET_CONFIGS)}"
+        )
+
+    preset = dict(PIXELDIT_T2I_PRESET_CONFIGS[model_type])
+    preset["in_channels"] = int(config.get("in_channels", 3))
+    preset["use_pixel_abs_pos"] = bool(config.get("use_pixel_abs_pos", True))
+    preset["model_type"] = model_type
+
+    for key in preset:
+        if config.get(key) is not None:
+            preset[key] = config[key]
+    if config.get("image_size") is not None:
+        preset["sample_size"] = int(config["image_size"])
+    return preset
+
+
+class MMDiTJointAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError("dim should be divisible by num_heads")
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.qkv_x = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.qkv_y = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm_x = RMSNorm(self.head_dim, eps=eps)
+        self.k_norm_x = RMSNorm(self.head_dim, eps=eps)
+        self.q_norm_y = RMSNorm(self.head_dim, eps=eps)
+        self.k_norm_y = RMSNorm(self.head_dim, eps=eps)
+        self.proj_x = nn.Linear(dim, dim)
+        self.proj_y = nn.Linear(dim, dim)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj_drop_x = nn.Dropout(proj_drop)
+        self.proj_drop_y = nn.Dropout(proj_drop)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        pos_img: torch.Tensor,
+        pos_txt: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_img_tokens, channels = x.shape
+        _, num_txt_tokens, txt_channels = y.shape
+        if batch_size != y.shape[0] or channels != txt_channels:
+            raise ValueError("x and y must share batch and channel dims")
+
+        qkv_x = (
+            self.qkv_x(x)
+            .reshape(batch_size, num_img_tokens, 3, self.num_heads, channels // self.num_heads)
+            .permute(2, 0, 1, 3, 4)
+        )
+        qx, kx, vx = qkv_x[0], qkv_x[1], qkv_x[2]
+        qx = self.q_norm_x(qx)
+        kx = self.k_norm_x(kx)
+
+        qkv_y = (
+            self.qkv_y(y)
+            .reshape(batch_size, num_txt_tokens, 3, self.num_heads, channels // self.num_heads)
+            .permute(2, 0, 1, 3, 4)
+        )
+        qy, ky, vy = qkv_y[0], qkv_y[1], qkv_y[2]
+        qy = self.q_norm_y(qy)
+        ky = self.k_norm_y(ky)
+
+        qx, kx = apply_rotary_emb(qx, kx, freqs_cis=pos_img)
+        if pos_txt is not None:
+            qy, ky = apply_rotary_emb(qy, ky, freqs_cis=pos_txt)
+
+        qx = qx.transpose(1, 2)
+        kx = kx.transpose(1, 2)
+        vx = vx.transpose(1, 2)
+        qy = qy.transpose(1, 2)
+        ky = ky.transpose(1, 2)
+        vy = vy.transpose(1, 2)
+
+        q_joint = torch.cat([qy, qx], dim=2)
+        k_joint = torch.cat([ky, kx], dim=2)
+        v_joint = torch.cat([vy, vx], dim=2)
+
+        out_joint = F.scaled_dot_product_attention(
+            q_joint, k_joint, v_joint, dropout_p=0.0, attn_mask=attn_mask
+        )
+        out_y = out_joint[:, :, :num_txt_tokens, :]
+        out_x = out_joint[:, :, num_txt_tokens:, :]
+
+        out_y = out_y.transpose(1, 2).reshape(batch_size, num_txt_tokens, channels)
+        out_x = out_x.transpose(1, 2).reshape(batch_size, num_img_tokens, channels)
+        out_x = self.proj_drop_x(self.proj_x(out_x))
+        out_y = self.proj_drop_y(self.proj_y(out_y))
+        return out_x, out_y
+
+
+class MMDiTBlockT2I(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        groups: int,
+        mlp_ratio: float = 4.0,
+        adaLN_modulation_img=None,
+        adaLN_modulation_txt=None,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.norm_x1 = RMSNorm(hidden_size, eps=eps)
+        self.norm_y1 = RMSNorm(hidden_size, eps=eps)
+        self.attn = MMDiTJointAttention(hidden_size, num_heads=groups, qkv_bias=False, eps=eps)
+        self.norm_x2 = RMSNorm(hidden_size, eps=eps)
+        self.norm_y2 = RMSNorm(hidden_size, eps=eps)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp_x = FeedForward(hidden_size, mlp_hidden_dim)
+        self.mlp_y = FeedForward(hidden_size, mlp_hidden_dim)
+        self.adaLN_modulation_img = adaLN_modulation_img or nn.Sequential(
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+        self.adaLN_modulation_txt = adaLN_modulation_txt or nn.Sequential(
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        conditioning: torch.Tensor,
+        pos_img: torch.Tensor,
+        pos_txt: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        shift_msa_x, scale_msa_x, gate_msa_x, shift_mlp_x, scale_mlp_x, gate_mlp_x = self.adaLN_modulation_img(
+            conditioning
+        ).chunk(6, dim=-1)
+        shift_msa_y, scale_msa_y, gate_msa_y, shift_mlp_y, scale_mlp_y, gate_mlp_y = self.adaLN_modulation_txt(
+            conditioning
+        ).chunk(6, dim=-1)
+
+        x_norm = apply_adaln(self.norm_x1(x), shift_msa_x, scale_msa_x)
+        y_norm = apply_adaln(self.norm_y1(y), shift_msa_y, scale_msa_y)
+        attn_x, attn_y = self.attn(x_norm, y_norm, pos_img, pos_txt, attn_mask)
+        x = x + gate_msa_x * attn_x
+        y = y + gate_msa_y * attn_y
+        x = x + gate_mlp_x * self.mlp_x(apply_adaln(self.norm_x2(x), shift_mlp_x, scale_mlp_x))
+        y = y + gate_mlp_y * self.mlp_y(apply_adaln(self.norm_y2(y), shift_mlp_y, scale_mlp_y))
+        return x, y
+
+
+class PixelDiTT2ITransformer2DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+    _skip_layerwise_casting_patterns = ["pos", "_pos_cache", "y_pos_embedding"]
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 1024,
+        in_channels: int = 3,
+        num_groups: int = 24,
+        hidden_size: int = 1536,
+        pixel_hidden_size: int = 16,
+        pixel_attn_hidden_size: int = 1152,
+        pixel_num_groups: int = 16,
+        patch_depth: int = 14,
+        pixel_depth: int = 2,
+        num_text_blocks: int = 4,
+        patch_size: int = 16,
+        txt_embed_dim: int = 2304,
+        txt_max_length: int = 300,
+        use_text_rope: bool = True,
+        text_rope_theta: float = 10000.0,
+        repa_encoder_index: int = 6,
+        use_pixel_abs_pos: bool = True,
+        norm_eps: float = 1e-6,
+        model_type: str | None = None,
+    ):
+        super().__init__()
+        if model_type in PIXELDIT_T2I_PRESET_CONFIGS:
+            preset = PIXELDIT_T2I_PRESET_CONFIGS[model_type]
+            sample_size = int(preset["sample_size"])
+            num_groups = int(preset["num_groups"])
+            hidden_size = int(preset["hidden_size"])
+            pixel_hidden_size = int(preset["pixel_hidden_size"])
+            pixel_attn_hidden_size = int(preset["pixel_attn_hidden_size"])
+            pixel_num_groups = int(preset["pixel_num_groups"])
+            patch_depth = int(preset["patch_depth"])
+            pixel_depth = int(preset["pixel_depth"])
+            num_text_blocks = int(preset["num_text_blocks"])
+            patch_size = int(preset["patch_size"])
+            txt_embed_dim = int(preset["txt_embed_dim"])
+            txt_max_length = int(preset["txt_max_length"])
+            use_text_rope = bool(preset["use_text_rope"])
+            text_rope_theta = float(preset["text_rope_theta"])
+            repa_encoder_index = int(preset["repa_encoder_index"])
+
+        self.sample_size = int(sample_size)
+        self.in_channels = int(in_channels)
+        self.out_channels = int(in_channels)
+        self.hidden_size = int(hidden_size)
+        self.num_groups = int(num_groups)
+        self.patch_depth = int(patch_depth)
+        self.pixel_depth = int(pixel_depth)
+        self.num_text_blocks = int(num_text_blocks)
+        self.patch_size = int(patch_size)
+        self.pixel_hidden_size = int(pixel_hidden_size)
+        self.pixel_attn_hidden_size = int(pixel_attn_hidden_size)
+        self.pixel_num_groups = int(pixel_num_groups)
+        self.txt_embed_dim = int(txt_embed_dim)
+        self.txt_max_length = int(txt_max_length)
+        self.use_text_rope = bool(use_text_rope)
+        self.text_rope_theta = float(text_rope_theta)
+        self.repa_encoder_index = int(repa_encoder_index)
+        self.use_pixel_abs_pos = bool(use_pixel_abs_pos)
+        self.norm_eps = float(norm_eps)
+        self.gradient_checkpointing = False
+
+        if self.pixel_depth <= 0:
+            raise ValueError("PixelDiT T2I expects pixel_depth > 0 to preserve the pixel pathway")
+
+        self.pixel_embedder = PixelTokenEmbedder(
+            self.in_channels, self.pixel_hidden_size, use_pixel_abs_pos=self.use_pixel_abs_pos
+        )
+        self.s_embedder = PatchTokenEmbedder(self.in_channels * self.patch_size**2, self.hidden_size, bias=True)
+        self.t_embedder = TimestepConditioner(self.hidden_size)
+        self.y_embedder = PatchTokenEmbedder(
+            self.txt_embed_dim,
+            self.hidden_size,
+            bias=True,
+            norm_layer=lambda dim: RMSNorm(dim, eps=self.norm_eps),
+        )
+        self.y_pos_embedding = nn.Parameter(torch.randn(1, self.txt_max_length, self.hidden_size))
+
+        self.patch_blocks = nn.ModuleList(
+            [MMDiTBlockT2I(self.hidden_size, self.num_groups, eps=self.norm_eps) for _ in range(self.patch_depth)]
+        )
+        self.pixel_blocks = nn.ModuleList(
+            [
+                PiTBlock(
+                    self.pixel_hidden_size,
+                    self.hidden_size,
+                    patch_size=self.patch_size,
+                    num_heads=self.num_groups,
+                    mlp_ratio=4.0,
+                    attn_hidden_size=self.pixel_attn_hidden_size,
+                    attn_num_heads=self.pixel_num_groups,
+                    rope_fn=precompute_freqs_cis_2d,
+                    eps=self.norm_eps,
+                )
+                for _ in range(self.pixel_depth)
+            ]
+        )
+        self.final_layer = FinalLayer(self.pixel_hidden_size, self.out_channels, eps=self.norm_eps)
+        self._precompute_pos: Dict[tuple[int, int], torch.Tensor] = {}
+        self._precompute_pos_txt: Dict[int, torch.Tensor] = {}
+        self._initialize_weights()
+
+    def _fetch_pos(self, height: int, width: int, device: torch.device):
+        key = (height, width)
+        if key in self._precompute_pos:
+            return self._precompute_pos[key].to(device)
+        pos = precompute_freqs_cis_2d(self.hidden_size // self.num_groups, height, width).to(device)
+        self._precompute_pos[key] = pos
+        return pos
+
+    def _fetch_pos_text(self, length: int, device: torch.device):
+        if length in self._precompute_pos_txt:
+            return self._precompute_pos_txt[length].to(device)
+        head_dim = self.hidden_size // self.num_groups
+        freqs = 1.0 / (
+            self.text_rope_theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
+        )
+        positions = torch.arange(0, length, device=device).float().unsqueeze(1)
+        angles = positions * freqs.unsqueeze(0)
+        freqs_cis = torch.polar(torch.ones_like(angles), angles)
+        self._precompute_pos_txt[length] = freqs_cis
+        return freqs_cis
+
+    def _initialize_weights(self) -> None:
+        weight = self.s_embedder.proj.weight.data
+        nn.init.xavier_uniform_(weight.view([weight.shape[0], -1]))
+        nn.init.constant_(self.s_embedder.proj.bias, 0)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        nn.init.zeros_(self.final_layer.linear.weight)
+        nn.init.zeros_(self.final_layer.linear.bias)
+
+    def _build_joint_attn_mask(
+        self,
+        encoder_attention_mask: Optional[torch.Tensor],
+        batch_size: int,
+        num_img_tokens: int,
+        num_txt_tokens: int,
+        device: torch.device,
+    ) -> Optional[torch.Tensor]:
+        if encoder_attention_mask is None:
+            return None
+        mask = encoder_attention_mask
+        while mask.dim() > 2 and mask.size(1) == 1:
+            mask = mask.squeeze(1)
+        if mask.dim() == 3 and mask.size(1) == 1:
+            mask = mask.squeeze(1)
+        if mask.dim() != 2:
+            return None
+        pad = mask == 0
+        pad_img = torch.zeros((batch_size, num_img_tokens), dtype=torch.bool, device=device)
+        return torch.cat([pad[:, :num_txt_tokens], pad_img], dim=1).view(batch_size, 1, 1, num_txt_tokens + num_img_tokens)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float],
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]:
+        if sample.dim() != 4:
+            raise ValueError("PixelDiTT2ITransformer2DModel expects sample of shape [B,C,H,W]")
+        batch_size, _, height, width = sample.shape
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+
+        timestep = torch.as_tensor(timestep, device=sample.device)
+        if timestep.ndim == 0:
+            timestep = timestep.repeat(batch_size)
+        else:
+            timestep = timestep.reshape(-1)
+            if timestep.shape[0] == 1 and batch_size > 1:
+                timestep = timestep.repeat(batch_size)
+
+        if encoder_hidden_states.dim() == 4:
+            encoder_hidden_states = encoder_hidden_states.squeeze(1)
+        if encoder_hidden_states.dim() != 3:
+            raise ValueError("encoder_hidden_states must be [B, L, D]")
+
+        height_tokens = height // self.patch_size
+        width_tokens = width // self.patch_size
+        num_img_tokens = height_tokens * width_tokens
+
+        pos = self._fetch_pos(height_tokens, width_tokens, sample.device)
+        x_patches = F.unfold(sample, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2)
+        t_emb = self.t_embedder(timestep.view(-1)).view(batch_size, -1, self.hidden_size)
+
+        txt_length = min(encoder_hidden_states.shape[1], self.txt_max_length)
+        text_states = encoder_hidden_states[:, :txt_length, :]
+        text_states = self.y_embedder(text_states).view(batch_size, txt_length, self.hidden_size)
+        text_states = text_states + self.y_pos_embedding[:, :txt_length, :].to(text_states.dtype)
+        conditioning = F.silu(t_emb)
+
+        pos_txt = self._fetch_pos_text(txt_length, sample.device) if self.use_text_rope else None
+        attn_mask = self._build_joint_attn_mask(
+            encoder_attention_mask, batch_size, num_img_tokens, txt_length, sample.device
+        )
+
+        patch_states = self.s_embedder(x_patches)
+        for block in self.patch_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(img_states, txt_states, cond, position_img, position_txt, mask):
+                    return block(img_states, txt_states, cond, position_img, position_txt, mask)
+
+                patch_states, text_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward,
+                    patch_states,
+                    text_states,
+                    conditioning,
+                    pos,
+                    pos_txt,
+                    attn_mask,
+                    use_reentrant=False,
+                )
+            else:
+                patch_states, text_states = block(
+                    patch_states, text_states, conditioning, pos, pos_txt, attn_mask
+                )
+
+        patch_states = F.silu(t_emb + patch_states)
+        if patch_states.shape[1] != num_img_tokens:
+            if patch_states.shape[1] > num_img_tokens:
+                patch_states = patch_states[:, :num_img_tokens, :]
+            else:
+                pad_len = num_img_tokens - patch_states.shape[1]
+                patch_states = torch.cat(
+                    [patch_states, patch_states.new_zeros(batch_size, pad_len, patch_states.shape[2])], dim=1
+                )
+
+        conditioning_states = patch_states.reshape(batch_size * num_img_tokens, self.hidden_size)
+        pixel_states = self.pixel_embedder(sample, img_height=height, img_width=width, patch_size=self.patch_size)
+        for block in self.pixel_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(hidden_states, cond):
+                    return block(hidden_states, cond, height, width, self.patch_size)
+
+                pixel_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward, pixel_states, conditioning_states, use_reentrant=False
+                )
+            else:
+                pixel_states = block(pixel_states, conditioning_states, height, width, self.patch_size)
+
+        pixel_states = self.final_layer(pixel_states)
+        patch_area = self.patch_size * self.patch_size
+        pixel_states = pixel_states.view(batch_size, num_img_tokens, patch_area, self.out_channels).permute(0, 3, 2, 1)
+        pixel_states = pixel_states.contiguous().view(batch_size, self.out_channels * patch_area, num_img_tokens)
+        output = F.fold(pixel_states, (height, width), kernel_size=self.patch_size, stride=self.patch_size)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    @classmethod
+    def from_pixeldit_t2i_checkpoint(
+        cls,
+        checkpoint_path: str,
+        model_type: Literal["pixeldit-t2i-1300m"] = "pixeldit-t2i-1300m",
+        map_location: str = "cpu",
+        strict: bool = True,
+    ) -> Tuple["PixelDiTT2ITransformer2DModel", Dict[str, object]]:
+        if model_type not in PIXELDIT_T2I_PRESET_CONFIGS:
+            raise ValueError(f"Unknown PixelDiT T2I preset '{model_type}'.")
+
+        if checkpoint_path.endswith(".safetensors"):
+            try:
+                from safetensors.torch import load_file
+            except ImportError as error:
+                raise ImportError("Install safetensors to load .safetensors checkpoints.") from error
+            state_dict = load_file(checkpoint_path, device=map_location)
+        else:
+            loaded = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
+            if isinstance(loaded, Mapping):
+                state_dict = loaded
+                for key in ("state_dict", "model", "module", "denoiser"):
+                    if key in state_dict and isinstance(state_dict[key], dict):
+                        state_dict = state_dict[key]
+                        break
+            else:
+                raise ValueError("Unsupported checkpoint format.")
+
+        config = dict(PIXELDIT_T2I_PRESET_CONFIGS[model_type])
+        config["model_type"] = model_type
+        model = cls(**config)
+        model.load_state_dict(remap_t2i_legacy_state_dict(state_dict), strict=strict)
+        metadata = {"checkpoint_path": checkpoint_path, "model_type": model_type}
+        return model, metadata
+
+
+PixelDiTT2IDiffusersModel = PixelDiTT2ITransformer2DModel

PixelDiT-XL-16-256/model_index.json

@@ -0,0 +1,1017 @@
+{
+  "_class_name": [
+    "pipeline",
+    "PixelDiTPipeline"
+  ],
+  "_diffusers_version": "0.35.1",
+  "id2label": {
+    "0": "tench, Tinca tinca",
+    "1": "goldfish, Carassius auratus",
+    "10": "brambling, Fringilla montifringilla",
+    "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",
+    "11": "goldfinch, Carduelis carduelis",
+    "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",
+    "12": "house finch, linnet, Carpodacus mexicanus",
+    "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",
+    "13": "junco, snowbird",
+    "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",
+    "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+    "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",
+    "15": "robin, American robin, Turdus migratorius",
+    "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",
+    "16": "bulbul",
+    "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",
+    "17": "jay",
+    "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",
+    "18": "magpie",
+    "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",
+    "19": "chickadee",
+    "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",
+    "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
+    "20": "water ouzel, dipper",
+    "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",
+    "21": "kite",
+    "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",
+    "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
+    "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",
+    "23": "vulture",
+    "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",
+    "24": "great grey owl, great gray owl, Strix nebulosa",
+    "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",
+    "25": "European fire salamander, Salamandra salamandra",
+    "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",
+    "26": "common newt, Triturus vulgaris",
+    "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",
+    "27": "eft",
+    "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",
+    "28": "spotted salamander, Ambystoma maculatum",
+    "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",
+    "29": "axolotl, mud puppy, Ambystoma mexicanum",
+    "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",
+    "3": "tiger shark, Galeocerdo cuvieri",
+    "30": "bullfrog, Rana catesbeiana",
+    "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",
+    "31": "tree frog, tree-frog",
+    "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",
+    "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+    "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",
+    "33": "loggerhead, loggerhead turtle, Caretta caretta",
+    "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",
+    "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
+    "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",
+    "35": "mud turtle",
+    "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",
+    "36": "terrapin",
+    "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",
+    "37": "box turtle, box tortoise",
+    "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",
+    "38": "banded gecko",
+    "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",
+    "39": "common iguana, iguana, Iguana iguana",
+    "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",
+    "4": "hammerhead, hammerhead shark",
+    "40": "American chameleon, anole, Anolis carolinensis",
+    "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",
+    "41": "whiptail, whiptail lizard",
+    "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",
+    "42": "agama",
+    "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",
+    "43": "frilled lizard, Chlamydosaurus kingi",
+    "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",
+    "44": "alligator lizard",
+    "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",
+    "45": "Gila monster, Heloderma suspectum",
+    "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",
+    "46": "green lizard, Lacerta viridis",
+    "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",
+    "47": "African chameleon, Chamaeleo chamaeleon",
+    "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",
+    "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
+    "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",
+    "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
+    "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",
+    "5": "electric ray, crampfish, numbfish, torpedo",
+    "50": "American alligator, Alligator mississipiensis",
+    "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",
+    "51": "triceratops",
+    "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",
+    "52": "thunder snake, worm snake, Carphophis amoenus",
+    "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",
+    "53": "ringneck snake, ring-necked snake, ring snake",
+    "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",
+    "54": "hognose snake, puff adder, sand viper",
+    "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",
+    "55": "green snake, grass snake",
+    "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",
+    "56": "king snake, kingsnake",
+    "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",
+    "57": "garter snake, grass snake",
+    "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",
+    "58": "water snake",
+    "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",
+    "59": "vine snake",
+    "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",
+    "6": "stingray",
+    "60": "night snake, Hypsiglena torquata",
+    "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",
+    "61": "boa constrictor, Constrictor constrictor",
+    "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",
+    "62": "rock python, rock snake, Python sebae",
+    "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",
+    "63": "Indian cobra, Naja naja",
+    "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",
+    "64": "green mamba",
+    "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",
+    "65": "sea snake",
+    "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",
+    "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+    "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",
+    "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+    "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",
+    "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
+    "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",
+    "69": "trilobite",
+    "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",
+    "7": "cock",
+    "70": "harvestman, daddy longlegs, Phalangium opilio",
+    "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",
+    "71": "scorpion",
+    "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",
+    "72": "black and gold garden spider, Argiope aurantia",
+    "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",
+    "73": "barn spider, Araneus cavaticus",
+    "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",
+    "74": "garden spider, Aranea diademata",
+    "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",
+    "75": "black widow, Latrodectus mactans",
+    "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",
+    "76": "tarantula",
+    "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",
+    "77": "wolf spider, hunting spider",
+    "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",
+    "78": "tick",
+    "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",
+    "79": "centipede",
+    "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",
+    "8": "hen",
+    "80": "black grouse",
+    "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",
+    "81": "ptarmigan",
+    "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",
+    "82": "ruffed grouse, partridge, Bonasa umbellus",
+    "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",
+    "83": "prairie chicken, prairie grouse, prairie fowl",
+    "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",
+    "84": "peacock",
+    "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",
+    "85": "quail",
+    "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",
+    "86": "partridge",
+    "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",
+    "87": "African grey, African gray, Psittacus erithacus",
+    "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",
+    "88": "macaw",
+    "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",
+    "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+    "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",
+    "9": "ostrich, Struthio camelus",
+    "90": "lorikeet",
+    "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",
+    "91": "coucal",
+    "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",
+    "92": "bee eater",
+    "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",
+    "93": "hornbill",
+    "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",
+    "94": "hummingbird",
+    "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",
+    "95": "jacamar",
+    "950": "orange",
+    "951": "lemon",
+    "952": "fig",
+    "953": "pineapple, ananas",
+    "954": "banana",
+    "955": "jackfruit, jak, jack",
+    "956": "custard apple",
+    "957": "pomegranate",
+    "958": "hay",
+    "959": "carbonara",
+    "96": "toucan",
+    "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",
+    "97": "drake",
+    "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",
+    "98": "red-breasted merganser, Mergus serrator",
+    "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",
+    "99": "goose",
+    "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"
+  },
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "transformer": [
+    "transformer_pixeldit",
+    "PixelDiTTransformer2DModel"
+  ]
+}

PixelDiT-XL-16-256/pipeline.py

@@ -0,0 +1,286 @@
+"""Hub custom pipeline: PixelDiTPipeline.
+Load with native Hugging Face diffusers and trust_remote_code=True.
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils.torch_utils import randn_tensor
+
+RECOMMENDED_GUIDANCE_BY_SIZE = {
+    256: 3.25,
+    512: 3.75,
+}
+
+RECOMMENDED_SCHEDULER_SHIFT_BY_SIZE = {
+    256: 1.0,
+    512: 3.0,
+}
+
+
+class PixelDiTPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for image generation using PixelDiT (Pixel Diffusion Transformer).
+
+    Parameters:
+        transformer ([`PixelDiTTransformer2DModel`]):
+            A class-conditioned `PixelDiTTransformer2DModel` that predicts flow-matching velocity in pixel space.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            Diffusers scheduler interface for PixelDiT generation (defaults to deterministic flow-matching Euler).
+        id2label (`dict[int, str]`, *optional*):
+            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
+    """
+
+    @staticmethod
+    def prepare_extra_step_kwargs(
+        scheduler,
+        generator=None,
+        eta: float | None = None,
+    ):
+        kwargs = {}
+        step_params = set(inspect.signature(scheduler.step).parameters.keys())
+        if "generator" in step_params:
+            kwargs["generator"] = generator
+        if eta is not None and "eta" in step_params:
+            kwargs["eta"] = eta
+        return kwargs
+
+    model_cpu_offload_seq = "transformer"
+
+    def __init__(
+        self,
+        transformer,
+        scheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
+    ):
+        super().__init__()
+        sample_size = int(getattr(transformer.config, "sample_size", 256))
+        default_shift = RECOMMENDED_SCHEDULER_SHIFT_BY_SIZE.get(sample_size, 1.0)
+        scheduler = scheduler or FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=1000,
+            shift=default_shift,
+            stochastic_sampling=False,
+        )
+        self.register_modules(transformer=transformer, scheduler=scheduler)
+        self._id2label = self._normalize_id2label(id2label)
+        self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = bool(self._id2label)
+
+    def _ensure_labels_loaded(self) -> None:
+        if self._labels_loaded_from_model_index:
+            return
+        loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None))
+        if loaded:
+            self._id2label = loaded
+            self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = True
+
+    @staticmethod
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]:
+        if not variant_path:
+            return {}
+        variant_dir = Path(variant_path).resolve()
+        model_index_path = variant_dir / "model_index.json"
+        if not model_index_path.exists():
+            return {}
+        raw = json.loads(model_index_path.read_text(encoding="utf-8"))
+        id2label = raw.get("id2label")
+        if not isinstance(id2label, dict):
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]:
+        label2id: Dict[str, int] = {}
+        for class_id, value in id2label.items():
+            for synonym in value.split(","):
+                synonym = synonym.strip()
+                if synonym:
+                    label2id[synonym] = int(class_id)
+        return dict(sorted(label2id.items()))
+
+    @property
+    def id2label(self) -> Dict[int, str]:
+        self._ensure_labels_loaded()
+        return self._id2label
+
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
+        self._ensure_labels_loaded()
+        label2id = self.labels
+        if not label2id:
+            raise ValueError(
+                "No English labels loaded. Ensure `id2label` exists in model_index.json."
+            )
+
+        if isinstance(label, str):
+            label = [label]
+
+        missing = [item for item in label if item not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
+        return [label2id[item] for item in label]
+
+    def _normalize_class_labels(
+        self,
+        class_labels: Union[int, str, List[Union[int, str]]],
+    ) -> List[int]:
+        if isinstance(class_labels, int):
+            return [class_labels]
+
+        if isinstance(class_labels, str):
+            return self.get_label_ids(class_labels)
+
+        if class_labels and isinstance(class_labels[0], str):
+            return self.get_label_ids(class_labels)
+
+        return list(class_labels)
+
+    @staticmethod
+    def _resolve_timeshift(scheduler, image_size: int) -> float:
+        shift = getattr(scheduler.config, "shift", None)
+        if shift is not None:
+            return float(shift)
+        return RECOMMENDED_SCHEDULER_SHIFT_BY_SIZE.get(image_size, 1.0)
+
+    @staticmethod
+    def _build_flow_timesteps(
+        num_inference_steps: int,
+        timeshift: float,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        last_step = 1.0 / num_inference_steps if num_inference_steps > 1 else 1.0
+        timesteps = torch.linspace(0.0, 1.0 - last_step, num_inference_steps, device=device, dtype=dtype)
+        timesteps = torch.cat([timesteps, torch.ones(1, device=device, dtype=dtype)], dim=0)
+        if timeshift != 1.0:
+            timesteps = timesteps / (timesteps + (1.0 - timesteps) * timeshift)
+        return timesteps
+
+    @staticmethod
+    def _apply_classifier_free_guidance(model_output: torch.Tensor, guidance_scale: float) -> torch.Tensor:
+        model_output_uncond, model_output_cond = model_output.chunk(2, dim=0)
+        return model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond)
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        class_labels: Union[int, str, List[Union[int, str]]],
+        guidance_scale: Optional[float] = None,
+        guidance_interval_min: float = 0.1,
+        guidance_interval_max: float = 1.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        num_inference_steps: int = 100,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        if num_inference_steps < 1:
+            raise ValueError("num_inference_steps must be >= 1.")
+        if output_type not in {"pil", "np", "pt"}:
+            raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")
+
+        class_label_ids = self._normalize_class_labels(class_labels)
+        do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
+
+        batch_size = len(class_label_ids)
+        image_size = int(getattr(self.transformer.config, "sample_size", 256))
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
+        channels = int(self.transformer.config.in_channels)
+        null_class_val = int(
+            getattr(self.transformer.config, "num_classes", getattr(self.transformer.config, "num_class_embeds", 1000))
+        )
+
+        if guidance_scale is None:
+            guidance_scale = RECOMMENDED_GUIDANCE_BY_SIZE.get(image_size, 3.25)
+
+        latents = randn_tensor(
+            shape=(batch_size, channels, height, width),
+            generator=generator,
+            device=self._execution_device,
+            dtype=self.transformer.dtype,
+        )
+
+        class_labels_t = torch.tensor(class_label_ids, device=self._execution_device, dtype=torch.long).reshape(-1)
+        class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
+        class_null = torch.full_like(class_labels_t, null_class_val)
+
+        timeshift = self._resolve_timeshift(self.scheduler, image_size)
+        flow_timesteps = self._build_flow_timesteps(
+            num_inference_steps,
+            timeshift,
+            device=self._execution_device,
+            dtype=torch.float32,
+        )
+        velocity_dtype = self.transformer.dtype
+        v_prev = None
+
+        for t_cur, t_next in self.progress_bar(list(zip(flow_timesteps[:-1], flow_timesteps[1:]))):
+            dt = t_next - t_cur
+            flow_time = float(t_cur)
+            effective_guidance = (
+                guidance_scale
+                if do_classifier_free_guidance
+                and guidance_interval_min < flow_time < guidance_interval_max
+                else 1.0
+            )
+
+            latent_model_input = torch.cat([latents, latents], dim=0)
+            labels = torch.cat([class_null, class_labels_t], dim=0)
+            timesteps = torch.full(
+                (latent_model_input.shape[0],),
+                flow_time,
+                device=self._execution_device,
+                dtype=velocity_dtype,
+            )
+            model_output = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=labels,
+            ).sample
+            velocity = self._apply_classifier_free_guidance(model_output, effective_guidance)
+
+            if v_prev is None:
+                latents = latents + velocity * dt
+            else:
+                latents = latents + dt * (1.5 * velocity - 0.5 * v_prev)
+            v_prev = velocity
+
+        images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
+        if output_type == "pt":
+            images = images_pt
+        elif output_type == "np":
+            images = images_pt.permute(0, 2, 3, 1).numpy()
+        else:
+            images = self.numpy_to_pil(images_pt.permute(0, 2, 3, 1).numpy())
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)
+
+
+PixelDiTPipelineOutput = ImagePipelineOutput

PixelDiT-XL-16-256/scheduler/config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.36.0",
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "stochastic_sampling": false
+}

PixelDiT-XL-16-256/scheduler/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.36.0",
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "stochastic_sampling": false
+}

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

@@ -0,0 +1,14 @@
+{
+  "_class_name": "PixelDiTTransformer2DModel",
+  "hidden_size": 1152,
+  "in_channels": 3,
+  "model_type": "pixeldit-xl",
+  "num_classes": 1000,
+  "num_groups": 16,
+  "patch_depth": 26,
+  "patch_size": 16,
+  "pixel_depth": 4,
+  "pixel_hidden_size": 16,
+  "sample_size": 256,
+  "use_pixel_abs_pos": true
+}

PixelDiT-XL-16-256/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07ea7a0cc50eb0318a2a100fc33d184c40153910e329398220eac4b9a632c3dd
+size 3189574228

PixelDiT-XL-16-256/transformer/transformer_pixeldit.py

@@ -0,0 +1,639 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from typing import Dict, Literal, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import RMSNorm
+
+
+PIXELDIT_PRESET_CONFIGS: Dict[str, Dict[str, object]] = {
+    "pixeldit-xl": {
+        "sample_size": 256,
+        "num_groups": 16,
+        "hidden_size": 1152,
+        "pixel_hidden_size": 16,
+        "patch_depth": 26,
+        "pixel_depth": 4,
+        "patch_size": 16,
+    },
+}
+
+
+def remap_legacy_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Map wrapper/backbone keys from legacy checkpoints to native PixelDiTTransformer2DModel keys."""
+    remapped: Dict[str, torch.Tensor] = {}
+    prefixes = ("transformer.", "model.", "module.", "denoiser.", "net.")
+    for key, value in state_dict.items():
+        new_key = key
+        for prefix in prefixes:
+            if new_key.startswith(prefix):
+                new_key = new_key[len(prefix) :]
+                break
+        remapped[new_key] = value
+    return remapped
+
+
+def config_from_legacy(config: Dict[str, object]) -> Dict[str, object]:
+    """Build native config kwargs from a legacy config.json dict."""
+    model_type = config.get("model_type") or config.get("model_name") or config.get("model_size")
+    if model_type not in PIXELDIT_PRESET_CONFIGS:
+        raise ValueError(f"Unknown PixelDiT preset '{model_type}'. Known: {list(PIXELDIT_PRESET_CONFIGS)}")
+
+    preset = dict(PIXELDIT_PRESET_CONFIGS[model_type])
+    preset["num_classes"] = int(config.get("num_classes") or config.get("num_class_embeds") or 1000)
+    preset["in_channels"] = int(config.get("in_channels", 3))
+    preset["use_pixel_abs_pos"] = bool(config.get("use_pixel_abs_pos", True))
+    preset["model_type"] = model_type
+
+    for key in ("sample_size", "num_groups", "hidden_size", "pixel_hidden_size", "patch_depth", "pixel_depth", "patch_size"):
+        if config.get(key) is not None:
+            preset[key] = config[key]
+
+    return preset
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int) -> np.ndarray:
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0)
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    return get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
+    if embed_dim % 2 != 0:
+        raise ValueError("Embedding dimension must be even for 2D sin/cos positional embeddings.")
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    if embed_dim % 2 != 0:
+        raise ValueError("Embedding dimension must be even for 1D sin/cos positional embeddings.")
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+def apply_adaln(hidden_states: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return hidden_states * (1 + scale) + shift
+
+
+def precompute_freqs_cis_2d(dim: int, height: int, width: int, theta: float = 10000.0, scale: float = 16.0):
+    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))
+    x_freqs = torch.outer(x_pos, freqs).float()
+    y_freqs = torch.outer(y_pos, freqs).float()
+    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)
+    return freqs_cis.reshape(height * width, -1)
+
+
+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_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 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)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class TimestepConditioner(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(timesteps: torch.Tensor, dim: int, max_period: int = 10):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half
+        )
+        args = timesteps[..., 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, timesteps: torch.Tensor) -> torch.Tensor:
+        timestep_freq = self.timestep_embedding(timesteps, self.frequency_embedding_size)
+        mlp_dtype = next(self.mlp.parameters()).dtype
+        if timestep_freq.dtype != mlp_dtype:
+            timestep_freq = timestep_freq.to(mlp_dtype)
+        return self.mlp(timestep_freq)
+
+
+class ClassEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int):
+        super().__init__()
+        self.embedding_table = nn.Embedding(num_classes, hidden_size)
+        self.num_classes = num_classes
+
+    def forward(self, labels: torch.Tensor) -> torch.Tensor:
+        return self.embedding_table(labels)
+
+
+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, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(hidden_states)) * self.w3(hidden_states))
+
+
+class RotaryAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError("dim should be divisible by num_heads")
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, eps=eps) 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, hidden_states: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        batch_size, length, channels = hidden_states.shape
+        qkv = (
+            self.qkv(hidden_states)
+            .reshape(batch_size, length, 3, self.num_heads, channels // self.num_heads)
+            .permute(2, 0, 1, 3, 4)
+        )
+        query, key, value = qkv[0], qkv[1], qkv[2]
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        query, key = apply_rotary_emb(query, key, freqs_cis=pos)
+        query = query.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2)
+        key = key.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2).contiguous()
+        value = value.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2).contiguous()
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, length, channels)
+        hidden_states = self.proj(hidden_states)
+        return self.proj_drop(hidden_states)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int, mlp_ratio: float = 4.0, drop: float = 0.0):
+        super().__init__()
+        hidden_dim = int(dim * mlp_ratio)
+        self.fc1 = nn.Linear(dim, hidden_dim)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_dim, dim)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.drop(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return self.drop(hidden_states)
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, out_channels: int, eps: float = 1e-6):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size, eps=eps)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(hidden_states)
+        return self.linear(hidden_states)
+
+
+class PatchTokenEmbedder(nn.Module):
+    def __init__(self, in_chans: int, embed_dim: int, 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, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.proj(hidden_states)
+        return self.norm(hidden_states)
+
+
+class PixelTokenEmbedder(nn.Module):
+    def __init__(self, in_channels: int, hidden_size_output: int, use_pixel_abs_pos: bool = True):
+        super().__init__()
+        self.in_channels = int(in_channels)
+        self.hidden_size_output = int(hidden_size_output)
+        self.use_pixel_abs_pos = bool(use_pixel_abs_pos)
+        self.proj = nn.Linear(self.in_channels, self.hidden_size_output, bias=True)
+        self._pos_cache: Dict[tuple[str, int, int], torch.Tensor] = {}
+
+    def _fetch_pixel_pos_image(self, height: int, width: int, device: torch.device, dtype: torch.dtype):
+        key = ("image", height, width)
+        if key in self._pos_cache:
+            return self._pos_cache[key].to(device=device, dtype=dtype)
+        if height == width:
+            pos_np = get_2d_sincos_pos_embed(self.hidden_size_output, height)
+        else:
+            grid_h = np.arange(height, dtype=np.float32)
+            grid_w = np.arange(width, dtype=np.float32)
+            grid = np.meshgrid(grid_w, grid_h)
+            grid = np.stack(grid, axis=0).reshape(2, 1, height, width)
+            pos_np = get_2d_sincos_pos_embed_from_grid(self.hidden_size_output, grid)
+        pos = torch.from_numpy(pos_np).to(device=device, dtype=dtype)
+        self._pos_cache[key] = pos
+        return pos
+
+    def forward(self, inputs: torch.Tensor, img_height: int, img_width: int, patch_size: int):
+        if inputs.dim() != 4:
+            raise ValueError("PixelTokenEmbedder expects inputs of shape [B,C,H,W]")
+        batch_size, channels, height, width = inputs.shape
+        if height != img_height or width != img_width:
+            raise ValueError("Input resolution does not match img_height/img_width.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+        h_tokens, w_tokens = height // patch_size, width // patch_size
+        patch_area = patch_size * patch_size
+        hidden_states = inputs.permute(0, 2, 3, 1).contiguous()
+        hidden_states = self.proj(hidden_states)
+        if self.use_pixel_abs_pos:
+            pos_full = self._fetch_pixel_pos_image(height, width, inputs.device, inputs.dtype)
+            hidden_states = hidden_states + pos_full.view(height, width, self.hidden_size_output).unsqueeze(0)
+        hidden_states = hidden_states.view(batch_size, h_tokens, patch_size, w_tokens, patch_size, self.hidden_size_output)
+        hidden_states = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous()
+        return hidden_states.view(batch_size * h_tokens * w_tokens, patch_area, self.hidden_size_output)
+
+
+class AugmentedDiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, groups: int, mlp_ratio: float = 4.0, adaLN_modulation=None, eps: float = 1e-6):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size, eps=eps)
+        self.attn = RotaryAttention(hidden_size, num_heads=groups, qkv_bias=False, eps=eps)
+        self.norm2 = RMSNorm(hidden_size, eps=eps)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp = FeedForward(hidden_size, mlp_hidden_dim)
+        self.adaLN_modulation = adaLN_modulation if adaLN_modulation is not None else nn.Sequential(
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, hidden_states: torch.Tensor, conditioning: torch.Tensor, pos: torch.Tensor):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(conditioning).chunk(
+            6, dim=-1
+        )
+        hidden_states = hidden_states + gate_msa * self.attn(
+            apply_adaln(self.norm1(hidden_states), shift_msa, scale_msa), pos
+        )
+        hidden_states = hidden_states + gate_mlp * self.mlp(
+            apply_adaln(self.norm2(hidden_states), shift_mlp, scale_mlp)
+        )
+        return hidden_states
+
+
+class PiTBlock(nn.Module):
+    def __init__(
+        self,
+        pixel_hidden_size: int,
+        patch_hidden_size: int,
+        patch_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_hidden_size: Optional[int] = None,
+        attn_num_heads: Optional[int] = None,
+        rope_fn=None,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.pixel_dim = int(pixel_hidden_size)
+        self.context_dim = int(patch_hidden_size)
+        self.patch_size = int(patch_size)
+        self.attn_dim = int(attn_hidden_size) if attn_hidden_size is not None else self.context_dim
+        self.num_heads = int(attn_num_heads) if attn_num_heads is not None else int(num_heads)
+        if self.attn_dim % self.num_heads != 0:
+            raise ValueError("pixel attention hidden size must be divisible by pixel num_heads")
+        patch_area = self.patch_size * self.patch_size
+        self.compress_to_attn = nn.Linear(patch_area * self.pixel_dim, self.attn_dim, bias=True)
+        self.expand_from_attn = nn.Linear(self.attn_dim, patch_area * self.pixel_dim, bias=True)
+        self.norm1 = RMSNorm(self.pixel_dim, eps=eps)
+        self.attn = RotaryAttention(self.attn_dim, num_heads=self.num_heads, qkv_bias=False, eps=eps)
+        self.norm2 = RMSNorm(self.pixel_dim, eps=eps)
+        self.mlp = MLP(self.pixel_dim, mlp_ratio=mlp_ratio, drop=0.0)
+        self.adaLN_modulation = nn.Sequential(nn.Linear(self.context_dim, 6 * self.pixel_dim * patch_area, bias=True))
+        self._pos_cache: Dict[tuple[int, int], torch.Tensor] = {}
+        self._rope_fn = rope_fn if rope_fn is not None else precompute_freqs_cis_2d
+
+    def _fetch_pos(self, height: int, width: int, device: torch.device):
+        key = (height, width)
+        if key in self._pos_cache:
+            return self._pos_cache[key].to(device)
+        pos = self._rope_fn(self.attn_dim // self.num_heads, height, width).to(device)
+        self._pos_cache[key] = pos
+        return pos
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        conditioning: torch.Tensor,
+        image_height: int,
+        image_width: int,
+        patch_size: int,
+    ) -> torch.Tensor:
+        batch_tokens, patch_area, channels = hidden_states.shape
+        if channels != self.pixel_dim:
+            raise ValueError(f"PiTBlock expected pixel_dim={self.pixel_dim}, got {channels}")
+        if image_height % patch_size != 0 or image_width % patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+        h_tokens, w_tokens = image_height // patch_size, image_width // patch_size
+        length = h_tokens * w_tokens
+        batch_size = batch_tokens // length
+        cond_params = self.adaLN_modulation(conditioning).view(batch_tokens, patch_area, 6 * self.pixel_dim)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(cond_params, 6, dim=-1)
+        hidden_norm = apply_adaln(self.norm1(hidden_states), shift_msa, scale_msa)
+        hidden_flat = hidden_norm.view(batch_tokens, patch_area * self.pixel_dim)
+        hidden_comp = self.compress_to_attn(hidden_flat).view(batch_size, length, self.attn_dim)
+        pos_comp = self._fetch_pos(h_tokens, w_tokens, hidden_states.device)
+        attn_out = self.attn(hidden_comp, pos_comp)
+        attn_flat = self.expand_from_attn(attn_out.view(batch_size * length, self.attn_dim))
+        attn_exp = attn_flat.view(batch_tokens, patch_area, self.pixel_dim)
+        hidden_states = hidden_states + gate_msa * attn_exp
+        mlp_out = self.mlp(apply_adaln(self.norm2(hidden_states), shift_mlp, scale_mlp))
+        hidden_states = hidden_states + gate_mlp * mlp_out
+        return hidden_states
+
+
+class PixelDiTTransformer2DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+    _skip_layerwise_casting_patterns = ["pos", "_pos_cache"]
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 256,
+        in_channels: int = 3,
+        num_groups: int = 16,
+        hidden_size: int = 1152,
+        pixel_hidden_size: int = 16,
+        patch_depth: int = 26,
+        pixel_depth: int = 4,
+        patch_size: int = 16,
+        num_classes: int = 1000,
+        use_pixel_abs_pos: bool = True,
+        norm_eps: float = 1e-6,
+        model_type: str | None = None,
+        num_class_embeds: int | None = None,
+    ):
+        super().__init__()
+        if num_class_embeds is not None:
+            num_classes = int(num_class_embeds)
+        if model_type in PIXELDIT_PRESET_CONFIGS:
+            preset = PIXELDIT_PRESET_CONFIGS[model_type]
+            sample_size = int(preset["sample_size"])
+            num_groups = int(preset["num_groups"])
+            hidden_size = int(preset["hidden_size"])
+            pixel_hidden_size = int(preset["pixel_hidden_size"])
+            patch_depth = int(preset["patch_depth"])
+            pixel_depth = int(preset["pixel_depth"])
+            patch_size = int(preset["patch_size"])
+
+        self.sample_size = int(sample_size)
+        self.in_channels = int(in_channels)
+        self.out_channels = int(in_channels)
+        self.hidden_size = int(hidden_size)
+        self.num_groups = int(num_groups)
+        self.patch_depth = int(patch_depth)
+        self.pixel_depth = int(pixel_depth)
+        self.patch_size = int(patch_size)
+        self.pixel_hidden_size = int(pixel_hidden_size)
+        self.num_classes = int(num_classes)
+        self.use_pixel_abs_pos = bool(use_pixel_abs_pos)
+        self.norm_eps = float(norm_eps)
+        self.gradient_checkpointing = False
+
+        if self.pixel_depth <= 0:
+            raise ValueError("PixelDiT expects pixel_depth > 0 to preserve the dual-level pipeline")
+
+        self.pixel_embedder = PixelTokenEmbedder(
+            self.in_channels, self.pixel_hidden_size, use_pixel_abs_pos=self.use_pixel_abs_pos
+        )
+        self.s_embedder = PatchTokenEmbedder(self.in_channels * self.patch_size**2, self.hidden_size, bias=True)
+        self.t_embedder = TimestepConditioner(self.hidden_size)
+        self.y_embedder = ClassEmbedder(self.num_classes + 1, self.hidden_size)
+
+        self.final_layer = FinalLayer(self.pixel_hidden_size, self.out_channels, eps=self.norm_eps)
+        self.patch_blocks = nn.ModuleList(
+            [AugmentedDiTBlock(self.hidden_size, self.num_groups, eps=self.norm_eps) for _ in range(self.patch_depth)]
+        )
+        self.pixel_blocks = nn.ModuleList(
+            [
+                PiTBlock(
+                    self.pixel_hidden_size,
+                    self.hidden_size,
+                    patch_size=self.patch_size,
+                    num_heads=self.num_groups,
+                    mlp_ratio=4.0,
+                    eps=self.norm_eps,
+                )
+                for _ in range(self.pixel_depth)
+            ]
+        )
+        self._precompute_pos: Dict[tuple[int, int], torch.Tensor] = {}
+        self._initialize_weights()
+
+    def _fetch_pos(self, height: int, width: int, device: torch.device):
+        key = (height, width)
+        if key in self._precompute_pos:
+            return self._precompute_pos[key].to(device)
+        pos = precompute_freqs_cis_2d(self.hidden_size // self.num_groups, height, width).to(device)
+        self._precompute_pos[key] = pos
+        return pos
+
+    def _initialize_weights(self) -> None:
+        weight = self.s_embedder.proj.weight.data
+        nn.init.xavier_uniform_(weight.view([weight.shape[0], -1]))
+        nn.init.constant_(self.s_embedder.proj.bias, 0)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        nn.init.zeros_(self.final_layer.linear.weight)
+        nn.init.zeros_(self.final_layer.linear.bias)
+        for block in self.patch_blocks:
+            nn.init.zeros_(block.adaLN_modulation[0].weight)
+            nn.init.zeros_(block.adaLN_modulation[0].bias)
+        for block in self.pixel_blocks:
+            nn.init.zeros_(block.adaLN_modulation[0].weight)
+            nn.init.zeros_(block.adaLN_modulation[0].bias)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float],
+        class_labels: Union[torch.Tensor, int],
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]:
+        if sample.dim() != 4:
+            raise ValueError("PixelDiTTransformer2DModel expects sample of shape [B,C,H,W]")
+        batch_size, _, height, width = sample.shape
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+
+        timestep = torch.as_tensor(timestep, device=sample.device)
+        if timestep.ndim == 0:
+            timestep = timestep.repeat(batch_size)
+        else:
+            timestep = timestep.reshape(-1)
+            if timestep.shape[0] == 1 and batch_size > 1:
+                timestep = timestep.repeat(batch_size)
+
+        if not torch.is_tensor(class_labels):
+            class_labels = torch.tensor(class_labels, device=sample.device, dtype=torch.long)
+        class_labels = class_labels.to(device=sample.device, dtype=torch.long).reshape(-1)
+        if class_labels.shape[0] == 1 and batch_size > 1:
+            class_labels = class_labels.repeat(batch_size)
+
+        pos = self._fetch_pos(height // self.patch_size, width // self.patch_size, sample.device)
+        x_patches = F.unfold(sample, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2)
+        t_emb = self.t_embedder(timestep.view(-1)).view(batch_size, -1, self.hidden_size)
+        y_emb = self.y_embedder(class_labels).view(batch_size, 1, self.hidden_size)
+        conditioning = F.silu(t_emb + y_emb)
+
+        patch_states = self.s_embedder(x_patches)
+        for block in self.patch_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(hidden_states, cond, position):
+                    return block(hidden_states, cond, position)
+
+                patch_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward, patch_states, conditioning, pos, use_reentrant=False
+                )
+            else:
+                patch_states = block(patch_states, conditioning, pos)
+        patch_states = F.silu(t_emb + patch_states)
+
+        length = patch_states.shape[1]
+        conditioning_states = patch_states.view(batch_size * length, self.hidden_size)
+        pixel_states = self.pixel_embedder(
+            sample, img_height=height, img_width=width, patch_size=self.patch_size
+        )
+        for block in self.pixel_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(hidden_states, cond):
+                    return block(hidden_states, cond, height, width, self.patch_size)
+
+                pixel_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward, pixel_states, conditioning_states, use_reentrant=False
+                )
+            else:
+                pixel_states = block(pixel_states, conditioning_states, height, width, self.patch_size)
+        pixel_states = self.final_layer(pixel_states)
+
+        patch_area = self.patch_size * self.patch_size
+        pixel_states = pixel_states.view(batch_size, length, patch_area, self.out_channels).permute(0, 3, 2, 1)
+        pixel_states = pixel_states.contiguous().view(batch_size, self.out_channels * patch_area, length)
+        output = F.fold(pixel_states, (height, width), kernel_size=self.patch_size, stride=self.patch_size)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    @classmethod
+    def from_pixeldit_checkpoint(
+        cls,
+        checkpoint_path: str,
+        model_type: Literal["pixeldit-xl"] = "pixeldit-xl",
+        map_location: str = "cpu",
+        strict: bool = True,
+    ) -> Tuple["PixelDiTTransformer2DModel", Dict[str, object]]:
+        if model_type not in PIXELDIT_PRESET_CONFIGS:
+            raise ValueError(f"Unknown PixelDiT preset '{model_type}'.")
+
+        if checkpoint_path.endswith(".safetensors"):
+            try:
+                from safetensors.torch import load_file
+            except ImportError as error:
+                raise ImportError("Install safetensors to load .safetensors checkpoints.") from error
+            state_dict = load_file(checkpoint_path, device=map_location)
+        else:
+            loaded = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
+            if isinstance(loaded, Mapping):
+                state_dict = loaded
+                for key in ("state_dict", "model", "module", "denoiser"):
+                    if key in state_dict and isinstance(state_dict[key], dict):
+                        state_dict = state_dict[key]
+                        break
+            else:
+                raise ValueError("Unsupported checkpoint format.")
+
+        config = dict(PIXELDIT_PRESET_CONFIGS[model_type])
+        config["model_type"] = model_type
+        model = cls(**config)
+        model.load_state_dict(remap_legacy_state_dict(state_dict), strict=strict)
+
+        metadata = {
+            "checkpoint_path": checkpoint_path,
+            "model_type": model_type,
+        }
+        return model, metadata
+
+    def to_pixeldit_checkpoint(self, prefix: str = "") -> Dict[str, torch.Tensor]:
+        checkpoint: Dict[str, torch.Tensor] = {}
+        for key, value in self.state_dict().items():
+            checkpoint[f"{prefix}{key}"] = value.detach().cpu()
+        return checkpoint
+
+
+PixelDiTDiffusersModel = PixelDiTTransformer2DModel

PixelDiT-XL-16-512/model_index.json

@@ -0,0 +1,1017 @@
+{
+  "_class_name": [
+    "pipeline",
+    "PixelDiTPipeline"
+  ],
+  "_diffusers_version": "0.35.1",
+  "id2label": {
+    "0": "tench, Tinca tinca",
+    "1": "goldfish, Carassius auratus",
+    "10": "brambling, Fringilla montifringilla",
+    "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",
+    "11": "goldfinch, Carduelis carduelis",
+    "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",
+    "12": "house finch, linnet, Carpodacus mexicanus",
+    "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",
+    "13": "junco, snowbird",
+    "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",
+    "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+    "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",
+    "15": "robin, American robin, Turdus migratorius",
+    "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",
+    "16": "bulbul",
+    "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",
+    "17": "jay",
+    "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",
+    "18": "magpie",
+    "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",
+    "19": "chickadee",
+    "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",
+    "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
+    "20": "water ouzel, dipper",
+    "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",
+    "21": "kite",
+    "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",
+    "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
+    "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",
+    "23": "vulture",
+    "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",
+    "24": "great grey owl, great gray owl, Strix nebulosa",
+    "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",
+    "25": "European fire salamander, Salamandra salamandra",
+    "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",
+    "26": "common newt, Triturus vulgaris",
+    "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",
+    "27": "eft",
+    "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",
+    "28": "spotted salamander, Ambystoma maculatum",
+    "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",
+    "29": "axolotl, mud puppy, Ambystoma mexicanum",
+    "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",
+    "3": "tiger shark, Galeocerdo cuvieri",
+    "30": "bullfrog, Rana catesbeiana",
+    "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",
+    "31": "tree frog, tree-frog",
+    "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",
+    "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+    "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",
+    "33": "loggerhead, loggerhead turtle, Caretta caretta",
+    "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",
+    "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
+    "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",
+    "35": "mud turtle",
+    "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",
+    "36": "terrapin",
+    "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",
+    "37": "box turtle, box tortoise",
+    "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",
+    "38": "banded gecko",
+    "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",
+    "39": "common iguana, iguana, Iguana iguana",
+    "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",
+    "4": "hammerhead, hammerhead shark",
+    "40": "American chameleon, anole, Anolis carolinensis",
+    "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",
+    "41": "whiptail, whiptail lizard",
+    "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",
+    "42": "agama",
+    "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",
+    "43": "frilled lizard, Chlamydosaurus kingi",
+    "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",
+    "44": "alligator lizard",
+    "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",
+    "45": "Gila monster, Heloderma suspectum",
+    "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",
+    "46": "green lizard, Lacerta viridis",
+    "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",
+    "47": "African chameleon, Chamaeleo chamaeleon",
+    "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",
+    "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
+    "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",
+    "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
+    "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",
+    "5": "electric ray, crampfish, numbfish, torpedo",
+    "50": "American alligator, Alligator mississipiensis",
+    "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",
+    "51": "triceratops",
+    "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",
+    "52": "thunder snake, worm snake, Carphophis amoenus",
+    "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",
+    "53": "ringneck snake, ring-necked snake, ring snake",
+    "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",
+    "54": "hognose snake, puff adder, sand viper",
+    "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",
+    "55": "green snake, grass snake",
+    "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",
+    "56": "king snake, kingsnake",
+    "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",
+    "57": "garter snake, grass snake",
+    "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",
+    "58": "water snake",
+    "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",
+    "59": "vine snake",
+    "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",
+    "6": "stingray",
+    "60": "night snake, Hypsiglena torquata",
+    "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",
+    "61": "boa constrictor, Constrictor constrictor",
+    "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",
+    "62": "rock python, rock snake, Python sebae",
+    "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",
+    "63": "Indian cobra, Naja naja",
+    "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",
+    "64": "green mamba",
+    "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",
+    "65": "sea snake",
+    "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",
+    "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+    "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",
+    "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+    "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",
+    "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
+    "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",
+    "69": "trilobite",
+    "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",
+    "7": "cock",
+    "70": "harvestman, daddy longlegs, Phalangium opilio",
+    "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",
+    "71": "scorpion",
+    "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",
+    "72": "black and gold garden spider, Argiope aurantia",
+    "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",
+    "73": "barn spider, Araneus cavaticus",
+    "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",
+    "74": "garden spider, Aranea diademata",
+    "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",
+    "75": "black widow, Latrodectus mactans",
+    "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",
+    "76": "tarantula",
+    "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",
+    "77": "wolf spider, hunting spider",
+    "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",
+    "78": "tick",
+    "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",
+    "79": "centipede",
+    "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",
+    "8": "hen",
+    "80": "black grouse",
+    "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",
+    "81": "ptarmigan",
+    "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",
+    "82": "ruffed grouse, partridge, Bonasa umbellus",
+    "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",
+    "83": "prairie chicken, prairie grouse, prairie fowl",
+    "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",
+    "84": "peacock",
+    "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",
+    "85": "quail",
+    "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",
+    "86": "partridge",
+    "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",
+    "87": "African grey, African gray, Psittacus erithacus",
+    "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",
+    "88": "macaw",
+    "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",
+    "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+    "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",
+    "9": "ostrich, Struthio camelus",
+    "90": "lorikeet",
+    "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",
+    "91": "coucal",
+    "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",
+    "92": "bee eater",
+    "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",
+    "93": "hornbill",
+    "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",
+    "94": "hummingbird",
+    "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",
+    "95": "jacamar",
+    "950": "orange",
+    "951": "lemon",
+    "952": "fig",
+    "953": "pineapple, ananas",
+    "954": "banana",
+    "955": "jackfruit, jak, jack",
+    "956": "custard apple",
+    "957": "pomegranate",
+    "958": "hay",
+    "959": "carbonara",
+    "96": "toucan",
+    "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",
+    "97": "drake",
+    "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",
+    "98": "red-breasted merganser, Mergus serrator",
+    "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",
+    "99": "goose",
+    "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"
+  },
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "transformer": [
+    "transformer_pixeldit",
+    "PixelDiTTransformer2DModel"
+  ]
+}

PixelDiT-XL-16-512/pipeline.py

@@ -0,0 +1,286 @@
+"""Hub custom pipeline: PixelDiTPipeline.
+Load with native Hugging Face diffusers and trust_remote_code=True.
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils.torch_utils import randn_tensor
+
+RECOMMENDED_GUIDANCE_BY_SIZE = {
+    256: 3.25,
+    512: 3.75,
+}
+
+RECOMMENDED_SCHEDULER_SHIFT_BY_SIZE = {
+    256: 1.0,
+    512: 3.0,
+}
+
+
+class PixelDiTPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for image generation using PixelDiT (Pixel Diffusion Transformer).
+
+    Parameters:
+        transformer ([`PixelDiTTransformer2DModel`]):
+            A class-conditioned `PixelDiTTransformer2DModel` that predicts flow-matching velocity in pixel space.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            Diffusers scheduler interface for PixelDiT generation (defaults to deterministic flow-matching Euler).
+        id2label (`dict[int, str]`, *optional*):
+            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
+    """
+
+    @staticmethod
+    def prepare_extra_step_kwargs(
+        scheduler,
+        generator=None,
+        eta: float | None = None,
+    ):
+        kwargs = {}
+        step_params = set(inspect.signature(scheduler.step).parameters.keys())
+        if "generator" in step_params:
+            kwargs["generator"] = generator
+        if eta is not None and "eta" in step_params:
+            kwargs["eta"] = eta
+        return kwargs
+
+    model_cpu_offload_seq = "transformer"
+
+    def __init__(
+        self,
+        transformer,
+        scheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
+    ):
+        super().__init__()
+        sample_size = int(getattr(transformer.config, "sample_size", 256))
+        default_shift = RECOMMENDED_SCHEDULER_SHIFT_BY_SIZE.get(sample_size, 1.0)
+        scheduler = scheduler or FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=1000,
+            shift=default_shift,
+            stochastic_sampling=False,
+        )
+        self.register_modules(transformer=transformer, scheduler=scheduler)
+        self._id2label = self._normalize_id2label(id2label)
+        self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = bool(self._id2label)
+
+    def _ensure_labels_loaded(self) -> None:
+        if self._labels_loaded_from_model_index:
+            return
+        loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None))
+        if loaded:
+            self._id2label = loaded
+            self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = True
+
+    @staticmethod
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]:
+        if not variant_path:
+            return {}
+        variant_dir = Path(variant_path).resolve()
+        model_index_path = variant_dir / "model_index.json"
+        if not model_index_path.exists():
+            return {}
+        raw = json.loads(model_index_path.read_text(encoding="utf-8"))
+        id2label = raw.get("id2label")
+        if not isinstance(id2label, dict):
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]:
+        label2id: Dict[str, int] = {}
+        for class_id, value in id2label.items():
+            for synonym in value.split(","):
+                synonym = synonym.strip()
+                if synonym:
+                    label2id[synonym] = int(class_id)
+        return dict(sorted(label2id.items()))
+
+    @property
+    def id2label(self) -> Dict[int, str]:
+        self._ensure_labels_loaded()
+        return self._id2label
+
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
+        self._ensure_labels_loaded()
+        label2id = self.labels
+        if not label2id:
+            raise ValueError(
+                "No English labels loaded. Ensure `id2label` exists in model_index.json."
+            )
+
+        if isinstance(label, str):
+            label = [label]
+
+        missing = [item for item in label if item not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
+        return [label2id[item] for item in label]
+
+    def _normalize_class_labels(
+        self,
+        class_labels: Union[int, str, List[Union[int, str]]],
+    ) -> List[int]:
+        if isinstance(class_labels, int):
+            return [class_labels]
+
+        if isinstance(class_labels, str):
+            return self.get_label_ids(class_labels)
+
+        if class_labels and isinstance(class_labels[0], str):
+            return self.get_label_ids(class_labels)
+
+        return list(class_labels)
+
+    @staticmethod
+    def _resolve_timeshift(scheduler, image_size: int) -> float:
+        shift = getattr(scheduler.config, "shift", None)
+        if shift is not None:
+            return float(shift)
+        return RECOMMENDED_SCHEDULER_SHIFT_BY_SIZE.get(image_size, 1.0)
+
+    @staticmethod
+    def _build_flow_timesteps(
+        num_inference_steps: int,
+        timeshift: float,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        last_step = 1.0 / num_inference_steps if num_inference_steps > 1 else 1.0
+        timesteps = torch.linspace(0.0, 1.0 - last_step, num_inference_steps, device=device, dtype=dtype)
+        timesteps = torch.cat([timesteps, torch.ones(1, device=device, dtype=dtype)], dim=0)
+        if timeshift != 1.0:
+            timesteps = timesteps / (timesteps + (1.0 - timesteps) * timeshift)
+        return timesteps
+
+    @staticmethod
+    def _apply_classifier_free_guidance(model_output: torch.Tensor, guidance_scale: float) -> torch.Tensor:
+        model_output_uncond, model_output_cond = model_output.chunk(2, dim=0)
+        return model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond)
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        class_labels: Union[int, str, List[Union[int, str]]],
+        guidance_scale: Optional[float] = None,
+        guidance_interval_min: float = 0.1,
+        guidance_interval_max: float = 1.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        num_inference_steps: int = 100,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        if num_inference_steps < 1:
+            raise ValueError("num_inference_steps must be >= 1.")
+        if output_type not in {"pil", "np", "pt"}:
+            raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")
+
+        class_label_ids = self._normalize_class_labels(class_labels)
+        do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
+
+        batch_size = len(class_label_ids)
+        image_size = int(getattr(self.transformer.config, "sample_size", 256))
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
+        channels = int(self.transformer.config.in_channels)
+        null_class_val = int(
+            getattr(self.transformer.config, "num_classes", getattr(self.transformer.config, "num_class_embeds", 1000))
+        )
+
+        if guidance_scale is None:
+            guidance_scale = RECOMMENDED_GUIDANCE_BY_SIZE.get(image_size, 3.25)
+
+        latents = randn_tensor(
+            shape=(batch_size, channels, height, width),
+            generator=generator,
+            device=self._execution_device,
+            dtype=self.transformer.dtype,
+        )
+
+        class_labels_t = torch.tensor(class_label_ids, device=self._execution_device, dtype=torch.long).reshape(-1)
+        class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
+        class_null = torch.full_like(class_labels_t, null_class_val)
+
+        timeshift = self._resolve_timeshift(self.scheduler, image_size)
+        flow_timesteps = self._build_flow_timesteps(
+            num_inference_steps,
+            timeshift,
+            device=self._execution_device,
+            dtype=torch.float32,
+        )
+        velocity_dtype = self.transformer.dtype
+        v_prev = None
+
+        for t_cur, t_next in self.progress_bar(list(zip(flow_timesteps[:-1], flow_timesteps[1:]))):
+            dt = t_next - t_cur
+            flow_time = float(t_cur)
+            effective_guidance = (
+                guidance_scale
+                if do_classifier_free_guidance
+                and guidance_interval_min < flow_time < guidance_interval_max
+                else 1.0
+            )
+
+            latent_model_input = torch.cat([latents, latents], dim=0)
+            labels = torch.cat([class_null, class_labels_t], dim=0)
+            timesteps = torch.full(
+                (latent_model_input.shape[0],),
+                flow_time,
+                device=self._execution_device,
+                dtype=velocity_dtype,
+            )
+            model_output = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=labels,
+            ).sample
+            velocity = self._apply_classifier_free_guidance(model_output, effective_guidance)
+
+            if v_prev is None:
+                latents = latents + velocity * dt
+            else:
+                latents = latents + dt * (1.5 * velocity - 0.5 * v_prev)
+            v_prev = velocity
+
+        images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
+        if output_type == "pt":
+            images = images_pt
+        elif output_type == "np":
+            images = images_pt.permute(0, 2, 3, 1).numpy()
+        else:
+            images = self.numpy_to_pil(images_pt.permute(0, 2, 3, 1).numpy())
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)
+
+
+PixelDiTPipelineOutput = ImagePipelineOutput

PixelDiT-XL-16-512/scheduler/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.36.0",
+  "num_train_timesteps": 1000,
+  "shift": 3.0,
+  "stochastic_sampling": false
+}

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

@@ -0,0 +1,14 @@
+{
+  "_class_name": "PixelDiTTransformer2DModel",
+  "hidden_size": 1152,
+  "in_channels": 3,
+  "model_type": "pixeldit-xl",
+  "num_classes": 1000,
+  "num_groups": 16,
+  "patch_depth": 26,
+  "patch_size": 16,
+  "pixel_depth": 4,
+  "pixel_hidden_size": 16,
+  "sample_size": 512,
+  "use_pixel_abs_pos": true
+}

PixelDiT-XL-16-512/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fc9fe08e4de9709a94818cd5519604715641280e47f3f41e48993f05e22fa99
+size 3189574228

PixelDiT-XL-16-512/transformer/transformer_pixeldit.py

@@ -0,0 +1,639 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from typing import Dict, Literal, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import RMSNorm
+
+
+PIXELDIT_PRESET_CONFIGS: Dict[str, Dict[str, object]] = {
+    "pixeldit-xl": {
+        "sample_size": 256,
+        "num_groups": 16,
+        "hidden_size": 1152,
+        "pixel_hidden_size": 16,
+        "patch_depth": 26,
+        "pixel_depth": 4,
+        "patch_size": 16,
+    },
+}
+
+
+def remap_legacy_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Map wrapper/backbone keys from legacy checkpoints to native PixelDiTTransformer2DModel keys."""
+    remapped: Dict[str, torch.Tensor] = {}
+    prefixes = ("transformer.", "model.", "module.", "denoiser.", "net.")
+    for key, value in state_dict.items():
+        new_key = key
+        for prefix in prefixes:
+            if new_key.startswith(prefix):
+                new_key = new_key[len(prefix) :]
+                break
+        remapped[new_key] = value
+    return remapped
+
+
+def config_from_legacy(config: Dict[str, object]) -> Dict[str, object]:
+    """Build native config kwargs from a legacy config.json dict."""
+    model_type = config.get("model_type") or config.get("model_name") or config.get("model_size")
+    if model_type not in PIXELDIT_PRESET_CONFIGS:
+        raise ValueError(f"Unknown PixelDiT preset '{model_type}'. Known: {list(PIXELDIT_PRESET_CONFIGS)}")
+
+    preset = dict(PIXELDIT_PRESET_CONFIGS[model_type])
+    preset["num_classes"] = int(config.get("num_classes") or config.get("num_class_embeds") or 1000)
+    preset["in_channels"] = int(config.get("in_channels", 3))
+    preset["use_pixel_abs_pos"] = bool(config.get("use_pixel_abs_pos", True))
+    preset["model_type"] = model_type
+
+    for key in ("sample_size", "num_groups", "hidden_size", "pixel_hidden_size", "patch_depth", "pixel_depth", "patch_size"):
+        if config.get(key) is not None:
+            preset[key] = config[key]
+
+    return preset
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int) -> np.ndarray:
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0)
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    return get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
+    if embed_dim % 2 != 0:
+        raise ValueError("Embedding dimension must be even for 2D sin/cos positional embeddings.")
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    if embed_dim % 2 != 0:
+        raise ValueError("Embedding dimension must be even for 1D sin/cos positional embeddings.")
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+def apply_adaln(hidden_states: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return hidden_states * (1 + scale) + shift
+
+
+def precompute_freqs_cis_2d(dim: int, height: int, width: int, theta: float = 10000.0, scale: float = 16.0):
+    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))
+    x_freqs = torch.outer(x_pos, freqs).float()
+    y_freqs = torch.outer(y_pos, freqs).float()
+    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)
+    return freqs_cis.reshape(height * width, -1)
+
+
+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_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 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)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class TimestepConditioner(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(timesteps: torch.Tensor, dim: int, max_period: int = 10):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half
+        )
+        args = timesteps[..., 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, timesteps: torch.Tensor) -> torch.Tensor:
+        timestep_freq = self.timestep_embedding(timesteps, self.frequency_embedding_size)
+        mlp_dtype = next(self.mlp.parameters()).dtype
+        if timestep_freq.dtype != mlp_dtype:
+            timestep_freq = timestep_freq.to(mlp_dtype)
+        return self.mlp(timestep_freq)
+
+
+class ClassEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int):
+        super().__init__()
+        self.embedding_table = nn.Embedding(num_classes, hidden_size)
+        self.num_classes = num_classes
+
+    def forward(self, labels: torch.Tensor) -> torch.Tensor:
+        return self.embedding_table(labels)
+
+
+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, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(hidden_states)) * self.w3(hidden_states))
+
+
+class RotaryAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError("dim should be divisible by num_heads")
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, eps=eps) 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, hidden_states: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        batch_size, length, channels = hidden_states.shape
+        qkv = (
+            self.qkv(hidden_states)
+            .reshape(batch_size, length, 3, self.num_heads, channels // self.num_heads)
+            .permute(2, 0, 1, 3, 4)
+        )
+        query, key, value = qkv[0], qkv[1], qkv[2]
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        query, key = apply_rotary_emb(query, key, freqs_cis=pos)
+        query = query.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2)
+        key = key.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2).contiguous()
+        value = value.view(batch_size, -1, self.num_heads, channels // self.num_heads).transpose(1, 2).contiguous()
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, length, channels)
+        hidden_states = self.proj(hidden_states)
+        return self.proj_drop(hidden_states)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int, mlp_ratio: float = 4.0, drop: float = 0.0):
+        super().__init__()
+        hidden_dim = int(dim * mlp_ratio)
+        self.fc1 = nn.Linear(dim, hidden_dim)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_dim, dim)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.drop(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return self.drop(hidden_states)
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, out_channels: int, eps: float = 1e-6):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size, eps=eps)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(hidden_states)
+        return self.linear(hidden_states)
+
+
+class PatchTokenEmbedder(nn.Module):
+    def __init__(self, in_chans: int, embed_dim: int, 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, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.proj(hidden_states)
+        return self.norm(hidden_states)
+
+
+class PixelTokenEmbedder(nn.Module):
+    def __init__(self, in_channels: int, hidden_size_output: int, use_pixel_abs_pos: bool = True):
+        super().__init__()
+        self.in_channels = int(in_channels)
+        self.hidden_size_output = int(hidden_size_output)
+        self.use_pixel_abs_pos = bool(use_pixel_abs_pos)
+        self.proj = nn.Linear(self.in_channels, self.hidden_size_output, bias=True)
+        self._pos_cache: Dict[tuple[str, int, int], torch.Tensor] = {}
+
+    def _fetch_pixel_pos_image(self, height: int, width: int, device: torch.device, dtype: torch.dtype):
+        key = ("image", height, width)
+        if key in self._pos_cache:
+            return self._pos_cache[key].to(device=device, dtype=dtype)
+        if height == width:
+            pos_np = get_2d_sincos_pos_embed(self.hidden_size_output, height)
+        else:
+            grid_h = np.arange(height, dtype=np.float32)
+            grid_w = np.arange(width, dtype=np.float32)
+            grid = np.meshgrid(grid_w, grid_h)
+            grid = np.stack(grid, axis=0).reshape(2, 1, height, width)
+            pos_np = get_2d_sincos_pos_embed_from_grid(self.hidden_size_output, grid)
+        pos = torch.from_numpy(pos_np).to(device=device, dtype=dtype)
+        self._pos_cache[key] = pos
+        return pos
+
+    def forward(self, inputs: torch.Tensor, img_height: int, img_width: int, patch_size: int):
+        if inputs.dim() != 4:
+            raise ValueError("PixelTokenEmbedder expects inputs of shape [B,C,H,W]")
+        batch_size, channels, height, width = inputs.shape
+        if height != img_height or width != img_width:
+            raise ValueError("Input resolution does not match img_height/img_width.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+        h_tokens, w_tokens = height // patch_size, width // patch_size
+        patch_area = patch_size * patch_size
+        hidden_states = inputs.permute(0, 2, 3, 1).contiguous()
+        hidden_states = self.proj(hidden_states)
+        if self.use_pixel_abs_pos:
+            pos_full = self._fetch_pixel_pos_image(height, width, inputs.device, inputs.dtype)
+            hidden_states = hidden_states + pos_full.view(height, width, self.hidden_size_output).unsqueeze(0)
+        hidden_states = hidden_states.view(batch_size, h_tokens, patch_size, w_tokens, patch_size, self.hidden_size_output)
+        hidden_states = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous()
+        return hidden_states.view(batch_size * h_tokens * w_tokens, patch_area, self.hidden_size_output)
+
+
+class AugmentedDiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, groups: int, mlp_ratio: float = 4.0, adaLN_modulation=None, eps: float = 1e-6):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size, eps=eps)
+        self.attn = RotaryAttention(hidden_size, num_heads=groups, qkv_bias=False, eps=eps)
+        self.norm2 = RMSNorm(hidden_size, eps=eps)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp = FeedForward(hidden_size, mlp_hidden_dim)
+        self.adaLN_modulation = adaLN_modulation if adaLN_modulation is not None else nn.Sequential(
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, hidden_states: torch.Tensor, conditioning: torch.Tensor, pos: torch.Tensor):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(conditioning).chunk(
+            6, dim=-1
+        )
+        hidden_states = hidden_states + gate_msa * self.attn(
+            apply_adaln(self.norm1(hidden_states), shift_msa, scale_msa), pos
+        )
+        hidden_states = hidden_states + gate_mlp * self.mlp(
+            apply_adaln(self.norm2(hidden_states), shift_mlp, scale_mlp)
+        )
+        return hidden_states
+
+
+class PiTBlock(nn.Module):
+    def __init__(
+        self,
+        pixel_hidden_size: int,
+        patch_hidden_size: int,
+        patch_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_hidden_size: Optional[int] = None,
+        attn_num_heads: Optional[int] = None,
+        rope_fn=None,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.pixel_dim = int(pixel_hidden_size)
+        self.context_dim = int(patch_hidden_size)
+        self.patch_size = int(patch_size)
+        self.attn_dim = int(attn_hidden_size) if attn_hidden_size is not None else self.context_dim
+        self.num_heads = int(attn_num_heads) if attn_num_heads is not None else int(num_heads)
+        if self.attn_dim % self.num_heads != 0:
+            raise ValueError("pixel attention hidden size must be divisible by pixel num_heads")
+        patch_area = self.patch_size * self.patch_size
+        self.compress_to_attn = nn.Linear(patch_area * self.pixel_dim, self.attn_dim, bias=True)
+        self.expand_from_attn = nn.Linear(self.attn_dim, patch_area * self.pixel_dim, bias=True)
+        self.norm1 = RMSNorm(self.pixel_dim, eps=eps)
+        self.attn = RotaryAttention(self.attn_dim, num_heads=self.num_heads, qkv_bias=False, eps=eps)
+        self.norm2 = RMSNorm(self.pixel_dim, eps=eps)
+        self.mlp = MLP(self.pixel_dim, mlp_ratio=mlp_ratio, drop=0.0)
+        self.adaLN_modulation = nn.Sequential(nn.Linear(self.context_dim, 6 * self.pixel_dim * patch_area, bias=True))
+        self._pos_cache: Dict[tuple[int, int], torch.Tensor] = {}
+        self._rope_fn = rope_fn if rope_fn is not None else precompute_freqs_cis_2d
+
+    def _fetch_pos(self, height: int, width: int, device: torch.device):
+        key = (height, width)
+        if key in self._pos_cache:
+            return self._pos_cache[key].to(device)
+        pos = self._rope_fn(self.attn_dim // self.num_heads, height, width).to(device)
+        self._pos_cache[key] = pos
+        return pos
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        conditioning: torch.Tensor,
+        image_height: int,
+        image_width: int,
+        patch_size: int,
+    ) -> torch.Tensor:
+        batch_tokens, patch_area, channels = hidden_states.shape
+        if channels != self.pixel_dim:
+            raise ValueError(f"PiTBlock expected pixel_dim={self.pixel_dim}, got {channels}")
+        if image_height % patch_size != 0 or image_width % patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+        h_tokens, w_tokens = image_height // patch_size, image_width // patch_size
+        length = h_tokens * w_tokens
+        batch_size = batch_tokens // length
+        cond_params = self.adaLN_modulation(conditioning).view(batch_tokens, patch_area, 6 * self.pixel_dim)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(cond_params, 6, dim=-1)
+        hidden_norm = apply_adaln(self.norm1(hidden_states), shift_msa, scale_msa)
+        hidden_flat = hidden_norm.view(batch_tokens, patch_area * self.pixel_dim)
+        hidden_comp = self.compress_to_attn(hidden_flat).view(batch_size, length, self.attn_dim)
+        pos_comp = self._fetch_pos(h_tokens, w_tokens, hidden_states.device)
+        attn_out = self.attn(hidden_comp, pos_comp)
+        attn_flat = self.expand_from_attn(attn_out.view(batch_size * length, self.attn_dim))
+        attn_exp = attn_flat.view(batch_tokens, patch_area, self.pixel_dim)
+        hidden_states = hidden_states + gate_msa * attn_exp
+        mlp_out = self.mlp(apply_adaln(self.norm2(hidden_states), shift_mlp, scale_mlp))
+        hidden_states = hidden_states + gate_mlp * mlp_out
+        return hidden_states
+
+
+class PixelDiTTransformer2DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+    _skip_layerwise_casting_patterns = ["pos", "_pos_cache"]
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 256,
+        in_channels: int = 3,
+        num_groups: int = 16,
+        hidden_size: int = 1152,
+        pixel_hidden_size: int = 16,
+        patch_depth: int = 26,
+        pixel_depth: int = 4,
+        patch_size: int = 16,
+        num_classes: int = 1000,
+        use_pixel_abs_pos: bool = True,
+        norm_eps: float = 1e-6,
+        model_type: str | None = None,
+        num_class_embeds: int | None = None,
+    ):
+        super().__init__()
+        if num_class_embeds is not None:
+            num_classes = int(num_class_embeds)
+        if model_type in PIXELDIT_PRESET_CONFIGS:
+            preset = PIXELDIT_PRESET_CONFIGS[model_type]
+            sample_size = int(preset["sample_size"])
+            num_groups = int(preset["num_groups"])
+            hidden_size = int(preset["hidden_size"])
+            pixel_hidden_size = int(preset["pixel_hidden_size"])
+            patch_depth = int(preset["patch_depth"])
+            pixel_depth = int(preset["pixel_depth"])
+            patch_size = int(preset["patch_size"])
+
+        self.sample_size = int(sample_size)
+        self.in_channels = int(in_channels)
+        self.out_channels = int(in_channels)
+        self.hidden_size = int(hidden_size)
+        self.num_groups = int(num_groups)
+        self.patch_depth = int(patch_depth)
+        self.pixel_depth = int(pixel_depth)
+        self.patch_size = int(patch_size)
+        self.pixel_hidden_size = int(pixel_hidden_size)
+        self.num_classes = int(num_classes)
+        self.use_pixel_abs_pos = bool(use_pixel_abs_pos)
+        self.norm_eps = float(norm_eps)
+        self.gradient_checkpointing = False
+
+        if self.pixel_depth <= 0:
+            raise ValueError("PixelDiT expects pixel_depth > 0 to preserve the dual-level pipeline")
+
+        self.pixel_embedder = PixelTokenEmbedder(
+            self.in_channels, self.pixel_hidden_size, use_pixel_abs_pos=self.use_pixel_abs_pos
+        )
+        self.s_embedder = PatchTokenEmbedder(self.in_channels * self.patch_size**2, self.hidden_size, bias=True)
+        self.t_embedder = TimestepConditioner(self.hidden_size)
+        self.y_embedder = ClassEmbedder(self.num_classes + 1, self.hidden_size)
+
+        self.final_layer = FinalLayer(self.pixel_hidden_size, self.out_channels, eps=self.norm_eps)
+        self.patch_blocks = nn.ModuleList(
+            [AugmentedDiTBlock(self.hidden_size, self.num_groups, eps=self.norm_eps) for _ in range(self.patch_depth)]
+        )
+        self.pixel_blocks = nn.ModuleList(
+            [
+                PiTBlock(
+                    self.pixel_hidden_size,
+                    self.hidden_size,
+                    patch_size=self.patch_size,
+                    num_heads=self.num_groups,
+                    mlp_ratio=4.0,
+                    eps=self.norm_eps,
+                )
+                for _ in range(self.pixel_depth)
+            ]
+        )
+        self._precompute_pos: Dict[tuple[int, int], torch.Tensor] = {}
+        self._initialize_weights()
+
+    def _fetch_pos(self, height: int, width: int, device: torch.device):
+        key = (height, width)
+        if key in self._precompute_pos:
+            return self._precompute_pos[key].to(device)
+        pos = precompute_freqs_cis_2d(self.hidden_size // self.num_groups, height, width).to(device)
+        self._precompute_pos[key] = pos
+        return pos
+
+    def _initialize_weights(self) -> None:
+        weight = self.s_embedder.proj.weight.data
+        nn.init.xavier_uniform_(weight.view([weight.shape[0], -1]))
+        nn.init.constant_(self.s_embedder.proj.bias, 0)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        nn.init.zeros_(self.final_layer.linear.weight)
+        nn.init.zeros_(self.final_layer.linear.bias)
+        for block in self.patch_blocks:
+            nn.init.zeros_(block.adaLN_modulation[0].weight)
+            nn.init.zeros_(block.adaLN_modulation[0].bias)
+        for block in self.pixel_blocks:
+            nn.init.zeros_(block.adaLN_modulation[0].weight)
+            nn.init.zeros_(block.adaLN_modulation[0].bias)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float],
+        class_labels: Union[torch.Tensor, int],
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]:
+        if sample.dim() != 4:
+            raise ValueError("PixelDiTTransformer2DModel expects sample of shape [B,C,H,W]")
+        batch_size, _, height, width = sample.shape
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError("Image height and width must be divisible by patch_size.")
+
+        timestep = torch.as_tensor(timestep, device=sample.device)
+        if timestep.ndim == 0:
+            timestep = timestep.repeat(batch_size)
+        else:
+            timestep = timestep.reshape(-1)
+            if timestep.shape[0] == 1 and batch_size > 1:
+                timestep = timestep.repeat(batch_size)
+
+        if not torch.is_tensor(class_labels):
+            class_labels = torch.tensor(class_labels, device=sample.device, dtype=torch.long)
+        class_labels = class_labels.to(device=sample.device, dtype=torch.long).reshape(-1)
+        if class_labels.shape[0] == 1 and batch_size > 1:
+            class_labels = class_labels.repeat(batch_size)
+
+        pos = self._fetch_pos(height // self.patch_size, width // self.patch_size, sample.device)
+        x_patches = F.unfold(sample, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2)
+        t_emb = self.t_embedder(timestep.view(-1)).view(batch_size, -1, self.hidden_size)
+        y_emb = self.y_embedder(class_labels).view(batch_size, 1, self.hidden_size)
+        conditioning = F.silu(t_emb + y_emb)
+
+        patch_states = self.s_embedder(x_patches)
+        for block in self.patch_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(hidden_states, cond, position):
+                    return block(hidden_states, cond, position)
+
+                patch_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward, patch_states, conditioning, pos, use_reentrant=False
+                )
+            else:
+                patch_states = block(patch_states, conditioning, pos)
+        patch_states = F.silu(t_emb + patch_states)
+
+        length = patch_states.shape[1]
+        conditioning_states = patch_states.view(batch_size * length, self.hidden_size)
+        pixel_states = self.pixel_embedder(
+            sample, img_height=height, img_width=width, patch_size=self.patch_size
+        )
+        for block in self.pixel_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def custom_forward(hidden_states, cond):
+                    return block(hidden_states, cond, height, width, self.patch_size)
+
+                pixel_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward, pixel_states, conditioning_states, use_reentrant=False
+                )
+            else:
+                pixel_states = block(pixel_states, conditioning_states, height, width, self.patch_size)
+        pixel_states = self.final_layer(pixel_states)
+
+        patch_area = self.patch_size * self.patch_size
+        pixel_states = pixel_states.view(batch_size, length, patch_area, self.out_channels).permute(0, 3, 2, 1)
+        pixel_states = pixel_states.contiguous().view(batch_size, self.out_channels * patch_area, length)
+        output = F.fold(pixel_states, (height, width), kernel_size=self.patch_size, stride=self.patch_size)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    @classmethod
+    def from_pixeldit_checkpoint(
+        cls,
+        checkpoint_path: str,
+        model_type: Literal["pixeldit-xl"] = "pixeldit-xl",
+        map_location: str = "cpu",
+        strict: bool = True,
+    ) -> Tuple["PixelDiTTransformer2DModel", Dict[str, object]]:
+        if model_type not in PIXELDIT_PRESET_CONFIGS:
+            raise ValueError(f"Unknown PixelDiT preset '{model_type}'.")
+
+        if checkpoint_path.endswith(".safetensors"):
+            try:
+                from safetensors.torch import load_file
+            except ImportError as error:
+                raise ImportError("Install safetensors to load .safetensors checkpoints.") from error
+            state_dict = load_file(checkpoint_path, device=map_location)
+        else:
+            loaded = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
+            if isinstance(loaded, Mapping):
+                state_dict = loaded
+                for key in ("state_dict", "model", "module", "denoiser"):
+                    if key in state_dict and isinstance(state_dict[key], dict):
+                        state_dict = state_dict[key]
+                        break
+            else:
+                raise ValueError("Unsupported checkpoint format.")
+
+        config = dict(PIXELDIT_PRESET_CONFIGS[model_type])
+        config["model_type"] = model_type
+        model = cls(**config)
+        model.load_state_dict(remap_legacy_state_dict(state_dict), strict=strict)
+
+        metadata = {
+            "checkpoint_path": checkpoint_path,
+            "model_type": model_type,
+        }
+        return model, metadata
+
+    def to_pixeldit_checkpoint(self, prefix: str = "") -> Dict[str, torch.Tensor]:
+        checkpoint: Dict[str, torch.Tensor] = {}
+        for key, value in self.state_dict().items():
+            checkpoint[f"{prefix}{key}"] = value.detach().cpu()
+        return checkpoint
+
+
+PixelDiTDiffusersModel = PixelDiTTransformer2DModel

README.md

@@ -0,0 +1,241 @@
+---
+license: other
+license_name: nsclv1
+license_link: https://huggingface.co/nvidia/PixelDiT-ImageNet/blob/main/LICENSE
+library_name: diffusers
+pipeline_tag: text-to-image
+tags:
+  - diffusers
+  - image-generation
+  - class-conditional
+  - text-to-image
+  - imagenet
+  - pixeldit
+  - flow-matching
+  - pixel-space
+  - dit
+widget:
+  - text: A golden retriever playing in a sunny garden
+    output:
+      url: PixelDiT-T2I-1024/demo.png
+  - text: golden retriever
+    output:
+      url: PixelDiT-XL-16-256/demo.png
+language:
+  - en
+---
+
+# BiliSakura/PixelDiT-diffusers
+
+Self-contained PixelDiT checkpoints for Hugging Face diffusers. Each variant folder ships its own `pipeline.py`, component modules, and weights.
+
+Converted from [nvidia/PixelDiT-ImageNet](https://huggingface.co/nvidia/PixelDiT-ImageNet) and [nvidia/PixelDiT-1300M-1024px](https://huggingface.co/nvidia/PixelDiT-1300M-1024px) using [PixelDiT-diffusers](https://github.com/BiliSakura/Visual-Generative-Foundation-Model-Collection/tree/main/libs/PixelDiT-diffusers).
+
+## Available checkpoints
+
+| Subfolder | Pipeline | Task | Resolution | Source checkpoint | gFID | Params |
+| --- | --- | --- | ---: | --- | ---: | ---: |
+| [`PixelDiT-T2I-1024/`](PixelDiT-T2I-1024/) | `PixelDiTT2IPipeline` | text-to-image | 1024×1024 | `pixeldit_t2i_v1.pth` | — | ~1.3B |
+| [`PixelDiT-XL-16-256/`](PixelDiT-XL-16-256/) | `PixelDiTPipeline` | class-to-image | 256×256 | `imagenet256_pixeldit_xl_epoch320.ckpt` | 1.61 | ~700M |
+| [`PixelDiT-XL-16-512/`](PixelDiT-XL-16-512/) | `PixelDiTPipeline` | class-to-image | 512×512 | `imagenet512_pixeldit_xl.ckpt` | 1.81 | ~700M |
+
+## Repo layout
+
+```text
+BiliSakura/PixelDiT-diffusers/
+├── README.md
+├── demo_inference.py
+├── PixelDiT-T2I-1024/
+│   ├── pipeline.py
+│   ├── model_index.json
+│   ├── demo.png
+│   ├── scheduler/scheduler_config.json
+│   └── transformer/
+├── PixelDiT-XL-16-256/
+│   ├── pipeline.py
+│   ├── model_index.json
+│   ├── demo.png
+│   ├── scheduler/scheduler_config.json
+│   └── transformer/
+└── PixelDiT-XL-16-512/
+    ├── pipeline.py
+    ├── model_index.json
+    ├── scheduler/scheduler_config.json
+    └── transformer/
+```
+
+Each variant is self-contained. The `scheduler/` folder uses built-in `FlowMatchEulerDiscreteScheduler` from PyPI diffusers. No shared helper modules at inference time beyond the local variant directory.
+
+## ImageNet class labels
+
+`id2label` is embedded in each variant's `model_index.json` (DiT-style).
+
+- `pipe.id2label` — inspect id → English label correspondence
+- `pipe.labels` — reverse map (English synonym → id)
+- `pipe.get_label_ids("golden retriever")`
+- `pipe(class_labels="golden retriever", ...)` — string labels resolved automatically
+
+## Demo
+
+![PixelDiT-T2I-1024 demo](PixelDiT-T2I-1024/demo.png)
+
+Text-to-image — "A golden retriever playing in a sunny garden", 1024×1024, 50 steps, `guidance_scale=2.75`.
+
+```bash
+python demo_inference_t2i.py
+```
+
+![PixelDiT-XL-16-256 demo](PixelDiT-XL-16-256/demo.png)
+
+Class 207 — golden retriever, 256×256, 100 steps, `guidance_scale=2.75`, CFG interval `[0.1, 0.9]`.
+
+```bash
+python demo_inference.py
+```
+
+## Load from a local clone
+
+### Text-to-image 1024×1024 (`PixelDiT-T2I-1024`)
+
+```python
+from pathlib import Path
+import torch
+from diffusers import DiffusionPipeline
+
+model_dir = Path("./PixelDiT-T2I-1024").resolve()
+pipe = DiffusionPipeline.from_pretrained(
+    str(model_dir),
+    local_files_only=True,
+    custom_pipeline=str(model_dir / "pipeline.py"),
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+)
+pipe.to("cuda")
+
+generator = torch.Generator(device="cuda").manual_seed(42)
+image = pipe(
+    prompt="A golden retriever playing in a sunny garden",
+    negative_prompt="low quality, worst quality, over-saturated, blurry, deformed, watermark",
+    height=1024,
+    width=1024,
+    num_inference_steps=50,
+    guidance_scale=2.75,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+Gemma text encoder (`google/gemma-2-2b-it`) is downloaded on first run unless bundled under `text_encoder/`.
+
+### ImageNet 256×256 (`PixelDiT-XL-16-256`)
+
+```python
+from pathlib import Path
+import torch
+from diffusers import DiffusionPipeline
+
+model_dir = Path("./PixelDiT-XL-16-256").resolve()
+pipe = DiffusionPipeline.from_pretrained(
+    str(model_dir),
+    local_files_only=True,
+    custom_pipeline=str(model_dir / "pipeline.py"),
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+)
+pipe.to("cuda")
+
+print(pipe.id2label[207])
+print(pipe.get_label_ids("golden retriever"))
+
+generator = torch.Generator(device="cuda").manual_seed(42)
+image = pipe(
+    class_labels="golden retriever",
+    height=256,
+    width=256,
+    num_inference_steps=100,
+    guidance_scale=2.75,
+    guidance_interval_min=0.1,
+    guidance_interval_max=0.9,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+### ImageNet 512×512 (`PixelDiT-XL-16-512`)
+
+```python
+from pathlib import Path
+import torch
+from diffusers import DiffusionPipeline
+
+model_dir = Path("./PixelDiT-XL-16-512").resolve()
+pipe = DiffusionPipeline.from_pretrained(
+    str(model_dir),
+    local_files_only=True,
+    custom_pipeline=str(model_dir / "pipeline.py"),
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+)
+pipe.to("cuda")
+
+generator = torch.Generator(device="cuda").manual_seed(42)
+image = pipe(
+    class_labels=207,
+    height=512,
+    width=512,
+    num_inference_steps=100,
+    guidance_scale=3.5,
+    guidance_interval_min=0.1,
+    guidance_interval_max=1.0,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+## Recommended inference settings
+
+| Variant | Steps | CFG scale | Scheduler shift | CFG interval |
+| --- | ---: | ---: | ---: | --- |
+| `PixelDiT-T2I-1024` | 50 | 2.75 | 4.0 | [0.0, 1.0] |
+| `PixelDiT-XL-16-256` | 100 | 2.75 | 1.0 | [0.1, 0.9] |
+| `PixelDiT-XL-16-512` | 100 | 3.5 | 2.0 | [0.1, 1.0] |
+
+PixelDiT denoises directly in pixel space (no VAE). `height` and `width` must be divisible by the patch size (16).
+
+## Conversion
+
+```bash
+cd libs/PixelDiT-diffusers
+
+python scripts/convert_pixeldit_t2i_to_diffusers.py \
+  --checkpoint /path/to/pixeldit_t2i_v1.pth \
+  --config /path/to/config.json \
+  --output /path/to/PixelDiT-T2I-1024 \
+  --sample-size 1024 \
+  --scheduler-shift 4.0 \
+  --check-load
+
+python scripts/convert_pixeldit_to_diffusers.py \
+  --checkpoint /path/to/imagenet256_pixeldit_xl_epoch320.ckpt \
+  --output /path/to/PixelDiT-XL-16-256 \
+  --model-size pixeldit-xl \
+  --sample-size 256 \
+  --scheduler-shift 1.0 \
+  --check-load \
+  --id2label /path/to/id2label_en.json
+```
+
+## Citation
+
+```bibtex
+@inproceedings{yu2025pixeldit,
+      title={PixelDiT: Pixel Diffusion Transformers for Image Generation},
+      author={Yongsheng Yu and Wei Xiong and Weili Nie and Yichen Sheng and Shiqiu Liu and Jiebo Luo},
+      booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+      year={2026},
+}
+```
+
+## License
+
+Weights are converted from NVIDIA checkpoints released under the [NSCLv1 License](https://huggingface.co/nvidia/PixelDiT-ImageNet/blob/main/LICENSE). Use for non-commercial research and evaluation only.

demo_inference.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+"""Generate demo images for PixelDiT class-conditional checkpoints."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from diffusers import DiffusionPipeline
+
+REPO_ROOT = Path(__file__).resolve().parent
+
+VARIANTS = {
+    "256": {
+        "dir": REPO_ROOT / "PixelDiT-XL-16-256",
+        "height": 256,
+        "width": 256,
+        "num_inference_steps": 100,
+        "guidance_scale": 3.25,
+        "class_label": "golden retriever",
+        "seed": 7,
+    },
+    "512": {
+        "dir": REPO_ROOT / "PixelDiT-XL-16-512",
+        "height": 512,
+        "width": 512,
+        "num_inference_steps": 100,
+        "guidance_scale": 3.75,
+        "class_label": "golden retriever",
+        "seed": 7,
+    },
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run PixelDiT demo inference.")
+    parser.add_argument(
+        "--variant",
+        choices=sorted(VARIANTS),
+        default="256",
+        help="Checkpoint resolution variant to sample.",
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Generate demo.png for every supported variant.",
+    )
+    return parser.parse_args()
+
+
+def run_variant(name: str) -> Path:
+    settings = VARIANTS[name]
+    model_dir = settings["dir"]
+    output_path = model_dir / "demo.png"
+
+    pipe = DiffusionPipeline.from_pretrained(
+        str(model_dir),
+        local_files_only=True,
+        custom_pipeline=str(model_dir / "pipeline.py"),
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+    )
+    pipe.to("cuda")
+
+    print(f"[{name}] {settings['class_label']} -> {pipe.get_label_ids(settings['class_label'])}")
+    print(f"[{name}] scheduler shift={pipe.scheduler.config.shift}")
+
+    generator = torch.Generator(device="cuda").manual_seed(settings["seed"])
+    image = pipe(
+        class_labels=settings["class_label"],
+        height=settings["height"],
+        width=settings["width"],
+        num_inference_steps=settings["num_inference_steps"],
+        guidance_scale=settings["guidance_scale"],
+        guidance_interval_min=0.1,
+        guidance_interval_max=1.0,
+        generator=generator,
+    ).images[0]
+    image.save(output_path)
+    print(f"[{name}] Saved demo image to {output_path}")
+    return output_path
+
+
+def main() -> None:
+    args = parse_args()
+    if args.all:
+        for name in VARIANTS:
+            run_variant(name)
+        return
+    run_variant(args.variant)
+
+
+if __name__ == "__main__":
+    main()

demo_inference_t2i.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""Generate a demo image with PixelDiT-T2I-1024."""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parent
+MODEL_DIR = REPO_ROOT / "PixelDiT-T2I-1024"
+OUTPUT_PATH = MODEL_DIR / "demo.png"
+
+DIFFUSERS_SRC = None
+for parent in Path(__file__).resolve().parents:
+    candidate = parent / "libs/diffusers/src"
+    if (candidate / "diffusers/schedulers/flow_dpm.py").is_file():
+        DIFFUSERS_SRC = candidate
+        break
+if DIFFUSERS_SRC is None:
+    fallback = Path("/data/projects/Visual-Generative-Foundation-Model-Collection/libs/diffusers/src")
+    if (fallback / "diffusers/schedulers/flow_dpm.py").is_file():
+        DIFFUSERS_SRC = fallback
+if DIFFUSERS_SRC is None:
+    raise ImportError("Could not locate libs/diffusers/src for PixelDiT flow DPM sampling.")
+if str(DIFFUSERS_SRC) not in sys.path:
+    sys.path.insert(0, str(DIFFUSERS_SRC))
+import os
+
+os.environ.setdefault("PIXELDIT_DIFFUSERS_SRC", str(DIFFUSERS_SRC))
+
+import torch
+
+
+def _load_pipeline_class():
+    spec = importlib.util.spec_from_file_location("pixeldit_t2i_pipeline", MODEL_DIR / "pipeline.py")
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Unable to load pipeline from {MODEL_DIR / 'pipeline.py'}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.PixelDiTT2IPipeline
+
+
+def main() -> None:
+    pipe_cls = _load_pipeline_class()
+    pipe = pipe_cls.from_pretrained(
+        str(MODEL_DIR),
+        local_files_only=True,
+        torch_dtype=torch.bfloat16,
+    )
+    pipe.to("cuda")
+
+    prompt = "A golden retriever playing in a sunny garden"
+    print("text_encoder:", type(pipe.text_encoder).__name__)
+    print("prompt:", prompt)
+
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    image = pipe(
+        prompt=prompt,
+        negative_prompt="low quality, worst quality, over-saturated, blurry, deformed, watermark",
+        height=1024,
+        width=1024,
+        num_inference_steps=50,
+        guidance_scale=2.75,
+        use_chi_prompt=True,
+        generator=generator,
+    ).images[0]
+    image.save(OUTPUT_PATH)
+    print(f"Saved demo image to {OUTPUT_PATH}")
+
+
+if __name__ == "__main__":
+    main()