Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,37 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz 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
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+Looped-DiT-B-16/demo.png filter=lfs diff=lfs merge=lfs -text
+Looped-DiT-B-32/demo.png filter=lfs diff=lfs merge=lfs -text

Looped-DiT-B-16/model_index.json

@@ -0,0 +1,23 @@
+{
+  "_class_name": [
+    "pipeline",
+    "LoopedDiTPipeline"
+  ],
+  "_diffusers_version": "0.39.0",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "T5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "T5Tokenizer"
+  ],
+  "transformer": [
+    "transformer_looped_dit",
+    "LoopedDiTTransformer2DModel"
+  ]
+}

Looped-DiT-B-16/pipeline.py

@@ -0,0 +1,763 @@
+# Copyright 2025 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.
+
+import inspect
+from typing import Any, Callable
+
+import torch
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import deprecate, is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+from transformers import AutoTokenizer, T5EncoderModel
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+# Training clamps the flow-matching denominator so the loss stays finite at t -> 1.
+VELOCITY_DENOM_MIN = 0.05
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from diffusers import DiffusionPipeline
+
+        >>> model_dir = Path("checkpoints/looped-dit-b16").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,
+        ... ).to("cuda")
+
+        >>> image = pipe(
+        ...     "a red cube on top of a blue sphere",
+        ...     num_inference_steps=100,
+        ...     guidance_scale=6.0,
+        ...     num_loops=4,
+        ...     generator=torch.Generator(device="cuda").manual_seed(0),
+        ... ).images[0]
+        >>> image.save("sample.png")
+
+        >>> # Hugging Face Hub style model id: UserID/RepoID
+        >>> # RepoID is usually like "modelname-diffusers"
+        >>> # Example: "your-user/Looped-DiT-diffusers"
+        ```
+"""
+
+
+def paper_euler_sigmas(num_inference_steps: int) -> list[float]:
+    r"""
+    Sigma grid of the training Euler sampler.
+
+    Training integrates flow time `t` from 0 (noise) to 1 (data) with
+    `torch.linspace(0, 1, steps + 1)`. Flow-match schedulers step in sigma
+    `1 - t` and append the terminal 0 themselves, so the returned list omits that 0.
+
+    Args:
+        num_inference_steps (`int`):
+            Number of Euler steps. Must be positive.
+
+    Returns:
+        `list[float]`: `num_inference_steps` sigmas starting at 1 and ending at `1 / steps`.
+    """
+    if num_inference_steps <= 0:
+        raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}.")
+    flow_time = torch.linspace(0.0, 1.0, num_inference_steps + 1)
+    return (1.0 - flow_time)[:-1].tolist()
+
+
+class LoopedDiTPipeline(DiffusionPipeline):
+    r"""
+    Text-to-image pipeline for Looped-DiT.
+
+    Looped-DiT denoises directly in RGB pixel space (no VAE). The transformer predicts the clean
+    image `x0`. This pipeline converts that prediction to a flow-matching velocity and integrates it
+    with a diffusers scheduler. The default scheduler is [`FlowMatchEulerDiscreteScheduler`] on the
+    same uniform grid as the paper (100 steps, shift 1). Any [`KarrasDiffusionSchedulers`] instance
+    can be assigned to `pipe.scheduler` without other code changes.
+
+    Classifier-free guidance uses the training null condition: an all-zero text mask, which the
+    denoiser replaces with its mask token. There is no separate negative-prompt encoder.
+
+    The pipeline inherits from [`DiffusionPipeline`]. Check the superclass documentation for the
+    generic methods (download, save, device placement, CPU offload).
+
+    Args:
+        transformer ([`ModelMixin`]):
+            Looped-DiT denoiser (`LoopedDiTTransformer2DModel`) that predicts `x0` in pixel space.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`] or [`KarrasDiffusionSchedulers`]):
+            Scheduler used to step the flow. The paper setting is
+            `FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=1.0)`.
+        tokenizer ([`~transformers.AutoTokenizer`], *optional*):
+            Tokenizer for the frozen text encoder. Loaded from `text_encoder_name` when missing.
+        text_encoder ([`~transformers.T5EncoderModel`], *optional*):
+            Frozen FLAN-T5 encoder. Loaded from `text_encoder_name` when missing.
+            `DiffusionPipeline.from_pretrained(..., torch_dtype=torch.bfloat16)` keeps this encoder in bf16 with the denoiser.
+        text_encoder_name (`str`, defaults to `"google/flan-t5-large"`):
+            Hub id or local path used when `tokenizer` / `text_encoder` are not passed.
+        prompt_length (`int`, *optional*):
+            Token length prompts are padded or truncated to. Defaults to `tokenizer.model_max_length`.
+        noise_scale (`float`, defaults to 2.0):
+            Standard deviation of the initial Gaussian, matching the training noise scale.
+        default_num_inference_steps (`int`, defaults to 100):
+            Step count used when `__call__` does not pass `num_inference_steps`.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer"
+    _optional_components = ["tokenizer", "text_encoder"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "prompt_attention_mask"]
+
+    def __init__(
+        self,
+        transformer: ModelMixin,
+        scheduler: KarrasDiffusionSchedulers | SchedulerMixin,
+        tokenizer: Any | None = None,
+        text_encoder: T5EncoderModel | None = None,
+        text_encoder_name: str = "google/flan-t5-large",
+        prompt_length: int | None = None,
+        noise_scale: float = 2.0,
+        default_num_inference_steps: int = 100,
+    ):
+        super().__init__()
+        if prompt_length is None and tokenizer is not None:
+            prompt_length = int(getattr(tokenizer, "model_max_length", 256))
+        if prompt_length is None:
+            prompt_length = 256
+        if scheduler is None:
+            scheduler = self._default_scheduler()
+        if noise_scale <= 0:
+            raise ValueError(f"`noise_scale` must be positive, got {noise_scale}.")
+        if prompt_length < 1:
+            raise ValueError(f"`prompt_length` must be positive, got {prompt_length}.")
+        if default_num_inference_steps < 1:
+            raise ValueError(f"`default_num_inference_steps` must be positive, got {default_num_inference_steps}.")
+
+        self.register_modules(
+            transformer=transformer,
+            scheduler=scheduler,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+        )
+        self.register_to_config(
+            text_encoder_name=text_encoder_name,
+            prompt_length=int(prompt_length),
+            noise_scale=float(noise_scale),
+            default_num_inference_steps=int(default_num_inference_steps),
+        )
+
+    @staticmethod
+    def _default_scheduler() -> FlowMatchEulerDiscreteScheduler:
+        r"""
+        Build the paper's Euler scheduler.
+
+        Returns:
+            [`FlowMatchEulerDiscreteScheduler`]: 1000 training timesteps, shift 1, deterministic.
+        """
+        kwargs: dict[str, Any] = {"num_train_timesteps": 1000, "shift": 1.0}
+        if "stochastic_sampling" in inspect.signature(FlowMatchEulerDiscreteScheduler.__init__).parameters:
+            kwargs["stochastic_sampling"] = False
+        return FlowMatchEulerDiscreteScheduler(**kwargs)
+
+    def _encode_prompt(
+        self,
+        prompt: str | list[str] | None,
+        device: torch.device,
+        num_images_per_prompt: int,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Deprecated alias of [`~LoopedDiTPipeline.encode_prompt`].
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt or prompts to tokenize and encode.
+            device (`torch.device`):
+                Device of the returned tensors.
+            num_images_per_prompt (`int`):
+                How many times to repeat each prompt embedding.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Already encoded prompts of shape `(batch, sequence, text_dim)`.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask of shape `(batch, sequence)` with 1 on real tokens. Required with `prompt_embeds`
+                only when padding should be replaced by the mask token; otherwise a mask of ones is used.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`: Prompt embeddings and the attention mask.
+        """
+        deprecation_message = (
+            "`_encode_prompt()` is deprecated and will be removed in a future version. Use `encode_prompt()` instead."
+        )
+        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
+        return self.encode_prompt(prompt, device, num_images_per_prompt, prompt_embeds, prompt_attention_mask)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str] | None,
+        device: torch.device,
+        num_images_per_prompt: int,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Encode prompts with the frozen FLAN-T5 encoder.
+
+        Prompts are padded or truncated to `config.prompt_length`. The unconditional branch of
+        classifier-free guidance is not encoded here: the denoiser builds it by zeroing this mask.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt or prompts to tokenize. Ignored when `prompt_embeds` is passed.
+            device (`torch.device`):
+                Device of the returned tensors.
+            num_images_per_prompt (`int`):
+                Number of times to repeat each encoded prompt along the batch dimension.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed embeddings of shape `(batch, sequence, text_dim)`. When set, `prompt` is ignored.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask of shape `(batch, sequence)`, 1 for tokens that should condition the model. When
+                `prompt_embeds` is set and this is omitted, every position is treated as a real token.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`:
+                Embeddings `(batch * num_images_per_prompt, sequence, text_dim)` and a mask of the same batch.
+        """
+        if num_images_per_prompt < 1:
+            raise ValueError(f"`num_images_per_prompt` must be >= 1, got {num_images_per_prompt}.")
+
+        if prompt_embeds is None:
+            if isinstance(prompt, str):
+                prompt = [prompt]
+            if self.tokenizer is None:
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    self.config.text_encoder_name, model_max_length=int(self.config.prompt_length)
+                )
+            if self.text_encoder is None:
+                self.text_encoder = T5EncoderModel.from_pretrained(self.config.text_encoder_name)
+                self.text_encoder.requires_grad_(False)
+            self.text_encoder.eval()
+            encoder_device = next(self.text_encoder.parameters()).device
+            if encoder_device != device:
+                self.text_encoder.to(device)
+            tokens = self.tokenizer(
+                prompt,
+                max_length=int(self.config.prompt_length),
+                padding="max_length",
+                truncation=True,
+                return_tensors="pt",
+            )
+            input_ids = tokens.input_ids.to(device)
+            prompt_attention_mask = tokens.attention_mask.to(device)
+            prompt_embeds = self.text_encoder(input_ids=input_ids, attention_mask=prompt_attention_mask).last_hidden_state
+        else:
+            prompt_embeds = prompt_embeds.to(device)
+            if prompt_attention_mask is None:
+                prompt_attention_mask = torch.ones(
+                    prompt_embeds.shape[:2], device=device, dtype=torch.long
+                )
+            else:
+                prompt_attention_mask = prompt_attention_mask.to(device)
+            if prompt_embeds.shape[0] != prompt_attention_mask.shape[0]:
+                raise ValueError(
+                    "`prompt_embeds` and `prompt_attention_mask` must have the same batch size, got "
+                    f"{prompt_embeds.shape[0]} and {prompt_attention_mask.shape[0]}."
+                )
+
+        if num_images_per_prompt != 1:
+            prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+            prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_images_per_prompt, dim=0)
+        return prompt_embeds, prompt_attention_mask
+
+    def prepare_extra_step_kwargs(
+        self, generator: torch.Generator | list[torch.Generator] | None, eta: float
+    ) -> dict[str, Any]:
+        r"""
+        Extra arguments forwarded to `scheduler.step`, depending on what that method accepts.
+
+        Args:
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Generator passed through when the scheduler step samples noise.
+            eta (`float`):
+                DDIM eta in `[0, 1]`. Ignored by schedulers whose `step` has no `eta` argument.
+
+        Returns:
+            `dict`: Keyword arguments for `scheduler.step`.
+        """
+        extra_step_kwargs: dict[str, Any] = {}
+        step_params = set(inspect.signature(self.scheduler.step).parameters.keys())
+        if "eta" in step_params:
+            extra_step_kwargs["eta"] = eta
+        if "generator" in step_params:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt: str | list[str] | None,
+        height: int,
+        width: int,
+        callback_steps: int | None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+        num_inference_steps: int = 100,
+        guidance_scale: float = 6.0,
+        num_loops: int | None = None,
+        output_type: str = "pil",
+    ) -> None:
+        r"""
+        Validate generation arguments and raise `ValueError` or `TypeError` on misuse.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt text. Mutually exclusive with `prompt_embeds`.
+            height (`int`):
+                Output height in pixels. Must equal the transformer's trained `image_size`.
+            width (`int`):
+                Output width in pixels. Must equal the transformer's trained `image_size`.
+            callback_steps (`int`, *optional*):
+                Deprecated callback period. When set, it must be a positive integer.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed text embeddings. Required when `prompt` is omitted.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask paired with `prompt_embeds`.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names the step callback may read. Each name must be listed on
+                `_callback_tensor_inputs`.
+            num_inference_steps (`int`):
+                Denoising steps. Must be positive.
+            guidance_scale (`float`):
+                Classifier-free guidance scale. Must be finite. `1` disables guidance.
+            num_loops (`int`, *optional*):
+                Loop depth. `None` uses the depth stored on the transformer. Otherwise `>= 1`, and
+                untied models cannot exceed the trained depth.
+            output_type (`str`):
+                One of `"pil"`, `"np"`, `"pt"`, or `"latent"`.
+        """
+        image_size = int(self.transformer.config.image_size)
+        patch_size = int(self.transformer.config.patch_size)
+        if height != image_size or width != image_size:
+            raise ValueError(
+                f"Looped-DiT uses a fixed positional grid of {image_size}x{image_size} "
+                f"(patch size {patch_size}). Got height={height}, width={width}."
+            )
+        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)}.")
+
+        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type {type(callback_steps)}."
+            )
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            key in self._callback_tensor_inputs for key in callback_on_step_end_tensor_inputs
+        ):
+            unexpected = [key for key in callback_on_step_end_tensor_inputs if key not in self._callback_tensor_inputs]
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {unexpected}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Pass only one of them.")
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
+        if prompt is not None and not isinstance(prompt, str) and not (
+            isinstance(prompt, list) and all(isinstance(item, str) for item in prompt)
+        ):
+            raise TypeError(f"`prompt` has to be a string or a list of strings, got {type(prompt)}.")
+        if prompt_embeds is not None and prompt_embeds.ndim != 3:
+            raise ValueError(f"`prompt_embeds` must have shape (batch, sequence, dim), got {tuple(prompt_embeds.shape)}.")
+        if prompt_attention_mask is not None and prompt_embeds is None:
+            raise ValueError("`prompt_attention_mask` was passed without `prompt_embeds`.")
+
+        if num_inference_steps <= 0:
+            raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}.")
+        if not torch.isfinite(torch.tensor(guidance_scale)):
+            raise ValueError(f"`guidance_scale` must be finite, got {guidance_scale}.")
+        if num_loops is not None:
+            if int(num_loops) < 1:
+                raise ValueError(f"`num_loops` must be >= 1, got {num_loops}.")
+            trained = int(self.transformer.config.num_loops)
+            if not bool(self.transformer.config.share_loop_weights) and int(num_loops) > trained:
+                raise ValueError(
+                    f"This checkpoint does not share loop weights, so `num_loops` cannot exceed the trained "
+                    f"depth {trained}. Got {num_loops}."
+                )
+        if output_type not in {"pil", "np", "pt", "latent"}:
+            raise ValueError(f"Unsupported `output_type` {output_type!r}. Choose from 'pil', 'np', 'pt', 'latent'.")
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        r"""
+        Sample the initial pixel-space noise, or validate a tensor the caller already sampled.
+
+        Looped-DiT has no VAE, so "latents" here are RGB images. Fresh noise is scaled by
+        `config.noise_scale` (2.0 in the paper). A provided `latents` tensor is not rescaled.
+
+        Args:
+            batch_size (`int`):
+                Number of images, including `num_images_per_prompt`.
+            num_channels (`int`):
+                Channel count. 3 for RGB.
+            height (`int`):
+                Image height in pixels.
+            width (`int`):
+                Image width in pixels.
+            dtype (`torch.dtype`):
+                Dtype of freshly sampled noise. The integration itself is accumulated in float32.
+            device (`torch.device`):
+                Device of the returned tensor.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Per-call RNG. A list must have length `batch_size`.
+            latents (`torch.Tensor`, *optional*):
+                Starting noise of shape `(batch_size, num_channels, height, width)`.
+
+        Returns:
+            `torch.Tensor`: Starting noise of shape `(batch_size, num_channels, height, width)`.
+        """
+        shape = (batch_size, num_channels, height, width)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You passed a list of {len(generator)} generators for a batch of {batch_size}. "
+                "The two lengths must match."
+            )
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = latents * float(self.config.noise_scale)
+        else:
+            latents = latents.to(device=device)
+            if tuple(latents.shape) != shape:
+                raise ValueError(f"`latents` shape {tuple(latents.shape)} does not match the expected {shape}.")
+        return latents
+
+    @property
+    def guidance_scale(self) -> float:
+        r"""
+        Classifier-free guidance scale of the call that is currently running.
+
+        Returns:
+            `float`: The scale set by the active `__call__`.
+        """
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self) -> bool:
+        r"""
+        Whether the active call runs a conditional and an unconditional forward.
+
+        Returns:
+            `bool`: True when `guidance_scale != 1`.
+        """
+        return self._guidance_scale != 1.0
+
+    @property
+    def num_timesteps(self) -> int:
+        r"""
+        Number of scheduler timesteps in the active call.
+
+        Returns:
+            `int`: Length of the timestep schedule.
+        """
+        return self._num_timesteps
+
+    @property
+    def interrupt(self) -> bool:
+        r"""
+        Whether the active denoising loop should skip remaining steps.
+
+        Returns:
+            `bool`: True after the caller sets `pipeline._interrupt = True`.
+        """
+        return self._interrupt
+
+    def _images_from_latents(self, latents: torch.Tensor, output_type: str) -> torch.Tensor | list[Image.Image] | Any:
+        r"""
+        Convert pixel-space samples in `[-1, 1]` to the requested output type.
+
+        Quantization matches the original sampler: `uint8(clamp(x, -1, 1) * 127.5 + 128)`.
+
+        Args:
+            latents (`torch.Tensor`):
+                Samples of shape `(batch, channels, height, width)` in model range `[-1, 1]`.
+            output_type (`str`):
+                `"latent"` returns `latents` unchanged. `"pt"` is float RGB in `[0, 1]`. `"np"` is
+                `uint8` HWC arrays. `"pil"` is a list of `PIL.Image.Image`.
+
+        Returns:
+            Images in the requested type.
+        """
+        if output_type == "latent":
+            return latents
+        images = (latents.float().clamp(-1, 1) * 127.5 + 128.0).clamp(0, 255).to(torch.uint8)
+        if output_type == "pt":
+            return images.float() / 255.0
+        arrays = images.permute(0, 2, 3, 1).cpu().numpy()
+        if output_type == "np":
+            return arrays
+        return [Image.fromarray(image) for image in arrays]
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        timesteps: list[int] | None = None,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 6.0,
+        num_images_per_prompt: int = 1,
+        num_loops: int | None = None,
+        eta: float = 0.0,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Callable[[int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        **kwargs,
+    ) -> ImagePipelineOutput | tuple:
+        r"""
+        Generate images from text prompts.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt or prompts to guide image generation. Required unless `prompt_embeds` is passed.
+            height (`int`, *optional*):
+                Image height in pixels. Defaults to the transformer's trained resolution (512).
+                Other resolutions are rejected: the positional embedding is a fixed grid.
+            width (`int`, *optional*):
+                Image width in pixels. Defaults to the trained resolution and must match `height`.
+            num_inference_steps (`int`, *optional*):
+                Denoising steps. Defaults to `config.default_num_inference_steps` (100).
+            timesteps (`list[int]`, *optional*):
+                Custom scheduler timesteps, descending. Mutually exclusive with `sigmas`. Ignored by the
+                paper Euler grid, which is selected only when both `timesteps` and `sigmas` are omitted
+                and the scheduler is [`FlowMatchEulerDiscreteScheduler`].
+            sigmas (`list[float]`, *optional*):
+                Custom sigmas passed to `scheduler.set_timesteps`. Mutually exclusive with `timesteps`.
+            guidance_scale (`float`, defaults to 6.0):
+                Classifier-free guidance scale from the paper. Guidance is on when this is not `1`.
+                The unconditional branch is an empty text mask, not a negative prompt.
+            num_images_per_prompt (`int`, defaults to 1):
+                How many images to sample for each prompt.
+            num_loops (`int`, *optional*):
+                How many times to run the shared middle blocks. `None` uses the trained depth. Other
+                depths work without retraining when loop weights are shared.
+            eta (`float`, defaults to 0.0):
+                DDIM eta. Ignored by the flow-match Euler scheduler.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG for the initial noise. `None` uses PyTorch's global generator, which is what
+                `torch.manual_seed` seeds.
+            latents (`torch.Tensor`, *optional*):
+                Initial noise `(batch, 3, height, width)`. Not multiplied by `noise_scale`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed FLAN-T5 states `(batch, sequence, text_dim)` in place of `prompt`.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask `(batch, sequence)` paired with `prompt_embeds`. 1 marks real tokens.
+            output_type (`str`, defaults to `"pil"`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"` (pixels in model range).
+            return_dict (`bool`, defaults to `True`):
+                Return [`ImagePipelineOutput`] when `True`, otherwise a one-tuple of images.
+            callback_on_step_end (`Callable` or `PipelineCallback`, *optional*):
+                Called as `callback_on_step_end(pipeline, step, timestep, callback_kwargs)` after each
+                scheduler step. Return a dict to replace tensors listed in
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`):
+                Tensor names passed to the step callback. Must be a subset of `_callback_tensor_inputs`.
+
+        Examples:
+
+        Returns:
+            [`ImagePipelineOutput`] or `tuple`:
+                When `return_dict` is `True`, [`ImagePipelineOutput`] with the images. Otherwise a tuple
+                whose first element is the images.
+        """
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        if kwargs:
+            raise TypeError(f"Unexpected arguments: {sorted(kwargs)}.")
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        image_size = int(self.transformer.config.image_size)
+        height = image_size if height is None else int(height)
+        width = image_size if width is None else int(width)
+        if num_inference_steps is None:
+            num_inference_steps = int(self.config.default_num_inference_steps)
+
+        # 1. Check inputs.
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            callback_steps,
+            prompt_embeds,
+            prompt_attention_mask,
+            callback_on_step_end_tensor_inputs,
+            num_inference_steps,
+            guidance_scale,
+            num_loops,
+            output_type,
+        )
+
+        self._guidance_scale = float(guidance_scale)
+        self._interrupt = False
+
+        # 2. Define call parameters.
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        device = self._execution_device
+
+        # 3. Encode input prompt.
+        prompt_embeds, prompt_attention_mask = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+        )
+        prompt_embeds = prompt_embeds.to(device=device, dtype=self.transformer.dtype)
+        prompt_attention_mask = prompt_attention_mask.to(device=device)
+
+        # 4. Prepare timesteps.
+        # The paper's Euler grid is the training linspace. Other schedulers keep their own spacing.
+        # A caller-supplied `timesteps` or `sigmas` always wins.
+        if timesteps is None and sigmas is None and isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+            sigmas = paper_euler_sigmas(num_inference_steps)
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, sigmas)
+        if getattr(self.scheduler.config, "stochastic_sampling", False):
+            raise ValueError(
+                "Looped-DiT's training sampler is deterministic. Set `stochastic_sampling=False` on "
+                "FlowMatchEulerDiscreteScheduler, or assign a different scheduler."
+            )
+
+        # 5. Prepare latent variables (pixel-space noise; there is no VAE).
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            int(self.transformer.config.in_channels),
+            height,
+            width,
+            self.transformer.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_train_timesteps = int(self.scheduler.config.num_train_timesteps)
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        # 7. Denoising loop.
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                model_latents = latents
+                if hasattr(self.scheduler, "scale_model_input"):
+                    model_latents = self.scheduler.scale_model_input(model_latents, t)
+                text = prompt_embeds
+                mask = prompt_attention_mask
+                if self.do_classifier_free_guidance:
+                    model_latents = torch.cat([model_latents, model_latents], dim=0)
+                    text = torch.cat([text, text], dim=0)
+                    mask = torch.cat([mask, torch.zeros_like(mask)], dim=0)
+
+                # fp32 latents with bf16 weights match the old sampler, which autocasts the forward.
+                amp_dtype = self.transformer.dtype
+                use_amp = model_latents.is_cuda and amp_dtype in (torch.float16, torch.bfloat16)
+                if not use_amp:
+                    model_latents = model_latents.to(dtype=amp_dtype)
+                with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp):
+                    x0 = self.transformer(model_latents, text, mask, num_loops=num_loops)
+                x0 = x0.float()
+                if self.do_classifier_free_guidance:
+                    x0_cond, x0_uncond = x0.chunk(2)
+                    x0 = x0_uncond + self.guidance_scale * (x0_cond - x0_uncond)
+
+                # sigma = 1 - t_flow. Passing -velocity makes `x + (sigma_next - sigma) * model_output`
+                # equal the training update `x + velocity * (t_next - t)`.
+                flow_time = (1.0 - t.to(device=latents.device, dtype=torch.float32) / num_train_timesteps)
+                velocity = (x0 - latents.float()) / (1.0 - flow_time).clamp_min(VELOCITY_DENOM_MIN)
+                latents = self.scheduler.step(-velocity, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {key: locals()[key] for key in callback_on_step_end_tensor_inputs}
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    prompt_attention_mask = callback_outputs.pop("prompt_attention_mask", prompt_attention_mask)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        images = self._images_from_latents(latents, output_type)
+
+        # Offload all models.
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

Looped-DiT-B-16/scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.39.0",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

Looped-DiT-B-16/text_encoder/README.md

@@ -0,0 +1,276 @@
+---
+language: 
+- en
+- fr
+- ro
+- de
+- multilingual
+
+widget:
+- text: "Translate to German:  My name is Arthur"
+  example_title: "Translation"
+- text: "Please answer to the following question. Who is going to be the next Ballon d'or?"
+  example_title: "Question Answering"
+- text: "Q: Can Geoffrey Hinton have a conversation with George Washington? Give the rationale before answering."
+  example_title: "Logical reasoning"
+- text: "Please answer the following question. What is the boiling point of Nitrogen?"
+  example_title: "Scientific knowledge"
+- text: "Answer the following yes/no question. Can you write a whole Haiku in a single tweet?"
+  example_title: "Yes/no question"
+- text: "Answer the following yes/no question by reasoning step-by-step. Can you write a whole Haiku in a single tweet?"
+  example_title: "Reasoning task"
+- text: "Q: ( False or not False or False ) is? A: Let's think step by step"
+  example_title: "Boolean Expressions"
+- text: "The square root of x is the cube root of y. What is y to the power of 2, if x = 4?"
+  example_title: "Math reasoning"
+- text: "Premise:  At my age you will probably have learnt one lesson. Hypothesis:  It's not certain how many lessons you'll learn by your thirties. Does the premise entail the hypothesis?"
+  example_title: "Premise and hypothesis"
+
+tags:
+- text2text-generation
+
+datasets:
+- svakulenk0/qrecc
+- taskmaster2
+- djaym7/wiki_dialog
+- deepmind/code_contests
+- lambada
+- gsm8k
+- aqua_rat
+- esnli
+- quasc
+- qed
+
+
+license: apache-2.0
+---
+
+# Model Card for FLAN-T5 large
+
+<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/flan2_architecture.jpg"
+alt="drawing" width="600"/>
+
+#  Table of Contents
+
+0. [TL;DR](#TL;DR)
+1. [Model Details](#model-details)
+2. [Usage](#usage)
+3. [Uses](#uses)
+4. [Bias, Risks, and Limitations](#bias-risks-and-limitations)
+5. [Training Details](#training-details)
+6. [Evaluation](#evaluation)
+7. [Environmental Impact](#environmental-impact)
+8. [Citation](#citation)
+9. [Model Card Authors](#model-card-authors)
+
+# TL;DR
+
+If you already know T5, FLAN-T5 is just better at everything. For the same number of parameters, these models have been fine-tuned on more than 1000 additional tasks covering also more languages. 
+As mentioned in the first few lines of the abstract : 
+>  Flan-PaLM 540B achieves state-of-the-art performance on several benchmarks, such as 75.2% on five-shot MMLU. We also publicly release Flan-T5 checkpoints,1 which achieve strong few-shot performance even compared to much larger models, such as PaLM 62B. Overall, instruction finetuning is a general method for improving the performance and usability of pretrained language models.
+
+**Disclaimer**: Content from **this** model card has been written by the Hugging Face team, and parts of it were copy pasted from the [T5 model card](https://huggingface.co/t5-large).
+
+# Model Details
+
+## Model Description
+
+
+- **Model type:** Language model
+- **Language(s) (NLP):** English, Spanish, Japanese, Persian, Hindi, French, Chinese, Bengali, Gujarati, German, Telugu, Italian, Arabic, Polish, Tamil, Marathi, Malayalam, Oriya, Panjabi, Portuguese, Urdu, Galician, Hebrew, Korean, Catalan, Thai, Dutch, Indonesian, Vietnamese, Bulgarian, Filipino, Central Khmer, Lao, Turkish, Russian, Croatian, Swedish, Yoruba, Kurdish, Burmese, Malay, Czech, Finnish, Somali, Tagalog, Swahili, Sinhala, Kannada, Zhuang, Igbo, Xhosa, Romanian, Haitian, Estonian, Slovak, Lithuanian, Greek, Nepali, Assamese, Norwegian
+- **License:** Apache 2.0
+- **Related Models:** [All FLAN-T5 Checkpoints](https://huggingface.co/models?search=flan-t5)
+- **Original Checkpoints:** [All Original FLAN-T5 Checkpoints](https://github.com/google-research/t5x/blob/main/docs/models.md#flan-t5-checkpoints)
+- **Resources for more information:**
+  - [Research paper](https://arxiv.org/pdf/2210.11416.pdf)
+  - [GitHub Repo](https://github.com/google-research/t5x)
+  - [Hugging Face FLAN-T5 Docs (Similar to T5) ](https://huggingface.co/docs/transformers/model_doc/t5)
+
+# Usage
+
+Find below some example scripts on how to use the model in `transformers`:
+
+## Using the Pytorch model
+
+### Running the model on a CPU
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+### Running the model on a GPU
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+# pip install accelerate
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto")
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+### Running the model on a GPU using different precisions
+
+#### FP16
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+# pip install accelerate
+import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto", torch_dtype=torch.float16)
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+#### INT8
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+# pip install bitsandbytes accelerate
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto", load_in_8bit=True)
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+# Uses
+
+## Direct Use and Downstream Use
+
+The authors write in [the original paper's model card](https://arxiv.org/pdf/2210.11416.pdf) that: 
+
+> The primary use is research on language models, including: research on zero-shot NLP tasks and in-context few-shot learning NLP tasks, such as reasoning, and question answering; advancing fairness and safety research, and understanding limitations of current large language models
+
+See the [research paper](https://arxiv.org/pdf/2210.11416.pdf) for further details.
+
+## Out-of-Scope Use
+
+More information needed.
+
+# Bias, Risks, and Limitations
+
+The information below in this section are copied from the model's [official model card](https://arxiv.org/pdf/2210.11416.pdf):
+
+> Language models, including Flan-T5, can potentially be used for language generation in a harmful way, according to Rae et al. (2021). Flan-T5 should not be used directly in any application, without a prior assessment of safety and fairness concerns specific to the application.
+
+## Ethical considerations and risks
+
+> Flan-T5 is fine-tuned on a large corpus of text data that was not filtered for explicit content or assessed for existing biases. As a result the model itself is potentially vulnerable to generating equivalently inappropriate content or replicating inherent biases in the underlying data.
+
+## Known Limitations
+
+> Flan-T5 has not been tested in real world applications.
+
+## Sensitive Use:
+
+> Flan-T5 should not be applied for any unacceptable use cases, e.g., generation of abusive speech.
+
+# Training Details
+
+## Training Data
+
+The model was trained on a mixture of tasks, that includes the tasks described in the table below (from the original paper, figure 2):
+
+![table.png](https://s3.amazonaws.com/moonup/production/uploads/1666363265279-62441d1d9fdefb55a0b7d12c.png)
+
+
+## Training Procedure
+
+According to the model card from the [original paper](https://arxiv.org/pdf/2210.11416.pdf):
+
+> These models are based on pretrained T5 (Raffel et al., 2020) and fine-tuned with instructions for better zero-shot and few-shot performance. There is one fine-tuned Flan model per T5 model size.
+
+The model has been trained on TPU v3 or TPU v4 pods, using [`t5x`](https://github.com/google-research/t5x) codebase together with [`jax`](https://github.com/google/jax).
+
+
+# Evaluation
+
+## Testing Data, Factors & Metrics
+
+The authors evaluated the model on various tasks covering several languages (1836 in total). See the table below for some quantitative evaluation:
+![image.png](https://s3.amazonaws.com/moonup/production/uploads/1668072995230-62441d1d9fdefb55a0b7d12c.png)
+For full details, please check the [research paper](https://arxiv.org/pdf/2210.11416.pdf).
+
+## Results 
+
+For full results for FLAN-T5-Large, see the [research paper](https://arxiv.org/pdf/2210.11416.pdf), Table 3.
+
+# Environmental Impact
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** Google Cloud TPU Pods - TPU v3 or TPU v4  | Number of chips ≥ 4.
+- **Hours used:** More information needed
+- **Cloud Provider:** GCP
+- **Compute Region:** More information needed
+- **Carbon Emitted:** More information needed
+
+# Citation
+
+**BibTeX:**
+
+```bibtex
+@misc{https://doi.org/10.48550/arxiv.2210.11416,
+  doi = {10.48550/ARXIV.2210.11416},
+  
+  url = {https://arxiv.org/abs/2210.11416},
+  
+  author = {Chung, Hyung Won and Hou, Le and Longpre, Shayne and Zoph, Barret and Tay, Yi and Fedus, William and Li, Eric and Wang, Xuezhi and Dehghani, Mostafa and Brahma, Siddhartha and Webson, Albert and Gu, Shixiang Shane and Dai, Zhuyun and Suzgun, Mirac and Chen, Xinyun and Chowdhery, Aakanksha and Narang, Sharan and Mishra, Gaurav and Yu, Adams and Zhao, Vincent and Huang, Yanping and Dai, Andrew and Yu, Hongkun and Petrov, Slav and Chi, Ed H. and Dean, Jeff and Devlin, Jacob and Roberts, Adam and Zhou, Denny and Le, Quoc V. and Wei, Jason},
+  
+  keywords = {Machine Learning (cs.LG), Computation and Language (cs.CL), FOS: Computer and information sciences, FOS: Computer and information sciences},
+  
+  title = {Scaling Instruction-Finetuned Language Models},
+  
+  publisher = {arXiv},
+  
+  year = {2022},
+  
+  copyright = {Creative Commons Attribution 4.0 International}
+}
+```

Looped-DiT-B-16/text_encoder/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 2816,
+  "d_kv": 64,
+  "d_model": 1024,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "n_positions": 512,
+  "num_decoder_layers": 24,
+  "num_heads": 16,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.23.1",
+  "use_cache": true,
+  "vocab_size": 32128
+}

Looped-DiT-B-16/text_encoder/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "decoder_start_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 0,
+  "transformers_version": "4.27.0.dev0"
+}

Looped-DiT-B-16/text_encoder/model.safetensors

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

Looped-DiT-B-16/text_encoder/special_tokens_map.json

@@ -0,0 +1,107 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-16/text_encoder/spiece.model

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

Looped-DiT-B-16/text_encoder/tokenizer_config.json

@@ -0,0 +1,113 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "model_max_length": 256,
+  "name_or_path": "google/t5-v1_1-large",
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "special_tokens_map_file": "/home/younes_huggingface_co/.cache/huggingface/hub/models--google--t5-v1_1-large/snapshots/314bc112b191ec17b625ba81438dc73d6c23659d/special_tokens_map.json",
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-16/tokenizer/special_tokens_map.json

@@ -0,0 +1,107 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-16/tokenizer/spiece.model

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

Looped-DiT-B-16/tokenizer/tokenizer_config.json

@@ -0,0 +1,113 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "model_max_length": 256,
+  "name_or_path": "google/t5-v1_1-large",
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "special_tokens_map_file": "/home/younes_huggingface_co/.cache/huggingface/hub/models--google--t5-v1_1-large/snapshots/314bc112b191ec17b625ba81438dc73d6c23659d/special_tokens_map.json",
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-16/transformer/config.json

@@ -0,0 +1,23 @@
+{
+  "_class_name": "LoopedDiTTransformer2DModel",
+  "_diffusers_version": "0.39.0",
+  "head_dim": 64,
+  "hidden_size": 768,
+  "image_size": 512,
+  "in_channels": 3,
+  "loop_split": [
+    6,
+    5,
+    6
+  ],
+  "mlp_ratio": 2.6667,
+  "num_heads": 12,
+  "num_loops": 4,
+  "patch_size": 16,
+  "pca_channels": 128,
+  "share_loop_weights": true,
+  "text_dim": 1024,
+  "text_preamble_depth": 2,
+  "use_attn_gate": false,
+  "use_xsa": true
+}

Looped-DiT-B-16/transformer/diffusion_pytorch_model.safetensors

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

Looped-DiT-B-16/transformer/transformer_looped_dit.py

@@ -0,0 +1,417 @@
+"""Looped-DiT denoiser as a diffusers model.
+
+The MiniT2I denoiser, a pixel-space variant of MMDiT (patchified image tokens
+and T5 text tokens with modality-specific weights and joint attention, no
+timestep conditioning), whose double-stream blocks are split into three stages:
+
+    pre-loop   A   blocks[:pre]              run once
+    looped     B   the next `core` blocks    run N times
+    post-loop  C   the last `post` blocks    run once, followed by the head
+
+    h_0 = A(x),   h_r = B(h_{r-1}),   x0_hat(r) = C(h_r),   r = 1..N
+
+With shared weights (the default) B is one set of blocks reused N times, so the
+loop adds depth but no weights. Any loop state h_r can be decoded through C:
+deep supervision trains those intermediate exits, and inference can run with a
+different loop depth. XSA and the attention gate act on the looped blocks only.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.models.modeling_utils import ModelMixin
+from torch import nn
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) * self.weight
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        hidden_dim = math.ceil(hidden_dim / 8) * 8
+        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)
+        for layer in (self.w1, self.w3, self.w2):
+            nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+# ---------------------------------------------------------------------------
+# Rotary position embeddings: 1D over text positions, 2D over the patch grid.
+# ---------------------------------------------------------------------------
+
+_ROPE_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+
+
+def _autocast_state() -> tuple:
+    try:
+        enabled = torch.is_autocast_enabled("cuda")
+        return enabled, torch.get_autocast_dtype("cuda") if enabled else None
+    except TypeError:  # torch < 2.4
+        enabled = torch.is_autocast_enabled()
+        return enabled, torch.get_autocast_gpu_dtype() if enabled else None
+
+
+def _rope_tables(grid: int | None, n: int, d: int, device, dtype, theta: float = 10000.0):
+    """cos/sin tables for 1D (grid=None) or 2D rotary embeddings.
+
+    Tables are built under whatever autocast state the caller runs in (under bf16
+    autocast the angle products are computed in bf16, which is what the models
+    were trained with), so that state is part of the cache key.
+    """
+    key = (grid, n, d, str(device), dtype, _autocast_state())
+    if key not in _ROPE_CACHE:
+        if grid is None:
+            inv = 1.0 / (theta ** (torch.arange(0, d, 2, device=device, dtype=torch.float32) / d))
+            pos = torch.arange(n, device=device, dtype=torch.float32)
+            angles = torch.einsum("n,f->nf", pos, inv)
+            angles = torch.cat([angles, angles], dim=-1)
+        else:
+            half = d // 2
+            inv = 1.0 / (theta ** (torch.arange(0, half, 2, device=device, dtype=torch.float32) / half))
+            freqs = torch.einsum("n,f->nf", torch.arange(grid, device=device, dtype=torch.float32), inv)
+            f_h, f_w = torch.broadcast_tensors(freqs[:, None, :], freqs[None, :, :])
+            angles = torch.cat([f_h, f_w], dim=-1)
+            angles = torch.cat([angles, angles], dim=-1).reshape(n, d)
+        _ROPE_CACHE[key] = (angles.cos()[None, None].to(dtype), angles.sin()[None, None].to(dtype))
+    return _ROPE_CACHE[key]
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x: torch.Tensor, grid: int | None = None) -> torch.Tensor:
+    """x: [batch, heads, tokens, head_dim]."""
+    cos, sin = _rope_tables(grid, x.shape[2], x.shape[3], x.device, x.dtype)
+    return x * cos + rotate_half(x) * sin
+
+
+def sincos_2d(dim: int, grid: int) -> torch.Tensor:
+    y, x = torch.meshgrid(torch.arange(grid), torch.arange(grid), indexing="ij")
+    omega = 1.0 / (10000 ** (torch.arange(dim // 4, dtype=torch.float32) / (dim // 4)))
+    out_y = torch.einsum("n,d->nd", y.flatten().float(), omega)
+    out_x = torch.einsum("n,d->nd", x.flatten().float(), omega)
+    return torch.cat([out_x.sin(), out_x.cos(), out_y.sin(), out_y.cos()], dim=1)
+
+
+def exclusive_self_attention(out: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """XSA (Zhai, 2026: https://arxiv.org/abs/2603.09078):
+    remove from each token's attention output the component along that token's
+    own value vector, so attention only writes content from other tokens.
+    `out` and `v` are token-aligned, heads first."""
+    v_hat = F.normalize(v.float(), dim=-1)
+    out_f = out.float()
+    return (out_f - (out_f * v_hat).sum(dim=-1, keepdim=True) * v_hat).to(out.dtype)
+
+
+# ---------------------------------------------------------------------------
+# Blocks
+# ---------------------------------------------------------------------------
+
+
+class PatchEmbed(nn.Module):
+    """Two-stage patch embedding: a low-rank patch projection, then a 1x1 conv."""
+
+    def __init__(self, patch_size: int, in_channels: int, hidden_size: int, bottleneck: int):
+        super().__init__()
+        self.proj1 = nn.Conv2d(in_channels, bottleneck, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(bottleneck, hidden_size, kernel_size=1, bias=True)
+        nn.init.xavier_uniform_(self.proj1.weight)
+        nn.init.xavier_uniform_(self.proj2.weight)
+        nn.init.zeros_(self.proj2.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj2(self.proj1(x)).flatten(2).transpose(1, 2)
+
+
+class TextBlock(nn.Module):
+    """Text-only transformer block that refines the T5 tokens before the joint blocks."""
+
+    def __init__(self, hidden_size: int, num_heads: int, head_dim: int, mlp_ratio: float):
+        super().__init__()
+        self.num_heads, self.head_dim = num_heads, head_dim
+        self.norm1 = RMSNorm(hidden_size)
+        self.norm2 = RMSNorm(hidden_size)
+        self.qkv = nn.Linear(hidden_size, num_heads * head_dim * 3)
+        self.proj = nn.Linear(num_heads * head_dim, hidden_size)
+        self.mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+    def forward(self, txt: torch.Tensor) -> torch.Tensor:
+        b, n, _ = txt.shape
+        q, k, v = self.qkv(self.norm1(txt)).view(b, n, 3, self.num_heads, self.head_dim).unbind(2)
+        q, k, v = (z.transpose(1, 2) for z in (self.q_norm(q), self.k_norm(k), v))
+        out = F.scaled_dot_product_attention(apply_rope(q), apply_rope(k), v, scale=self.head_dim**-0.5)
+        txt = txt + self.proj(out.transpose(1, 2).reshape(b, n, -1))
+        return txt + self.mlp(self.norm2(txt))
+
+
+class DoubleStreamBlock(nn.Module):
+    """MMDiT block: separate image/text weights, one joint attention over both.
+
+    `use_xsa` / `use_attn_gate` turn on self-modulating attention (set only for the
+    looped blocks). The attention gate (Qiu et al., 2026:
+    https://arxiv.org/abs/2505.06708) is head-wise:
+    y_i <- y_i * sigmoid(W_g u_i + b_g), with u_i the block's normed input.
+    
+    `update_text=False` skips the text-stream update, for the last block, whose
+    text output is never read.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        head_dim: int,
+        mlp_ratio: float,
+        grid: int,
+        use_xsa: bool = False,
+        use_attn_gate: bool = False,
+        update_text: bool = True,
+    ):
+        super().__init__()
+        self.num_heads, self.head_dim, self.grid = num_heads, head_dim, grid
+        self.use_xsa, self.use_attn_gate, self.update_text = use_xsa, use_attn_gate, update_text
+        inner = num_heads * head_dim
+        self.img_norm1 = RMSNorm(hidden_size)
+        self.img_norm2 = RMSNorm(hidden_size)
+        self.txt_norm1 = RMSNorm(hidden_size)
+        self.txt_norm2 = RMSNorm(hidden_size)
+        self.img_qkv = nn.Linear(hidden_size, inner * 3)
+        self.txt_qkv = nn.Linear(hidden_size, inner * 3)
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+        self.img_proj = nn.Linear(inner, hidden_size)
+        self.txt_proj = nn.Linear(inner, hidden_size)
+        if use_attn_gate:
+            # Zero bias: the gates start half open on average.
+            self.img_gate = nn.Linear(hidden_size, num_heads)
+            self.txt_gate = nn.Linear(hidden_size, num_heads)
+            nn.init.zeros_(self.img_gate.bias)
+            nn.init.zeros_(self.txt_gate.bias)
+        self.img_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        self.txt_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+
+    def forward(self, img: torch.Tensor, txt: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        b, li, _ = img.shape
+        lt = txt.shape[1]
+        img_n, txt_n = self.img_norm1(img), self.txt_norm1(txt)
+        qi, ki, vi = self.img_qkv(img_n).view(b, li, 3, self.num_heads, self.head_dim).unbind(2)
+        qt, kt, vt = self.txt_qkv(txt_n).view(b, lt, 3, self.num_heads, self.head_dim).unbind(2)
+        # Joint sequence [text; image], heads first.
+        q = torch.cat([qt, qi], dim=1).transpose(1, 2)
+        k = torch.cat([kt, ki], dim=1).transpose(1, 2)
+        v = torch.cat([vt, vi], dim=1).transpose(1, 2)
+        q = torch.cat([apply_rope(self.q_norm(q[:, :, :lt])), apply_rope(self.q_norm(q[:, :, lt:]), self.grid)], dim=2)
+        k = torch.cat([apply_rope(self.k_norm(k[:, :, :lt])), apply_rope(self.k_norm(k[:, :, lt:]), self.grid)], dim=2)
+        out = F.scaled_dot_product_attention(q, k, v, scale=self.head_dim**-0.5)
+        if self.use_xsa:
+            out = exclusive_self_attention(out, v)
+        out = out.transpose(1, 2)  # [b, tokens, heads, head_dim]
+        out_t, out_i = out[:, :lt], out[:, lt:]
+        if self.use_attn_gate:
+            out_i = out_i * torch.sigmoid(self.img_gate(img_n)).unsqueeze(-1)
+        img = img + self.img_proj(out_i.reshape(b, li, -1))
+        img = img + self.img_mlp(self.img_norm2(img))
+        if self.update_text:
+            if self.use_attn_gate:
+                out_t = out_t * torch.sigmoid(self.txt_gate(txt_n)).unsqueeze(-1)
+            txt = txt + self.txt_proj(out_t.reshape(b, lt, -1))
+            txt = txt + self.txt_mlp(self.txt_norm2(txt))
+        return img, txt
+
+
+# ---------------------------------------------------------------------------
+# Looped MMDiT
+# ---------------------------------------------------------------------------
+
+
+class LoopedDiTTransformer2DModel(ModelMixin, ConfigMixin):
+    """Predicts the clean image x0 from a noisy image and T5 text embeddings.
+
+    This is a diffusers [`ModelMixin`]: `save_pretrained` / `from_pretrained` round-trip the
+    architecture in `config.json`. `loop_split` is `(pre, core, post)` and `num_loops` is the
+    trained loop depth (`1` is the MiniT2I model without looping). With `share_loop_weights=False`
+    every pass gets its own copy of the core blocks: the compute-matched "deeper" baseline with
+    the same exits.
+
+    The forward pass does not take a timestep. Flow-matching time is applied by the pipeline when
+    it converts the x0 prediction into a velocity.
+    """
+
+    def __init__(
+        self,
+        image_size: int = 512,
+        patch_size: int = 32,
+        in_channels: int = 3,
+        hidden_size: int = 768,
+        num_heads: int = 12,
+        head_dim: int = 64,
+        mlp_ratio: float = 2.6667,
+        pca_channels: int = 128,
+        text_dim: int = 1024,
+        text_preamble_depth: int = 2,
+        loop_split: tuple[int, int, int] | list[int] = (6, 5, 6),
+        num_loops: int = 4,
+        share_loop_weights: bool = True,
+        use_xsa: bool = False,
+        use_attn_gate: bool = False,
+    ):
+        super().__init__()
+        loop_split = [int(n) for n in loop_split]
+        num_loops = int(num_loops)
+        share_loop_weights = bool(share_loop_weights)
+        use_xsa = bool(use_xsa)
+        use_attn_gate = bool(use_attn_gate)
+        if len(loop_split) != 3 or min(loop_split) < 1:
+            raise ValueError(f"loop_split must be three positive block counts (pre, core, post), got {loop_split}")
+        if num_loops < 1:
+            raise ValueError(f"num_loops must be >= 1, got {num_loops}")
+        # Lists (not tuples) so config.json stays valid JSON.
+        self.register_to_config(
+            image_size=image_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            mlp_ratio=mlp_ratio,
+            pca_channels=pca_channels,
+            text_dim=text_dim,
+            text_preamble_depth=text_preamble_depth,
+            loop_split=loop_split,
+            num_loops=num_loops,
+            share_loop_weights=share_loop_weights,
+            use_xsa=use_xsa,
+            use_attn_gate=use_attn_gate,
+        )
+        pre, core, post = loop_split
+        self.patch_size, self.in_channels = patch_size, in_channels
+        self.grid = image_size // patch_size
+        self.pre, self.core, self.post = pre, core, post
+        self.num_loops = num_loops
+        self.share_loop_weights = share_loop_weights
+
+        self.img_embed = PatchEmbed(patch_size, in_channels, hidden_size, pca_channels)
+        self.txt_embed = nn.Linear(text_dim, hidden_size, bias=False)
+        # Replaces the T5 embedding at padded prompt positions (and everywhere
+        # for the unconditional branch of classifier-free guidance).
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, text_dim))
+        nn.init.normal_(self.mask_token, std=0.02)
+        # MiniT2I's timestep and pooled-text embedders. The model has no timestep
+        # conditioning and never uses them; they are kept (frozen, see below) so that
+        # the model and its checkpoints match MiniT2I and the paper.
+        self.t_embed = nn.ModuleDict(
+            {"mlp": nn.Sequential(nn.Linear(256, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size))}
+        )
+        for layer in (self.t_embed.mlp[0], self.t_embed.mlp[2]):
+            nn.init.normal_(layer.weight, std=0.02)
+            nn.init.zeros_(layer.bias)
+        self.pooled_embed = nn.Linear(text_dim, hidden_size, bias=False)
+        self.register_buffer("pos_embed", sincos_2d(hidden_size, self.grid)[None], persistent=False)
+        self.txt_blocks = nn.ModuleList(
+            TextBlock(hidden_size, num_heads, head_dim, mlp_ratio) for _ in range(text_preamble_depth)
+        )
+        looped = core if self.share_loop_weights else core * self.num_loops
+        depth = pre + looped + post
+        self.blocks = nn.ModuleList(
+            DoubleStreamBlock(
+                hidden_size,
+                num_heads,
+                head_dim,
+                mlp_ratio,
+                self.grid,
+                use_xsa=use_xsa and pre <= i < pre + looped,
+                use_attn_gate=use_attn_gate and pre <= i < pre + looped,
+                update_text=i < depth - 1,
+            )
+            for i in range(depth)
+        )
+        self.final_norm = RMSNorm(hidden_size)
+        self.final = nn.Linear(hidden_size, patch_size * patch_size * in_channels)
+        nn.init.zeros_(self.final.weight)
+        nn.init.zeros_(self.final.bias)
+        # Frozen because they never receive a gradient: the unused embedders and the
+        # text-stream update of the last block (whose text output is never read).
+        last = self.blocks[-1]
+        for module in (self.t_embed, self.pooled_embed, last.txt_norm2, last.txt_proj, last.txt_mlp):
+            module.requires_grad_(False)
+
+    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
+        b, n, _ = x.shape
+        p, c, g = self.patch_size, self.in_channels, int(n**0.5)
+        x = x.view(b, g, g, p, p, c).permute(0, 5, 1, 3, 2, 4).contiguous()
+        return x.view(b, c, g * p, g * p)
+
+    def loop_blocks(self, r: int) -> nn.ModuleList:
+        """Core blocks run on loop pass r (1-based)."""
+        start = self.pre if self.share_loop_weights else self.pre + (r - 1) * self.core
+        return self.blocks[start : start + self.core]
+
+    def decode(self, img: torch.Tensor, txt: torch.Tensor) -> torch.Tensor:
+        """Post-loop blocks and output head: a loop state -> x0 prediction."""
+        for block in self.blocks[len(self.blocks) - self.post :]:
+            img, txt = block(img, txt)
+        return self.unpatchify(self.final(self.final_norm(img))).float()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text: torch.Tensor,
+        text_mask: torch.Tensor,
+        num_loops: int | None = None,
+        exit_loops: tuple[int, ...] = (),
+    ) -> torch.Tensor | tuple[torch.Tensor, dict[int, torch.Tensor]]:
+        """x: noisy images [B, C, H, W]; text: T5 states [B, L, text_dim];
+        text_mask: [B, L], 1 for prompt tokens (all 0 = unconditional).
+
+        num_loops overrides the loop depth at inference. exit_loops lists
+        intermediate depths r < num_loops to decode as well; the call then
+        returns (final prediction, {r: prediction after r loops}).
+        """
+        n = self.num_loops if num_loops is None else int(num_loops)
+        if n < 1 or (not self.share_loop_weights and n > self.num_loops):
+            raise ValueError(f"num_loops={n} is not available for this model (trained with {self.num_loops})")
+        exits = sorted({int(r) for r in exit_loops})
+        if any(not 1 <= r < n for r in exits):
+            raise ValueError(f"exit_loops must lie in [1, {n}), got {exits}")
+
+        text = torch.where(text_mask.to(torch.bool)[:, :, None], text, self.mask_token.to(text.dtype))
+        img = self.img_embed(x) + self.pos_embed.to(device=x.device, dtype=x.dtype)
+        txt = self.txt_embed(text)
+        for block in self.txt_blocks:
+            txt = block(txt)
+        for block in self.blocks[: self.pre]:
+            img, txt = block(img, txt)
+        states = {}
+        for r in range(1, n + 1):
+            for block in self.loop_blocks(r):
+                img, txt = block(img, txt)
+            if r in exits:
+                states[r] = (img, txt)
+        out = self.decode(img, txt)
+        if not exits:
+            return out
+        return out, {r: self.decode(*states[r]) for r in exits}
+
+
+# Training imports this name. The diffusers class name is the one stored in checkpoints.
+LoopedMMDiT = LoopedDiTTransformer2DModel

Looped-DiT-B-32/model_index.json

@@ -0,0 +1,23 @@
+{
+  "_class_name": [
+    "pipeline",
+    "LoopedDiTPipeline"
+  ],
+  "_diffusers_version": "0.39.0",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "T5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "T5Tokenizer"
+  ],
+  "transformer": [
+    "transformer_looped_dit",
+    "LoopedDiTTransformer2DModel"
+  ]
+}

Looped-DiT-B-32/pipeline.py

@@ -0,0 +1,763 @@
+# Copyright 2025 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.
+
+import inspect
+from typing import Any, Callable
+
+import torch
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import deprecate, is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+from transformers import AutoTokenizer, T5EncoderModel
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+# Training clamps the flow-matching denominator so the loss stays finite at t -> 1.
+VELOCITY_DENOM_MIN = 0.05
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from diffusers import DiffusionPipeline
+
+        >>> model_dir = Path("checkpoints/looped-dit-b16").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,
+        ... ).to("cuda")
+
+        >>> image = pipe(
+        ...     "a red cube on top of a blue sphere",
+        ...     num_inference_steps=100,
+        ...     guidance_scale=6.0,
+        ...     num_loops=4,
+        ...     generator=torch.Generator(device="cuda").manual_seed(0),
+        ... ).images[0]
+        >>> image.save("sample.png")
+
+        >>> # Hugging Face Hub style model id: UserID/RepoID
+        >>> # RepoID is usually like "modelname-diffusers"
+        >>> # Example: "your-user/Looped-DiT-diffusers"
+        ```
+"""
+
+
+def paper_euler_sigmas(num_inference_steps: int) -> list[float]:
+    r"""
+    Sigma grid of the training Euler sampler.
+
+    Training integrates flow time `t` from 0 (noise) to 1 (data) with
+    `torch.linspace(0, 1, steps + 1)`. Flow-match schedulers step in sigma
+    `1 - t` and append the terminal 0 themselves, so the returned list omits that 0.
+
+    Args:
+        num_inference_steps (`int`):
+            Number of Euler steps. Must be positive.
+
+    Returns:
+        `list[float]`: `num_inference_steps` sigmas starting at 1 and ending at `1 / steps`.
+    """
+    if num_inference_steps <= 0:
+        raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}.")
+    flow_time = torch.linspace(0.0, 1.0, num_inference_steps + 1)
+    return (1.0 - flow_time)[:-1].tolist()
+
+
+class LoopedDiTPipeline(DiffusionPipeline):
+    r"""
+    Text-to-image pipeline for Looped-DiT.
+
+    Looped-DiT denoises directly in RGB pixel space (no VAE). The transformer predicts the clean
+    image `x0`. This pipeline converts that prediction to a flow-matching velocity and integrates it
+    with a diffusers scheduler. The default scheduler is [`FlowMatchEulerDiscreteScheduler`] on the
+    same uniform grid as the paper (100 steps, shift 1). Any [`KarrasDiffusionSchedulers`] instance
+    can be assigned to `pipe.scheduler` without other code changes.
+
+    Classifier-free guidance uses the training null condition: an all-zero text mask, which the
+    denoiser replaces with its mask token. There is no separate negative-prompt encoder.
+
+    The pipeline inherits from [`DiffusionPipeline`]. Check the superclass documentation for the
+    generic methods (download, save, device placement, CPU offload).
+
+    Args:
+        transformer ([`ModelMixin`]):
+            Looped-DiT denoiser (`LoopedDiTTransformer2DModel`) that predicts `x0` in pixel space.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`] or [`KarrasDiffusionSchedulers`]):
+            Scheduler used to step the flow. The paper setting is
+            `FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=1.0)`.
+        tokenizer ([`~transformers.AutoTokenizer`], *optional*):
+            Tokenizer for the frozen text encoder. Loaded from `text_encoder_name` when missing.
+        text_encoder ([`~transformers.T5EncoderModel`], *optional*):
+            Frozen FLAN-T5 encoder. Loaded from `text_encoder_name` when missing.
+            `DiffusionPipeline.from_pretrained(..., torch_dtype=torch.bfloat16)` keeps this encoder in bf16 with the denoiser.
+        text_encoder_name (`str`, defaults to `"google/flan-t5-large"`):
+            Hub id or local path used when `tokenizer` / `text_encoder` are not passed.
+        prompt_length (`int`, *optional*):
+            Token length prompts are padded or truncated to. Defaults to `tokenizer.model_max_length`.
+        noise_scale (`float`, defaults to 2.0):
+            Standard deviation of the initial Gaussian, matching the training noise scale.
+        default_num_inference_steps (`int`, defaults to 100):
+            Step count used when `__call__` does not pass `num_inference_steps`.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer"
+    _optional_components = ["tokenizer", "text_encoder"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "prompt_attention_mask"]
+
+    def __init__(
+        self,
+        transformer: ModelMixin,
+        scheduler: KarrasDiffusionSchedulers | SchedulerMixin,
+        tokenizer: Any | None = None,
+        text_encoder: T5EncoderModel | None = None,
+        text_encoder_name: str = "google/flan-t5-large",
+        prompt_length: int | None = None,
+        noise_scale: float = 2.0,
+        default_num_inference_steps: int = 100,
+    ):
+        super().__init__()
+        if prompt_length is None and tokenizer is not None:
+            prompt_length = int(getattr(tokenizer, "model_max_length", 256))
+        if prompt_length is None:
+            prompt_length = 256
+        if scheduler is None:
+            scheduler = self._default_scheduler()
+        if noise_scale <= 0:
+            raise ValueError(f"`noise_scale` must be positive, got {noise_scale}.")
+        if prompt_length < 1:
+            raise ValueError(f"`prompt_length` must be positive, got {prompt_length}.")
+        if default_num_inference_steps < 1:
+            raise ValueError(f"`default_num_inference_steps` must be positive, got {default_num_inference_steps}.")
+
+        self.register_modules(
+            transformer=transformer,
+            scheduler=scheduler,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+        )
+        self.register_to_config(
+            text_encoder_name=text_encoder_name,
+            prompt_length=int(prompt_length),
+            noise_scale=float(noise_scale),
+            default_num_inference_steps=int(default_num_inference_steps),
+        )
+
+    @staticmethod
+    def _default_scheduler() -> FlowMatchEulerDiscreteScheduler:
+        r"""
+        Build the paper's Euler scheduler.
+
+        Returns:
+            [`FlowMatchEulerDiscreteScheduler`]: 1000 training timesteps, shift 1, deterministic.
+        """
+        kwargs: dict[str, Any] = {"num_train_timesteps": 1000, "shift": 1.0}
+        if "stochastic_sampling" in inspect.signature(FlowMatchEulerDiscreteScheduler.__init__).parameters:
+            kwargs["stochastic_sampling"] = False
+        return FlowMatchEulerDiscreteScheduler(**kwargs)
+
+    def _encode_prompt(
+        self,
+        prompt: str | list[str] | None,
+        device: torch.device,
+        num_images_per_prompt: int,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Deprecated alias of [`~LoopedDiTPipeline.encode_prompt`].
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt or prompts to tokenize and encode.
+            device (`torch.device`):
+                Device of the returned tensors.
+            num_images_per_prompt (`int`):
+                How many times to repeat each prompt embedding.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Already encoded prompts of shape `(batch, sequence, text_dim)`.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask of shape `(batch, sequence)` with 1 on real tokens. Required with `prompt_embeds`
+                only when padding should be replaced by the mask token; otherwise a mask of ones is used.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`: Prompt embeddings and the attention mask.
+        """
+        deprecation_message = (
+            "`_encode_prompt()` is deprecated and will be removed in a future version. Use `encode_prompt()` instead."
+        )
+        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
+        return self.encode_prompt(prompt, device, num_images_per_prompt, prompt_embeds, prompt_attention_mask)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str] | None,
+        device: torch.device,
+        num_images_per_prompt: int,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Encode prompts with the frozen FLAN-T5 encoder.
+
+        Prompts are padded or truncated to `config.prompt_length`. The unconditional branch of
+        classifier-free guidance is not encoded here: the denoiser builds it by zeroing this mask.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt or prompts to tokenize. Ignored when `prompt_embeds` is passed.
+            device (`torch.device`):
+                Device of the returned tensors.
+            num_images_per_prompt (`int`):
+                Number of times to repeat each encoded prompt along the batch dimension.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed embeddings of shape `(batch, sequence, text_dim)`. When set, `prompt` is ignored.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask of shape `(batch, sequence)`, 1 for tokens that should condition the model. When
+                `prompt_embeds` is set and this is omitted, every position is treated as a real token.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`:
+                Embeddings `(batch * num_images_per_prompt, sequence, text_dim)` and a mask of the same batch.
+        """
+        if num_images_per_prompt < 1:
+            raise ValueError(f"`num_images_per_prompt` must be >= 1, got {num_images_per_prompt}.")
+
+        if prompt_embeds is None:
+            if isinstance(prompt, str):
+                prompt = [prompt]
+            if self.tokenizer is None:
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    self.config.text_encoder_name, model_max_length=int(self.config.prompt_length)
+                )
+            if self.text_encoder is None:
+                self.text_encoder = T5EncoderModel.from_pretrained(self.config.text_encoder_name)
+                self.text_encoder.requires_grad_(False)
+            self.text_encoder.eval()
+            encoder_device = next(self.text_encoder.parameters()).device
+            if encoder_device != device:
+                self.text_encoder.to(device)
+            tokens = self.tokenizer(
+                prompt,
+                max_length=int(self.config.prompt_length),
+                padding="max_length",
+                truncation=True,
+                return_tensors="pt",
+            )
+            input_ids = tokens.input_ids.to(device)
+            prompt_attention_mask = tokens.attention_mask.to(device)
+            prompt_embeds = self.text_encoder(input_ids=input_ids, attention_mask=prompt_attention_mask).last_hidden_state
+        else:
+            prompt_embeds = prompt_embeds.to(device)
+            if prompt_attention_mask is None:
+                prompt_attention_mask = torch.ones(
+                    prompt_embeds.shape[:2], device=device, dtype=torch.long
+                )
+            else:
+                prompt_attention_mask = prompt_attention_mask.to(device)
+            if prompt_embeds.shape[0] != prompt_attention_mask.shape[0]:
+                raise ValueError(
+                    "`prompt_embeds` and `prompt_attention_mask` must have the same batch size, got "
+                    f"{prompt_embeds.shape[0]} and {prompt_attention_mask.shape[0]}."
+                )
+
+        if num_images_per_prompt != 1:
+            prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+            prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_images_per_prompt, dim=0)
+        return prompt_embeds, prompt_attention_mask
+
+    def prepare_extra_step_kwargs(
+        self, generator: torch.Generator | list[torch.Generator] | None, eta: float
+    ) -> dict[str, Any]:
+        r"""
+        Extra arguments forwarded to `scheduler.step`, depending on what that method accepts.
+
+        Args:
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Generator passed through when the scheduler step samples noise.
+            eta (`float`):
+                DDIM eta in `[0, 1]`. Ignored by schedulers whose `step` has no `eta` argument.
+
+        Returns:
+            `dict`: Keyword arguments for `scheduler.step`.
+        """
+        extra_step_kwargs: dict[str, Any] = {}
+        step_params = set(inspect.signature(self.scheduler.step).parameters.keys())
+        if "eta" in step_params:
+            extra_step_kwargs["eta"] = eta
+        if "generator" in step_params:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt: str | list[str] | None,
+        height: int,
+        width: int,
+        callback_steps: int | None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+        num_inference_steps: int = 100,
+        guidance_scale: float = 6.0,
+        num_loops: int | None = None,
+        output_type: str = "pil",
+    ) -> None:
+        r"""
+        Validate generation arguments and raise `ValueError` or `TypeError` on misuse.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt text. Mutually exclusive with `prompt_embeds`.
+            height (`int`):
+                Output height in pixels. Must equal the transformer's trained `image_size`.
+            width (`int`):
+                Output width in pixels. Must equal the transformer's trained `image_size`.
+            callback_steps (`int`, *optional*):
+                Deprecated callback period. When set, it must be a positive integer.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed text embeddings. Required when `prompt` is omitted.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask paired with `prompt_embeds`.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names the step callback may read. Each name must be listed on
+                `_callback_tensor_inputs`.
+            num_inference_steps (`int`):
+                Denoising steps. Must be positive.
+            guidance_scale (`float`):
+                Classifier-free guidance scale. Must be finite. `1` disables guidance.
+            num_loops (`int`, *optional*):
+                Loop depth. `None` uses the depth stored on the transformer. Otherwise `>= 1`, and
+                untied models cannot exceed the trained depth.
+            output_type (`str`):
+                One of `"pil"`, `"np"`, `"pt"`, or `"latent"`.
+        """
+        image_size = int(self.transformer.config.image_size)
+        patch_size = int(self.transformer.config.patch_size)
+        if height != image_size or width != image_size:
+            raise ValueError(
+                f"Looped-DiT uses a fixed positional grid of {image_size}x{image_size} "
+                f"(patch size {patch_size}). Got height={height}, width={width}."
+            )
+        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)}.")
+
+        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type {type(callback_steps)}."
+            )
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            key in self._callback_tensor_inputs for key in callback_on_step_end_tensor_inputs
+        ):
+            unexpected = [key for key in callback_on_step_end_tensor_inputs if key not in self._callback_tensor_inputs]
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {unexpected}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Pass only one of them.")
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
+        if prompt is not None and not isinstance(prompt, str) and not (
+            isinstance(prompt, list) and all(isinstance(item, str) for item in prompt)
+        ):
+            raise TypeError(f"`prompt` has to be a string or a list of strings, got {type(prompt)}.")
+        if prompt_embeds is not None and prompt_embeds.ndim != 3:
+            raise ValueError(f"`prompt_embeds` must have shape (batch, sequence, dim), got {tuple(prompt_embeds.shape)}.")
+        if prompt_attention_mask is not None and prompt_embeds is None:
+            raise ValueError("`prompt_attention_mask` was passed without `prompt_embeds`.")
+
+        if num_inference_steps <= 0:
+            raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}.")
+        if not torch.isfinite(torch.tensor(guidance_scale)):
+            raise ValueError(f"`guidance_scale` must be finite, got {guidance_scale}.")
+        if num_loops is not None:
+            if int(num_loops) < 1:
+                raise ValueError(f"`num_loops` must be >= 1, got {num_loops}.")
+            trained = int(self.transformer.config.num_loops)
+            if not bool(self.transformer.config.share_loop_weights) and int(num_loops) > trained:
+                raise ValueError(
+                    f"This checkpoint does not share loop weights, so `num_loops` cannot exceed the trained "
+                    f"depth {trained}. Got {num_loops}."
+                )
+        if output_type not in {"pil", "np", "pt", "latent"}:
+            raise ValueError(f"Unsupported `output_type` {output_type!r}. Choose from 'pil', 'np', 'pt', 'latent'.")
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        r"""
+        Sample the initial pixel-space noise, or validate a tensor the caller already sampled.
+
+        Looped-DiT has no VAE, so "latents" here are RGB images. Fresh noise is scaled by
+        `config.noise_scale` (2.0 in the paper). A provided `latents` tensor is not rescaled.
+
+        Args:
+            batch_size (`int`):
+                Number of images, including `num_images_per_prompt`.
+            num_channels (`int`):
+                Channel count. 3 for RGB.
+            height (`int`):
+                Image height in pixels.
+            width (`int`):
+                Image width in pixels.
+            dtype (`torch.dtype`):
+                Dtype of freshly sampled noise. The integration itself is accumulated in float32.
+            device (`torch.device`):
+                Device of the returned tensor.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Per-call RNG. A list must have length `batch_size`.
+            latents (`torch.Tensor`, *optional*):
+                Starting noise of shape `(batch_size, num_channels, height, width)`.
+
+        Returns:
+            `torch.Tensor`: Starting noise of shape `(batch_size, num_channels, height, width)`.
+        """
+        shape = (batch_size, num_channels, height, width)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You passed a list of {len(generator)} generators for a batch of {batch_size}. "
+                "The two lengths must match."
+            )
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = latents * float(self.config.noise_scale)
+        else:
+            latents = latents.to(device=device)
+            if tuple(latents.shape) != shape:
+                raise ValueError(f"`latents` shape {tuple(latents.shape)} does not match the expected {shape}.")
+        return latents
+
+    @property
+    def guidance_scale(self) -> float:
+        r"""
+        Classifier-free guidance scale of the call that is currently running.
+
+        Returns:
+            `float`: The scale set by the active `__call__`.
+        """
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self) -> bool:
+        r"""
+        Whether the active call runs a conditional and an unconditional forward.
+
+        Returns:
+            `bool`: True when `guidance_scale != 1`.
+        """
+        return self._guidance_scale != 1.0
+
+    @property
+    def num_timesteps(self) -> int:
+        r"""
+        Number of scheduler timesteps in the active call.
+
+        Returns:
+            `int`: Length of the timestep schedule.
+        """
+        return self._num_timesteps
+
+    @property
+    def interrupt(self) -> bool:
+        r"""
+        Whether the active denoising loop should skip remaining steps.
+
+        Returns:
+            `bool`: True after the caller sets `pipeline._interrupt = True`.
+        """
+        return self._interrupt
+
+    def _images_from_latents(self, latents: torch.Tensor, output_type: str) -> torch.Tensor | list[Image.Image] | Any:
+        r"""
+        Convert pixel-space samples in `[-1, 1]` to the requested output type.
+
+        Quantization matches the original sampler: `uint8(clamp(x, -1, 1) * 127.5 + 128)`.
+
+        Args:
+            latents (`torch.Tensor`):
+                Samples of shape `(batch, channels, height, width)` in model range `[-1, 1]`.
+            output_type (`str`):
+                `"latent"` returns `latents` unchanged. `"pt"` is float RGB in `[0, 1]`. `"np"` is
+                `uint8` HWC arrays. `"pil"` is a list of `PIL.Image.Image`.
+
+        Returns:
+            Images in the requested type.
+        """
+        if output_type == "latent":
+            return latents
+        images = (latents.float().clamp(-1, 1) * 127.5 + 128.0).clamp(0, 255).to(torch.uint8)
+        if output_type == "pt":
+            return images.float() / 255.0
+        arrays = images.permute(0, 2, 3, 1).cpu().numpy()
+        if output_type == "np":
+            return arrays
+        return [Image.fromarray(image) for image in arrays]
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        timesteps: list[int] | None = None,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 6.0,
+        num_images_per_prompt: int = 1,
+        num_loops: int | None = None,
+        eta: float = 0.0,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Callable[[int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        **kwargs,
+    ) -> ImagePipelineOutput | tuple:
+        r"""
+        Generate images from text prompts.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt or prompts to guide image generation. Required unless `prompt_embeds` is passed.
+            height (`int`, *optional*):
+                Image height in pixels. Defaults to the transformer's trained resolution (512).
+                Other resolutions are rejected: the positional embedding is a fixed grid.
+            width (`int`, *optional*):
+                Image width in pixels. Defaults to the trained resolution and must match `height`.
+            num_inference_steps (`int`, *optional*):
+                Denoising steps. Defaults to `config.default_num_inference_steps` (100).
+            timesteps (`list[int]`, *optional*):
+                Custom scheduler timesteps, descending. Mutually exclusive with `sigmas`. Ignored by the
+                paper Euler grid, which is selected only when both `timesteps` and `sigmas` are omitted
+                and the scheduler is [`FlowMatchEulerDiscreteScheduler`].
+            sigmas (`list[float]`, *optional*):
+                Custom sigmas passed to `scheduler.set_timesteps`. Mutually exclusive with `timesteps`.
+            guidance_scale (`float`, defaults to 6.0):
+                Classifier-free guidance scale from the paper. Guidance is on when this is not `1`.
+                The unconditional branch is an empty text mask, not a negative prompt.
+            num_images_per_prompt (`int`, defaults to 1):
+                How many images to sample for each prompt.
+            num_loops (`int`, *optional*):
+                How many times to run the shared middle blocks. `None` uses the trained depth. Other
+                depths work without retraining when loop weights are shared.
+            eta (`float`, defaults to 0.0):
+                DDIM eta. Ignored by the flow-match Euler scheduler.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG for the initial noise. `None` uses PyTorch's global generator, which is what
+                `torch.manual_seed` seeds.
+            latents (`torch.Tensor`, *optional*):
+                Initial noise `(batch, 3, height, width)`. Not multiplied by `noise_scale`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed FLAN-T5 states `(batch, sequence, text_dim)` in place of `prompt`.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Mask `(batch, sequence)` paired with `prompt_embeds`. 1 marks real tokens.
+            output_type (`str`, defaults to `"pil"`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"` (pixels in model range).
+            return_dict (`bool`, defaults to `True`):
+                Return [`ImagePipelineOutput`] when `True`, otherwise a one-tuple of images.
+            callback_on_step_end (`Callable` or `PipelineCallback`, *optional*):
+                Called as `callback_on_step_end(pipeline, step, timestep, callback_kwargs)` after each
+                scheduler step. Return a dict to replace tensors listed in
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`):
+                Tensor names passed to the step callback. Must be a subset of `_callback_tensor_inputs`.
+
+        Examples:
+
+        Returns:
+            [`ImagePipelineOutput`] or `tuple`:
+                When `return_dict` is `True`, [`ImagePipelineOutput`] with the images. Otherwise a tuple
+                whose first element is the images.
+        """
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        if kwargs:
+            raise TypeError(f"Unexpected arguments: {sorted(kwargs)}.")
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        image_size = int(self.transformer.config.image_size)
+        height = image_size if height is None else int(height)
+        width = image_size if width is None else int(width)
+        if num_inference_steps is None:
+            num_inference_steps = int(self.config.default_num_inference_steps)
+
+        # 1. Check inputs.
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            callback_steps,
+            prompt_embeds,
+            prompt_attention_mask,
+            callback_on_step_end_tensor_inputs,
+            num_inference_steps,
+            guidance_scale,
+            num_loops,
+            output_type,
+        )
+
+        self._guidance_scale = float(guidance_scale)
+        self._interrupt = False
+
+        # 2. Define call parameters.
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        device = self._execution_device
+
+        # 3. Encode input prompt.
+        prompt_embeds, prompt_attention_mask = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+        )
+        prompt_embeds = prompt_embeds.to(device=device, dtype=self.transformer.dtype)
+        prompt_attention_mask = prompt_attention_mask.to(device=device)
+
+        # 4. Prepare timesteps.
+        # The paper's Euler grid is the training linspace. Other schedulers keep their own spacing.
+        # A caller-supplied `timesteps` or `sigmas` always wins.
+        if timesteps is None and sigmas is None and isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+            sigmas = paper_euler_sigmas(num_inference_steps)
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, sigmas)
+        if getattr(self.scheduler.config, "stochastic_sampling", False):
+            raise ValueError(
+                "Looped-DiT's training sampler is deterministic. Set `stochastic_sampling=False` on "
+                "FlowMatchEulerDiscreteScheduler, or assign a different scheduler."
+            )
+
+        # 5. Prepare latent variables (pixel-space noise; there is no VAE).
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            int(self.transformer.config.in_channels),
+            height,
+            width,
+            self.transformer.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_train_timesteps = int(self.scheduler.config.num_train_timesteps)
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        # 7. Denoising loop.
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                model_latents = latents
+                if hasattr(self.scheduler, "scale_model_input"):
+                    model_latents = self.scheduler.scale_model_input(model_latents, t)
+                text = prompt_embeds
+                mask = prompt_attention_mask
+                if self.do_classifier_free_guidance:
+                    model_latents = torch.cat([model_latents, model_latents], dim=0)
+                    text = torch.cat([text, text], dim=0)
+                    mask = torch.cat([mask, torch.zeros_like(mask)], dim=0)
+
+                # fp32 latents with bf16 weights match the old sampler, which autocasts the forward.
+                amp_dtype = self.transformer.dtype
+                use_amp = model_latents.is_cuda and amp_dtype in (torch.float16, torch.bfloat16)
+                if not use_amp:
+                    model_latents = model_latents.to(dtype=amp_dtype)
+                with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp):
+                    x0 = self.transformer(model_latents, text, mask, num_loops=num_loops)
+                x0 = x0.float()
+                if self.do_classifier_free_guidance:
+                    x0_cond, x0_uncond = x0.chunk(2)
+                    x0 = x0_uncond + self.guidance_scale * (x0_cond - x0_uncond)
+
+                # sigma = 1 - t_flow. Passing -velocity makes `x + (sigma_next - sigma) * model_output`
+                # equal the training update `x + velocity * (t_next - t)`.
+                flow_time = (1.0 - t.to(device=latents.device, dtype=torch.float32) / num_train_timesteps)
+                velocity = (x0 - latents.float()) / (1.0 - flow_time).clamp_min(VELOCITY_DENOM_MIN)
+                latents = self.scheduler.step(-velocity, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {key: locals()[key] for key in callback_on_step_end_tensor_inputs}
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    prompt_attention_mask = callback_outputs.pop("prompt_attention_mask", prompt_attention_mask)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        images = self._images_from_latents(latents, output_type)
+
+        # Offload all models.
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

Looped-DiT-B-32/scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.39.0",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

Looped-DiT-B-32/text_encoder/README.md

@@ -0,0 +1,276 @@
+---
+language: 
+- en
+- fr
+- ro
+- de
+- multilingual
+
+widget:
+- text: "Translate to German:  My name is Arthur"
+  example_title: "Translation"
+- text: "Please answer to the following question. Who is going to be the next Ballon d'or?"
+  example_title: "Question Answering"
+- text: "Q: Can Geoffrey Hinton have a conversation with George Washington? Give the rationale before answering."
+  example_title: "Logical reasoning"
+- text: "Please answer the following question. What is the boiling point of Nitrogen?"
+  example_title: "Scientific knowledge"
+- text: "Answer the following yes/no question. Can you write a whole Haiku in a single tweet?"
+  example_title: "Yes/no question"
+- text: "Answer the following yes/no question by reasoning step-by-step. Can you write a whole Haiku in a single tweet?"
+  example_title: "Reasoning task"
+- text: "Q: ( False or not False or False ) is? A: Let's think step by step"
+  example_title: "Boolean Expressions"
+- text: "The square root of x is the cube root of y. What is y to the power of 2, if x = 4?"
+  example_title: "Math reasoning"
+- text: "Premise:  At my age you will probably have learnt one lesson. Hypothesis:  It's not certain how many lessons you'll learn by your thirties. Does the premise entail the hypothesis?"
+  example_title: "Premise and hypothesis"
+
+tags:
+- text2text-generation
+
+datasets:
+- svakulenk0/qrecc
+- taskmaster2
+- djaym7/wiki_dialog
+- deepmind/code_contests
+- lambada
+- gsm8k
+- aqua_rat
+- esnli
+- quasc
+- qed
+
+
+license: apache-2.0
+---
+
+# Model Card for FLAN-T5 large
+
+<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/flan2_architecture.jpg"
+alt="drawing" width="600"/>
+
+#  Table of Contents
+
+0. [TL;DR](#TL;DR)
+1. [Model Details](#model-details)
+2. [Usage](#usage)
+3. [Uses](#uses)
+4. [Bias, Risks, and Limitations](#bias-risks-and-limitations)
+5. [Training Details](#training-details)
+6. [Evaluation](#evaluation)
+7. [Environmental Impact](#environmental-impact)
+8. [Citation](#citation)
+9. [Model Card Authors](#model-card-authors)
+
+# TL;DR
+
+If you already know T5, FLAN-T5 is just better at everything. For the same number of parameters, these models have been fine-tuned on more than 1000 additional tasks covering also more languages. 
+As mentioned in the first few lines of the abstract : 
+>  Flan-PaLM 540B achieves state-of-the-art performance on several benchmarks, such as 75.2% on five-shot MMLU. We also publicly release Flan-T5 checkpoints,1 which achieve strong few-shot performance even compared to much larger models, such as PaLM 62B. Overall, instruction finetuning is a general method for improving the performance and usability of pretrained language models.
+
+**Disclaimer**: Content from **this** model card has been written by the Hugging Face team, and parts of it were copy pasted from the [T5 model card](https://huggingface.co/t5-large).
+
+# Model Details
+
+## Model Description
+
+
+- **Model type:** Language model
+- **Language(s) (NLP):** English, Spanish, Japanese, Persian, Hindi, French, Chinese, Bengali, Gujarati, German, Telugu, Italian, Arabic, Polish, Tamil, Marathi, Malayalam, Oriya, Panjabi, Portuguese, Urdu, Galician, Hebrew, Korean, Catalan, Thai, Dutch, Indonesian, Vietnamese, Bulgarian, Filipino, Central Khmer, Lao, Turkish, Russian, Croatian, Swedish, Yoruba, Kurdish, Burmese, Malay, Czech, Finnish, Somali, Tagalog, Swahili, Sinhala, Kannada, Zhuang, Igbo, Xhosa, Romanian, Haitian, Estonian, Slovak, Lithuanian, Greek, Nepali, Assamese, Norwegian
+- **License:** Apache 2.0
+- **Related Models:** [All FLAN-T5 Checkpoints](https://huggingface.co/models?search=flan-t5)
+- **Original Checkpoints:** [All Original FLAN-T5 Checkpoints](https://github.com/google-research/t5x/blob/main/docs/models.md#flan-t5-checkpoints)
+- **Resources for more information:**
+  - [Research paper](https://arxiv.org/pdf/2210.11416.pdf)
+  - [GitHub Repo](https://github.com/google-research/t5x)
+  - [Hugging Face FLAN-T5 Docs (Similar to T5) ](https://huggingface.co/docs/transformers/model_doc/t5)
+
+# Usage
+
+Find below some example scripts on how to use the model in `transformers`:
+
+## Using the Pytorch model
+
+### Running the model on a CPU
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+### Running the model on a GPU
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+# pip install accelerate
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto")
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+### Running the model on a GPU using different precisions
+
+#### FP16
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+# pip install accelerate
+import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto", torch_dtype=torch.float16)
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+#### INT8
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+# pip install bitsandbytes accelerate
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto", load_in_8bit=True)
+
+input_text = "translate English to German: How old are you?"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+# Uses
+
+## Direct Use and Downstream Use
+
+The authors write in [the original paper's model card](https://arxiv.org/pdf/2210.11416.pdf) that: 
+
+> The primary use is research on language models, including: research on zero-shot NLP tasks and in-context few-shot learning NLP tasks, such as reasoning, and question answering; advancing fairness and safety research, and understanding limitations of current large language models
+
+See the [research paper](https://arxiv.org/pdf/2210.11416.pdf) for further details.
+
+## Out-of-Scope Use
+
+More information needed.
+
+# Bias, Risks, and Limitations
+
+The information below in this section are copied from the model's [official model card](https://arxiv.org/pdf/2210.11416.pdf):
+
+> Language models, including Flan-T5, can potentially be used for language generation in a harmful way, according to Rae et al. (2021). Flan-T5 should not be used directly in any application, without a prior assessment of safety and fairness concerns specific to the application.
+
+## Ethical considerations and risks
+
+> Flan-T5 is fine-tuned on a large corpus of text data that was not filtered for explicit content or assessed for existing biases. As a result the model itself is potentially vulnerable to generating equivalently inappropriate content or replicating inherent biases in the underlying data.
+
+## Known Limitations
+
+> Flan-T5 has not been tested in real world applications.
+
+## Sensitive Use:
+
+> Flan-T5 should not be applied for any unacceptable use cases, e.g., generation of abusive speech.
+
+# Training Details
+
+## Training Data
+
+The model was trained on a mixture of tasks, that includes the tasks described in the table below (from the original paper, figure 2):
+
+![table.png](https://s3.amazonaws.com/moonup/production/uploads/1666363265279-62441d1d9fdefb55a0b7d12c.png)
+
+
+## Training Procedure
+
+According to the model card from the [original paper](https://arxiv.org/pdf/2210.11416.pdf):
+
+> These models are based on pretrained T5 (Raffel et al., 2020) and fine-tuned with instructions for better zero-shot and few-shot performance. There is one fine-tuned Flan model per T5 model size.
+
+The model has been trained on TPU v3 or TPU v4 pods, using [`t5x`](https://github.com/google-research/t5x) codebase together with [`jax`](https://github.com/google/jax).
+
+
+# Evaluation
+
+## Testing Data, Factors & Metrics
+
+The authors evaluated the model on various tasks covering several languages (1836 in total). See the table below for some quantitative evaluation:
+![image.png](https://s3.amazonaws.com/moonup/production/uploads/1668072995230-62441d1d9fdefb55a0b7d12c.png)
+For full details, please check the [research paper](https://arxiv.org/pdf/2210.11416.pdf).
+
+## Results 
+
+For full results for FLAN-T5-Large, see the [research paper](https://arxiv.org/pdf/2210.11416.pdf), Table 3.
+
+# Environmental Impact
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** Google Cloud TPU Pods - TPU v3 or TPU v4  | Number of chips ≥ 4.
+- **Hours used:** More information needed
+- **Cloud Provider:** GCP
+- **Compute Region:** More information needed
+- **Carbon Emitted:** More information needed
+
+# Citation
+
+**BibTeX:**
+
+```bibtex
+@misc{https://doi.org/10.48550/arxiv.2210.11416,
+  doi = {10.48550/ARXIV.2210.11416},
+  
+  url = {https://arxiv.org/abs/2210.11416},
+  
+  author = {Chung, Hyung Won and Hou, Le and Longpre, Shayne and Zoph, Barret and Tay, Yi and Fedus, William and Li, Eric and Wang, Xuezhi and Dehghani, Mostafa and Brahma, Siddhartha and Webson, Albert and Gu, Shixiang Shane and Dai, Zhuyun and Suzgun, Mirac and Chen, Xinyun and Chowdhery, Aakanksha and Narang, Sharan and Mishra, Gaurav and Yu, Adams and Zhao, Vincent and Huang, Yanping and Dai, Andrew and Yu, Hongkun and Petrov, Slav and Chi, Ed H. and Dean, Jeff and Devlin, Jacob and Roberts, Adam and Zhou, Denny and Le, Quoc V. and Wei, Jason},
+  
+  keywords = {Machine Learning (cs.LG), Computation and Language (cs.CL), FOS: Computer and information sciences, FOS: Computer and information sciences},
+  
+  title = {Scaling Instruction-Finetuned Language Models},
+  
+  publisher = {arXiv},
+  
+  year = {2022},
+  
+  copyright = {Creative Commons Attribution 4.0 International}
+}
+```

Looped-DiT-B-32/text_encoder/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 2816,
+  "d_kv": 64,
+  "d_model": 1024,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "n_positions": 512,
+  "num_decoder_layers": 24,
+  "num_heads": 16,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.23.1",
+  "use_cache": true,
+  "vocab_size": 32128
+}

Looped-DiT-B-32/text_encoder/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "decoder_start_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 0,
+  "transformers_version": "4.27.0.dev0"
+}

Looped-DiT-B-32/text_encoder/model.safetensors

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

Looped-DiT-B-32/text_encoder/special_tokens_map.json

@@ -0,0 +1,107 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-32/text_encoder/spiece.model

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

Looped-DiT-B-32/text_encoder/tokenizer_config.json

@@ -0,0 +1,113 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "model_max_length": 256,
+  "name_or_path": "google/t5-v1_1-large",
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "special_tokens_map_file": "/home/younes_huggingface_co/.cache/huggingface/hub/models--google--t5-v1_1-large/snapshots/314bc112b191ec17b625ba81438dc73d6c23659d/special_tokens_map.json",
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-32/tokenizer/special_tokens_map.json

@@ -0,0 +1,107 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-32/tokenizer/spiece.model

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

Looped-DiT-B-32/tokenizer/tokenizer_config.json

@@ -0,0 +1,113 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "model_max_length": 256,
+  "name_or_path": "google/t5-v1_1-large",
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "special_tokens_map_file": "/home/younes_huggingface_co/.cache/huggingface/hub/models--google--t5-v1_1-large/snapshots/314bc112b191ec17b625ba81438dc73d6c23659d/special_tokens_map.json",
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

Looped-DiT-B-32/transformer/config.json

@@ -0,0 +1,23 @@
+{
+  "_class_name": "LoopedDiTTransformer2DModel",
+  "_diffusers_version": "0.39.0",
+  "head_dim": 64,
+  "hidden_size": 768,
+  "image_size": 512,
+  "in_channels": 3,
+  "loop_split": [
+    6,
+    5,
+    6
+  ],
+  "mlp_ratio": 2.6667,
+  "num_heads": 12,
+  "num_loops": 4,
+  "patch_size": 32,
+  "pca_channels": 128,
+  "share_loop_weights": true,
+  "text_dim": 1024,
+  "text_preamble_depth": 2,
+  "use_attn_gate": false,
+  "use_xsa": true
+}

Looped-DiT-B-32/transformer/diffusion_pytorch_model.safetensors

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

Looped-DiT-B-32/transformer/transformer_looped_dit.py

@@ -0,0 +1,417 @@
+"""Looped-DiT denoiser as a diffusers model.
+
+The MiniT2I denoiser, a pixel-space variant of MMDiT (patchified image tokens
+and T5 text tokens with modality-specific weights and joint attention, no
+timestep conditioning), whose double-stream blocks are split into three stages:
+
+    pre-loop   A   blocks[:pre]              run once
+    looped     B   the next `core` blocks    run N times
+    post-loop  C   the last `post` blocks    run once, followed by the head
+
+    h_0 = A(x),   h_r = B(h_{r-1}),   x0_hat(r) = C(h_r),   r = 1..N
+
+With shared weights (the default) B is one set of blocks reused N times, so the
+loop adds depth but no weights. Any loop state h_r can be decoded through C:
+deep supervision trains those intermediate exits, and inference can run with a
+different loop depth. XSA and the attention gate act on the looped blocks only.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.models.modeling_utils import ModelMixin
+from torch import nn
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) * self.weight
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        hidden_dim = math.ceil(hidden_dim / 8) * 8
+        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)
+        for layer in (self.w1, self.w3, self.w2):
+            nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+# ---------------------------------------------------------------------------
+# Rotary position embeddings: 1D over text positions, 2D over the patch grid.
+# ---------------------------------------------------------------------------
+
+_ROPE_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+
+
+def _autocast_state() -> tuple:
+    try:
+        enabled = torch.is_autocast_enabled("cuda")
+        return enabled, torch.get_autocast_dtype("cuda") if enabled else None
+    except TypeError:  # torch < 2.4
+        enabled = torch.is_autocast_enabled()
+        return enabled, torch.get_autocast_gpu_dtype() if enabled else None
+
+
+def _rope_tables(grid: int | None, n: int, d: int, device, dtype, theta: float = 10000.0):
+    """cos/sin tables for 1D (grid=None) or 2D rotary embeddings.
+
+    Tables are built under whatever autocast state the caller runs in (under bf16
+    autocast the angle products are computed in bf16, which is what the models
+    were trained with), so that state is part of the cache key.
+    """
+    key = (grid, n, d, str(device), dtype, _autocast_state())
+    if key not in _ROPE_CACHE:
+        if grid is None:
+            inv = 1.0 / (theta ** (torch.arange(0, d, 2, device=device, dtype=torch.float32) / d))
+            pos = torch.arange(n, device=device, dtype=torch.float32)
+            angles = torch.einsum("n,f->nf", pos, inv)
+            angles = torch.cat([angles, angles], dim=-1)
+        else:
+            half = d // 2
+            inv = 1.0 / (theta ** (torch.arange(0, half, 2, device=device, dtype=torch.float32) / half))
+            freqs = torch.einsum("n,f->nf", torch.arange(grid, device=device, dtype=torch.float32), inv)
+            f_h, f_w = torch.broadcast_tensors(freqs[:, None, :], freqs[None, :, :])
+            angles = torch.cat([f_h, f_w], dim=-1)
+            angles = torch.cat([angles, angles], dim=-1).reshape(n, d)
+        _ROPE_CACHE[key] = (angles.cos()[None, None].to(dtype), angles.sin()[None, None].to(dtype))
+    return _ROPE_CACHE[key]
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x: torch.Tensor, grid: int | None = None) -> torch.Tensor:
+    """x: [batch, heads, tokens, head_dim]."""
+    cos, sin = _rope_tables(grid, x.shape[2], x.shape[3], x.device, x.dtype)
+    return x * cos + rotate_half(x) * sin
+
+
+def sincos_2d(dim: int, grid: int) -> torch.Tensor:
+    y, x = torch.meshgrid(torch.arange(grid), torch.arange(grid), indexing="ij")
+    omega = 1.0 / (10000 ** (torch.arange(dim // 4, dtype=torch.float32) / (dim // 4)))
+    out_y = torch.einsum("n,d->nd", y.flatten().float(), omega)
+    out_x = torch.einsum("n,d->nd", x.flatten().float(), omega)
+    return torch.cat([out_x.sin(), out_x.cos(), out_y.sin(), out_y.cos()], dim=1)
+
+
+def exclusive_self_attention(out: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """XSA (Zhai, 2026: https://arxiv.org/abs/2603.09078):
+    remove from each token's attention output the component along that token's
+    own value vector, so attention only writes content from other tokens.
+    `out` and `v` are token-aligned, heads first."""
+    v_hat = F.normalize(v.float(), dim=-1)
+    out_f = out.float()
+    return (out_f - (out_f * v_hat).sum(dim=-1, keepdim=True) * v_hat).to(out.dtype)
+
+
+# ---------------------------------------------------------------------------
+# Blocks
+# ---------------------------------------------------------------------------
+
+
+class PatchEmbed(nn.Module):
+    """Two-stage patch embedding: a low-rank patch projection, then a 1x1 conv."""
+
+    def __init__(self, patch_size: int, in_channels: int, hidden_size: int, bottleneck: int):
+        super().__init__()
+        self.proj1 = nn.Conv2d(in_channels, bottleneck, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(bottleneck, hidden_size, kernel_size=1, bias=True)
+        nn.init.xavier_uniform_(self.proj1.weight)
+        nn.init.xavier_uniform_(self.proj2.weight)
+        nn.init.zeros_(self.proj2.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj2(self.proj1(x)).flatten(2).transpose(1, 2)
+
+
+class TextBlock(nn.Module):
+    """Text-only transformer block that refines the T5 tokens before the joint blocks."""
+
+    def __init__(self, hidden_size: int, num_heads: int, head_dim: int, mlp_ratio: float):
+        super().__init__()
+        self.num_heads, self.head_dim = num_heads, head_dim
+        self.norm1 = RMSNorm(hidden_size)
+        self.norm2 = RMSNorm(hidden_size)
+        self.qkv = nn.Linear(hidden_size, num_heads * head_dim * 3)
+        self.proj = nn.Linear(num_heads * head_dim, hidden_size)
+        self.mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+    def forward(self, txt: torch.Tensor) -> torch.Tensor:
+        b, n, _ = txt.shape
+        q, k, v = self.qkv(self.norm1(txt)).view(b, n, 3, self.num_heads, self.head_dim).unbind(2)
+        q, k, v = (z.transpose(1, 2) for z in (self.q_norm(q), self.k_norm(k), v))
+        out = F.scaled_dot_product_attention(apply_rope(q), apply_rope(k), v, scale=self.head_dim**-0.5)
+        txt = txt + self.proj(out.transpose(1, 2).reshape(b, n, -1))
+        return txt + self.mlp(self.norm2(txt))
+
+
+class DoubleStreamBlock(nn.Module):
+    """MMDiT block: separate image/text weights, one joint attention over both.
+
+    `use_xsa` / `use_attn_gate` turn on self-modulating attention (set only for the
+    looped blocks). The attention gate (Qiu et al., 2026:
+    https://arxiv.org/abs/2505.06708) is head-wise:
+    y_i <- y_i * sigmoid(W_g u_i + b_g), with u_i the block's normed input.
+    
+    `update_text=False` skips the text-stream update, for the last block, whose
+    text output is never read.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        head_dim: int,
+        mlp_ratio: float,
+        grid: int,
+        use_xsa: bool = False,
+        use_attn_gate: bool = False,
+        update_text: bool = True,
+    ):
+        super().__init__()
+        self.num_heads, self.head_dim, self.grid = num_heads, head_dim, grid
+        self.use_xsa, self.use_attn_gate, self.update_text = use_xsa, use_attn_gate, update_text
+        inner = num_heads * head_dim
+        self.img_norm1 = RMSNorm(hidden_size)
+        self.img_norm2 = RMSNorm(hidden_size)
+        self.txt_norm1 = RMSNorm(hidden_size)
+        self.txt_norm2 = RMSNorm(hidden_size)
+        self.img_qkv = nn.Linear(hidden_size, inner * 3)
+        self.txt_qkv = nn.Linear(hidden_size, inner * 3)
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+        self.img_proj = nn.Linear(inner, hidden_size)
+        self.txt_proj = nn.Linear(inner, hidden_size)
+        if use_attn_gate:
+            # Zero bias: the gates start half open on average.
+            self.img_gate = nn.Linear(hidden_size, num_heads)
+            self.txt_gate = nn.Linear(hidden_size, num_heads)
+            nn.init.zeros_(self.img_gate.bias)
+            nn.init.zeros_(self.txt_gate.bias)
+        self.img_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        self.txt_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+
+    def forward(self, img: torch.Tensor, txt: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        b, li, _ = img.shape
+        lt = txt.shape[1]
+        img_n, txt_n = self.img_norm1(img), self.txt_norm1(txt)
+        qi, ki, vi = self.img_qkv(img_n).view(b, li, 3, self.num_heads, self.head_dim).unbind(2)
+        qt, kt, vt = self.txt_qkv(txt_n).view(b, lt, 3, self.num_heads, self.head_dim).unbind(2)
+        # Joint sequence [text; image], heads first.
+        q = torch.cat([qt, qi], dim=1).transpose(1, 2)
+        k = torch.cat([kt, ki], dim=1).transpose(1, 2)
+        v = torch.cat([vt, vi], dim=1).transpose(1, 2)
+        q = torch.cat([apply_rope(self.q_norm(q[:, :, :lt])), apply_rope(self.q_norm(q[:, :, lt:]), self.grid)], dim=2)
+        k = torch.cat([apply_rope(self.k_norm(k[:, :, :lt])), apply_rope(self.k_norm(k[:, :, lt:]), self.grid)], dim=2)
+        out = F.scaled_dot_product_attention(q, k, v, scale=self.head_dim**-0.5)
+        if self.use_xsa:
+            out = exclusive_self_attention(out, v)
+        out = out.transpose(1, 2)  # [b, tokens, heads, head_dim]
+        out_t, out_i = out[:, :lt], out[:, lt:]
+        if self.use_attn_gate:
+            out_i = out_i * torch.sigmoid(self.img_gate(img_n)).unsqueeze(-1)
+        img = img + self.img_proj(out_i.reshape(b, li, -1))
+        img = img + self.img_mlp(self.img_norm2(img))
+        if self.update_text:
+            if self.use_attn_gate:
+                out_t = out_t * torch.sigmoid(self.txt_gate(txt_n)).unsqueeze(-1)
+            txt = txt + self.txt_proj(out_t.reshape(b, lt, -1))
+            txt = txt + self.txt_mlp(self.txt_norm2(txt))
+        return img, txt
+
+
+# ---------------------------------------------------------------------------
+# Looped MMDiT
+# ---------------------------------------------------------------------------
+
+
+class LoopedDiTTransformer2DModel(ModelMixin, ConfigMixin):
+    """Predicts the clean image x0 from a noisy image and T5 text embeddings.
+
+    This is a diffusers [`ModelMixin`]: `save_pretrained` / `from_pretrained` round-trip the
+    architecture in `config.json`. `loop_split` is `(pre, core, post)` and `num_loops` is the
+    trained loop depth (`1` is the MiniT2I model without looping). With `share_loop_weights=False`
+    every pass gets its own copy of the core blocks: the compute-matched "deeper" baseline with
+    the same exits.
+
+    The forward pass does not take a timestep. Flow-matching time is applied by the pipeline when
+    it converts the x0 prediction into a velocity.
+    """
+
+    def __init__(
+        self,
+        image_size: int = 512,
+        patch_size: int = 32,
+        in_channels: int = 3,
+        hidden_size: int = 768,
+        num_heads: int = 12,
+        head_dim: int = 64,
+        mlp_ratio: float = 2.6667,
+        pca_channels: int = 128,
+        text_dim: int = 1024,
+        text_preamble_depth: int = 2,
+        loop_split: tuple[int, int, int] | list[int] = (6, 5, 6),
+        num_loops: int = 4,
+        share_loop_weights: bool = True,
+        use_xsa: bool = False,
+        use_attn_gate: bool = False,
+    ):
+        super().__init__()
+        loop_split = [int(n) for n in loop_split]
+        num_loops = int(num_loops)
+        share_loop_weights = bool(share_loop_weights)
+        use_xsa = bool(use_xsa)
+        use_attn_gate = bool(use_attn_gate)
+        if len(loop_split) != 3 or min(loop_split) < 1:
+            raise ValueError(f"loop_split must be three positive block counts (pre, core, post), got {loop_split}")
+        if num_loops < 1:
+            raise ValueError(f"num_loops must be >= 1, got {num_loops}")
+        # Lists (not tuples) so config.json stays valid JSON.
+        self.register_to_config(
+            image_size=image_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            mlp_ratio=mlp_ratio,
+            pca_channels=pca_channels,
+            text_dim=text_dim,
+            text_preamble_depth=text_preamble_depth,
+            loop_split=loop_split,
+            num_loops=num_loops,
+            share_loop_weights=share_loop_weights,
+            use_xsa=use_xsa,
+            use_attn_gate=use_attn_gate,
+        )
+        pre, core, post = loop_split
+        self.patch_size, self.in_channels = patch_size, in_channels
+        self.grid = image_size // patch_size
+        self.pre, self.core, self.post = pre, core, post
+        self.num_loops = num_loops
+        self.share_loop_weights = share_loop_weights
+
+        self.img_embed = PatchEmbed(patch_size, in_channels, hidden_size, pca_channels)
+        self.txt_embed = nn.Linear(text_dim, hidden_size, bias=False)
+        # Replaces the T5 embedding at padded prompt positions (and everywhere
+        # for the unconditional branch of classifier-free guidance).
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, text_dim))
+        nn.init.normal_(self.mask_token, std=0.02)
+        # MiniT2I's timestep and pooled-text embedders. The model has no timestep
+        # conditioning and never uses them; they are kept (frozen, see below) so that
+        # the model and its checkpoints match MiniT2I and the paper.
+        self.t_embed = nn.ModuleDict(
+            {"mlp": nn.Sequential(nn.Linear(256, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size))}
+        )
+        for layer in (self.t_embed.mlp[0], self.t_embed.mlp[2]):
+            nn.init.normal_(layer.weight, std=0.02)
+            nn.init.zeros_(layer.bias)
+        self.pooled_embed = nn.Linear(text_dim, hidden_size, bias=False)
+        self.register_buffer("pos_embed", sincos_2d(hidden_size, self.grid)[None], persistent=False)
+        self.txt_blocks = nn.ModuleList(
+            TextBlock(hidden_size, num_heads, head_dim, mlp_ratio) for _ in range(text_preamble_depth)
+        )
+        looped = core if self.share_loop_weights else core * self.num_loops
+        depth = pre + looped + post
+        self.blocks = nn.ModuleList(
+            DoubleStreamBlock(
+                hidden_size,
+                num_heads,
+                head_dim,
+                mlp_ratio,
+                self.grid,
+                use_xsa=use_xsa and pre <= i < pre + looped,
+                use_attn_gate=use_attn_gate and pre <= i < pre + looped,
+                update_text=i < depth - 1,
+            )
+            for i in range(depth)
+        )
+        self.final_norm = RMSNorm(hidden_size)
+        self.final = nn.Linear(hidden_size, patch_size * patch_size * in_channels)
+        nn.init.zeros_(self.final.weight)
+        nn.init.zeros_(self.final.bias)
+        # Frozen because they never receive a gradient: the unused embedders and the
+        # text-stream update of the last block (whose text output is never read).
+        last = self.blocks[-1]
+        for module in (self.t_embed, self.pooled_embed, last.txt_norm2, last.txt_proj, last.txt_mlp):
+            module.requires_grad_(False)
+
+    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
+        b, n, _ = x.shape
+        p, c, g = self.patch_size, self.in_channels, int(n**0.5)
+        x = x.view(b, g, g, p, p, c).permute(0, 5, 1, 3, 2, 4).contiguous()
+        return x.view(b, c, g * p, g * p)
+
+    def loop_blocks(self, r: int) -> nn.ModuleList:
+        """Core blocks run on loop pass r (1-based)."""
+        start = self.pre if self.share_loop_weights else self.pre + (r - 1) * self.core
+        return self.blocks[start : start + self.core]
+
+    def decode(self, img: torch.Tensor, txt: torch.Tensor) -> torch.Tensor:
+        """Post-loop blocks and output head: a loop state -> x0 prediction."""
+        for block in self.blocks[len(self.blocks) - self.post :]:
+            img, txt = block(img, txt)
+        return self.unpatchify(self.final(self.final_norm(img))).float()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text: torch.Tensor,
+        text_mask: torch.Tensor,
+        num_loops: int | None = None,
+        exit_loops: tuple[int, ...] = (),
+    ) -> torch.Tensor | tuple[torch.Tensor, dict[int, torch.Tensor]]:
+        """x: noisy images [B, C, H, W]; text: T5 states [B, L, text_dim];
+        text_mask: [B, L], 1 for prompt tokens (all 0 = unconditional).
+
+        num_loops overrides the loop depth at inference. exit_loops lists
+        intermediate depths r < num_loops to decode as well; the call then
+        returns (final prediction, {r: prediction after r loops}).
+        """
+        n = self.num_loops if num_loops is None else int(num_loops)
+        if n < 1 or (not self.share_loop_weights and n > self.num_loops):
+            raise ValueError(f"num_loops={n} is not available for this model (trained with {self.num_loops})")
+        exits = sorted({int(r) for r in exit_loops})
+        if any(not 1 <= r < n for r in exits):
+            raise ValueError(f"exit_loops must lie in [1, {n}), got {exits}")
+
+        text = torch.where(text_mask.to(torch.bool)[:, :, None], text, self.mask_token.to(text.dtype))
+        img = self.img_embed(x) + self.pos_embed.to(device=x.device, dtype=x.dtype)
+        txt = self.txt_embed(text)
+        for block in self.txt_blocks:
+            txt = block(txt)
+        for block in self.blocks[: self.pre]:
+            img, txt = block(img, txt)
+        states = {}
+        for r in range(1, n + 1):
+            for block in self.loop_blocks(r):
+                img, txt = block(img, txt)
+            if r in exits:
+                states[r] = (img, txt)
+        out = self.decode(img, txt)
+        if not exits:
+            return out
+        return out, {r: self.decode(*states[r]) for r in exits}
+
+
+# Training imports this name. The diffusers class name is the one stored in checkpoints.
+LoopedMMDiT = LoopedDiTTransformer2DModel

README.md

@@ -0,0 +1,151 @@
+---
+license: mit
+library_name: diffusers
+pipeline_tag: text-to-image
+tags:
+  - diffusers
+  - looped-dit
+  - image-generation
+  - text-to-image
+  - flow-matching
+  - pixel-space
+inference: true
+widget:
+  - text: a red cube on top of a blue sphere
+    output:
+      url: Looped-DiT-B-16/demo.png
+language:
+  - en
+---
+
+# BiliSakura/Looped-DiT-diffusers
+
+Self-contained Looped-DiT text-to-image checkpoints for Hugging Face diffusers. Each variant folder ships its own pipeline code, component modules, bundled FLAN-T5-Large text encoder, and transformer weights.
+
+## Available checkpoints
+
+| Subfolder | Model | Params (denoiser + text encoder) | Patch | Loop depth | CFG |
+| --- | --- | --- | ---: | ---: | ---: |
+| [`Looped-DiT-B-32/`](Looped-DiT-B-32/) | Looped-DiT B/32 | 260M + 341M | 32 | 4 | 6.0 |
+| [`Looped-DiT-B-16/`](Looped-DiT-B-16/) | Looped-DiT B/16 | 258M + 341M | 16 | 4 | 6.0 |
+
+Benchmark scores (100 Euler steps, CFG 6.0, loop depth 4):
+
+| Model | GenEval | DPG-Bench | PRISM | CoReBench | SpatialGenEval | TIIF-Short | Avg |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| B/32 (290k) | 85.1 | 85.3 | 54.4 | 44.5 | 52.3 | 76.1 | 66.3 |
+| B/16 (580k) | 87.4 | 87.0 | 67.0 | 53.5 | 54.6 | 79.7 | 71.5 |
+
+## Repo layout
+
+```text
+BiliSakura/Looped-DiT-diffusers/
+├── README.md
+├── .gitattributes
+├── Looped-DiT-B-32/
+│   ├── pipeline.py
+│   ├── model_index.json
+│   ├── demo.png
+│   ├── scheduler/
+│   ├── text_encoder/
+│   ├── tokenizer/
+│   └── transformer/
+└── Looped-DiT-B-16/
+    └── ...
+```
+
+Each variant is self-contained: load with `custom_pipeline` pointing at that folder’s `pipeline.py` and `trust_remote_code=True`. Looped-DiT denoises directly in RGB pixel space (no VAE).
+
+## Demo
+
+![Looped-DiT-B-16 demo](Looped-DiT-B-16/demo.png)
+
+Prompt: *"a red cube on top of a blue sphere."* — Looped-DiT B/16 at 512×512, 100 steps, `guidance_scale=6.0`, `num_loops=4`, `torch_dtype=bfloat16`, seed 42.
+
+![Looped-DiT-B-32 demo](Looped-DiT-B-32/demo.png)
+
+Same prompt and settings with Looped-DiT B/32.
+
+## Load from Hugging Face
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/Looped-DiT-diffusers",
+    subfolder="Looped-DiT-B-16",
+    custom_pipeline="pipeline.py",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+generator = torch.Generator(device="cuda").manual_seed(42)
+image = pipe(
+    "a red cube on top of a blue sphere",
+    num_inference_steps=100,
+    guidance_scale=6.0,
+    num_loops=4,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+For B/32, set `subfolder="Looped-DiT-B-32"`.
+
+## Load from a local clone
+
+```python
+from pathlib import Path
+import torch
+from diffusers import DiffusionPipeline
+
+model_dir = Path("./Looped-DiT-B-16").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,
+).to("cuda")
+
+generator = torch.Generator(device="cuda").manual_seed(42)
+image = pipe(
+    "a red cube on top of a blue sphere",
+    num_inference_steps=100,
+    guidance_scale=6.0,
+    num_loops=4,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+Use `./Looped-DiT-B-32` instead of `./Looped-DiT-B-16` for the B/32 checkpoint.
+
+## Recommended inference settings
+
+| Variant | Resolution | Steps | CFG scale | `num_loops` | `torch_dtype` |
+| --- | --- | ---: | ---: | ---: | --- |
+| `Looped-DiT-B-32` | 512×512 | 100 | 6.0 | 4 (default) | `bfloat16` (full pipeline) |
+| `Looped-DiT-B-16` | 512×512 | 100 | 6.0 | 4 (default) | `bfloat16` (full pipeline) |
+
+Other loop depths work at inference when loop weights are shared (the default for released models).
+
+## Interface notes
+
+- Text conditioning uses bundled `google/flan-t5-large` (`T5EncoderModel` + `T5Tokenizer`) in **bfloat16**, the same dtype as the denoiser. Prompt length is the tokenizer `model_max_length` (256).
+- `torch_dtype=torch.bfloat16` on `from_pretrained` sets both. Do not cast `pipe.text_encoder` back to float32.
+- Set `custom_pipeline` to the variant’s `pipeline.py` (Hub: `"pipeline.py"` with `subfolder`; local: absolute path).
+- Scheduler is `FlowMatchEulerDiscreteScheduler` with 1000 training timesteps and `shift=1.0`.
+- `guidance_scale > 1.0` enables classifier-free guidance with an empty-string null prompt.
+- Output resolution is fixed at 512×512.
+
+## Links
+
+- Upstream B/32 weights: [sensenova/Looped-DiT-B32](https://huggingface.co/sensenova/Looped-DiT-B32)
+- Upstream B/16 weights: [sensenova/Looped-DiT-B16](https://huggingface.co/sensenova/Looped-DiT-B16)
+- Backbone: [MiniT2I](https://github.com/PeppaKing8/minit2i-jax) · [BiliSakura/MiniT2I-diffusers](https://huggingface.co/BiliSakura/MiniT2I-diffusers)
+
+## License
+
+MIT (same as upstream Looped-DiT and MiniT2I).