Add Boomer FLA fine-tuned checkpoint (step 055000, ema weights)

.gitattributes

@@ -0,0 +1 @@
+transformer/diffusion_pytorch_model.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,199 @@
+---
+license: other
+tags:
+  - text-to-image
+  - diffusion
+  - linear-attention
+  - pytorch
+  - safetensors
+language:
+  - en
+pipeline_tag: text-to-image
+---
+
+# Boomer FLA
+
+Boomer FLA is a 657M parameter text-to-image diffusion model that generates **1024×1024px** images from text prompts.
+
+Instead of standard quadratic self-attention, it uses **GatedDeltaNet** — a bidirectional Flash Linear Attention mixer — as the backbone of its transformer blocks. This keeps memory flat regardless of sequence length. Every 6th block adds a full SDPA layer for global spatial coherence.
+
+Text conditioning uses **Gemma 4 2B** (1536-dim embeddings, up to 300 tokens). Decoding uses the **DC-AE f32c32** VAE with 32× spatial compression, producing 32×32 latents from 1024px images. Sampling uses **STORK-2**, a high-order Runge–Kutta flow matching solver that converges in 32 steps.
+
+Fine-tuned from a JourneyDB-pretrained base on 600k high-resolution images at 1024px.
+
+---
+
+## Sample outputs
+
+![Model grid — portraits and landscapes](boomer_model_grid.png)
+
+*Portraits (top) and landscapes (bottom) generated at 1024×1024px, 32 STORK-2 steps, CFG 4.5.*
+
+---
+
+## Architecture
+
+| Property | Value |
+|---|---|
+| Parameters | 657M |
+| Backbone | Bidirectional GatedDeltaNet (Flash Linear Attention) |
+| Depth | 24 layers |
+| Hidden dim | 896 |
+| Heads | 14 |
+| Image attention | Every 6th layer (full SDPA + 2D RoPE) |
+| Patch size | 1 — one token per latent pixel (256 tokens at 512px, 1024 tokens at 1024px) |
+| Text encoder | Gemma 4 2B (`google/gemma-4-E2B-it`) |
+| VAE | DC-AE f32c32 (`mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers`) |
+| Sampler | STORK-2, 32 steps |
+| Dtype | bfloat16 |
+
+---
+
+## Training details
+
+| Setting | Value |
+|---|---|
+| Pre-train dataset | JourneyDB (~3.8M images, 512px, patch size 1) |
+| Fine-tune dataset | FineT2I (~600k images, 1024px, patch size 1) |
+| Fine-tune steps | 55,000 |
+| Batch size | 24 |
+| Learning rate | 1e-4, linear warmup (1000 steps) → cosine decay |
+| Flow shift | 1.5 |
+| Timestep sampler | Plateau logit-normal (μ=0, σ=1) |
+| Min-SNR γ | 5.0 |
+| CFG dropout | 0.1 |
+| EMA decay | 0.999 |
+| Gradient clip | 0.3 |
+| Optimizer | Fused AdamW |
+| Hardware | NVIDIA A100 (Google Colab) |
+| Precision | bfloat16 |
+
+---
+
+## VRAM and RAM requirements
+
+Measured at 1024×1024px, bfloat16, STORK-2, CFG batch=2.
+
+| Component | VRAM |
+|---|---|
+| DiT weights (EMA, bf16) | 1.25 GB |
+| Gemma 4 2B text encoder | 8.62 GB |
+| Denoising peak (CFG on) | 1.36 GB |
+| VAE decode peak | 3.51 GB |
+
+| Mode | Peak VRAM | Minimum GPU |
+|---|---|---|
+| **Condition-cache** — pre-encoded embeddings, no text encoder in VRAM | **4.76 GB** | RTX 3060 8GB, T4 |
+| **Fresh-prompt** — text encoder + DiT + VAE together | **13.38 GB** | RTX 3090, A100 |
+
+**System RAM**: loading the text encoder (Gemma 4 2B) requires ~9 GB of system RAM even when using GPU. For condition-cache mode, encoding can be done on CPU with ~9 GB RAM — the generation step then needs only 5 GB VRAM.
+
+---
+
+## Usage
+
+### Install
+
+```bash
+pip install torch diffusers transformers accelerate safetensors
+pip install git+https://github.com/fla-org/flash-linear-attention.git
+# STORK is bundled with the model — no separate install needed
+```
+
+### Generate
+
+```python
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "akrao9/Boomer-T2I",
+    trust_remote_code=True,
+    torch_dtype="auto",
+).to("cuda")
+
+image = pipe("a photorealistic portrait of a woman with dark hair")[0]
+image.save("output.png")
+```
+
+### Parameters
+
+```python
+image = pipe(
+    "a rocky coastline at sunset with crashing waves",
+    steps=32,        # denoising steps — 32 is recommended with STORK-2
+    cfg_scale=4.5,   # classifier-free guidance scale (4.0–5.0)
+    cfg_rescale=0.5, # reduces over-saturation at high CFG
+    seed=42,
+)[0]
+```
+
+### Low VRAM — condition cache mode
+
+Encode prompts once on any machine (including CPU), save the embedding, then generate with only the 1.25 GB DiT loaded. Peak VRAM drops from 13.38 GB → 4.76 GB.
+
+```python
+# Step 1 — encode on any machine (even CPU with 9GB RAM)
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+TE_REPO      = "google/gemma-4-E2B-it"
+tokenizer    = AutoProcessor.from_pretrained(TE_REPO)
+text_encoder = AutoModelForCausalLM.from_pretrained(
+    TE_REPO, torch_dtype=torch.bfloat16
+).get_decoder()
+
+tokens = tokenizer(
+    "a mountain lake surrounded by alpine peaks",
+    max_length=300, padding="max_length",
+    truncation=True, return_tensors="pt",
+)
+with torch.inference_mode():
+    hidden = text_encoder(
+        tokens["input_ids"], attention_mask=tokens["attention_mask"]
+    )[0]
+    idx  = [0] + list(range(-299, 0))
+    emb  = hidden[:, idx]
+    mask = tokens["attention_mask"][:, idx]
+
+torch.save({"emb": emb.cpu(), "mask": mask.cpu()}, "condition.pt")
+```
+
+```python
+# Step 2 — generate on low-VRAM GPU (no text encoder needed in VRAM)
+from diffusers import DiffusionPipeline
+import torch
+
+pipe = DiffusionPipeline.from_pretrained(
+    "akrao9/Boomer-T2I",
+    trust_remote_code=True,
+    torch_dtype="auto",
+).to("cuda")
+
+saved = torch.load("condition.pt")
+image = pipe(
+    prompt="",
+    _preencoded_emb=saved["emb"].cuda(),
+    _preencoded_mask=saved["mask"].cuda(),
+)[0]
+image.save("output.png")
+```
+
+---
+
+## Limitations
+
+- **Strong at** — photorealistic human portraits, dramatic landscapes, architectural scenes
+- **Weak at** — animals, text rendering, small detailed objects (limited training data coverage)
+- Landscapes have a painterly/HDR bias inherited from heavily post-processed stock images in the training set
+- Not safety filtered — outputs may reflect biases in the training data
+- Maximum tested resolution: **1024×1024px**
+
+---
+
+## License
+
+The Boomer FLA model weights are released for research and personal use. Commercial use is not permitted without explicit permission.
+
+Upstream component licenses:
+- DC-AE VAE: [mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers](https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers)
+- Gemma 4 text encoder: [Gemma Terms of Use](https://ai.google.dev/gemma/terms)

STORKScheduler.py

@@ -0,0 +1,1641 @@
+# 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 math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+import numpy as np
+import torch
+from scipy.io import loadmat
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
+from diffusers.utils import BaseOutput, is_scipy_available, logging
+from pathlib import Path
+
+
+
+@dataclass
+class STORKSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+current_file = Path(__file__)
+CONSTANTSFOLDER = f"{current_file.parent}/STORK_constants"
+
+
+
+
+
+class STORKScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `STORKScheduler` uses modified stabilized Runge-Kutta method for the backward ODE in the diffusion or flow matching models.
+    This include the original STORK method and the modified STORK++ methods.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+        use_dynamic_shifting (`bool`, defaults to False):
+            Whether to apply timestep shifting on-the-fly based on the image resolution.
+        base_shift (`float`, defaults to 0.5):
+            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
+            with desired output.
+        max_shift (`float`, defaults to 1.15):
+            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
+            more exaggerated or stylized.
+        base_image_seq_len (`int`, defaults to 256):
+            The base image sequence length.
+        max_image_seq_len (`int`, defaults to 4096):
+            The maximum image sequence length.
+        invert_sigmas (`bool`, defaults to False):
+            Whether to invert the sigmas.
+        shift_terminal (`float`, defaults to None):
+            The end value of the shifted timestep schedule.
+        use_karras_sigmas (`bool`, defaults to False):
+            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
+        use_exponential_sigmas (`bool`, defaults to False):
+            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
+        use_beta_sigmas (`bool`, defaults to False):
+            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
+        solver_order (`int`, defaults to 2):
+            The STORK order which can be `2` or `4`. It is recommended to use `solver_order=2` uniformly.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process) or `flow_prediction`.
+        time_shift_type (`str`, defaults to "exponential"):
+            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
+        derivative_order (`int`, defaults to 1):
+            The order of the Taylor expansion derivative to use for the sub-step velocity approximation. Only supports 1, 2 or 3.
+        s (`int`, defaults to 50):
+            The number of sub-steps to use in the STORK.
+        precision (`str`, defaults to "float32"):
+            The precision to use for the scheduler; supports "float32", "bfloat16", or "float16".
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        use_dynamic_shifting: bool = False,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        stopping_eps: float = 1e-2,
+        solver_order: int = 4,
+        prediction_type: str = "epsilon",
+        time_shift_type: str = "exponential",
+        derivative_order: int = 1,
+        s: int = 50,
+        base_shift: Optional[float] = 0.5,
+        max_shift: Optional[float] = 1.15,
+        base_image_seq_len: Optional[int] = 256,
+        max_image_seq_len: Optional[int] = 4096,
+        invert_sigmas: bool = False,
+        shift_terminal: Optional[float] = None,
+        use_karras_sigmas: Optional[bool] = False,
+        use_exponential_sigmas: Optional[bool] = False,
+        use_beta_sigmas: Optional[bool] = False,
+    ):
+        
+        super().__init__()
+        # if prediction_type == "flow_prediction" and sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
+        #     raise ValueError(
+        #         "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
+        #     )
+        if time_shift_type not in {"exponential", "linear"}:
+            raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.")
+ 
+        # We manually enforce precision to float32 for numerical issues.Add commentMore actions
+        self.np_dtype = np.float32
+        self.dtype = torch.float32
+
+
+        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=self.np_dtype)[::-1].copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=self.dtype)
+        sigmas = timesteps / num_train_timesteps
+
+
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = None    #sigmas * num_train_timesteps
+        self._step_index = None
+        self._begin_index = None
+        self._shift = shift
+        self.sigmas = sigmas #.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+        # Store the predictions for the velocity/noise for higher order derivative approximations
+        self.velocity_predictions = []
+        self.noise_predictions = []
+        self.s = s
+        self.derivative_order = derivative_order
+
+        self.solver_order = solver_order
+        self.prediction_type = prediction_type
+
+
+        # Set the betas for noise-based models
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+        else:
+            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+        
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+        
+        # Noise-based models epsilon to avoid numerical issues
+        self.stopping_eps = stopping_eps
+
+
+
+
+    def set_timesteps(
+        self,
+        num_inference_steps: Optional[int] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[float] = None,
+        timesteps: Optional[List[float]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`, *optional*):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            sigmas (`List[float]`, *optional*):
+                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
+                automatically.
+            mu (`float`, *optional*):
+                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
+                shifting.
+            timesteps (`List[float]`, *optional*):
+                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
+                automatically.
+        """
+        
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`")
+
+        if sigmas is not None and timesteps is not None:
+            if len(sigmas) != len(timesteps):
+                raise ValueError("`sigmas` and `timesteps` should have the same length")
+
+        if num_inference_steps is not None:
+            if (sigmas is not None and len(sigmas) != num_inference_steps) or (
+                timesteps is not None and len(timesteps) != num_inference_steps
+            ):
+                raise ValueError(
+                    "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided"
+                )
+        else:
+            num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps)
+
+        self.num_inference_steps = num_inference_steps
+
+        if self.prediction_type == "epsilon":
+            self.set_timesteps_noise(num_inference_steps, device)
+        elif self.prediction_type == "flow_prediction":
+            self.set_timesteps_flow_matching(num_inference_steps, device, sigmas, mu, timesteps)
+        else:
+            raise ValueError(f"Prediction type {self.prediction_type} is not yet supported")
+        
+        # Reset the step index and begin index
+        self._step_index = None
+        self._begin_index = None
+
+        
+
+    def set_timesteps_noise(self,
+        num_inference_steps: Optional[int] = None,
+        device: Union[str, torch.device] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference), for noise-based models.
+
+        Args:
+            num_inference_steps (`int`, *optional*):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+        seq = np.linspace(0, 1, self.num_inference_steps+1)
+        seq[0] = self.stopping_eps
+        seq = seq[:-1]
+        seq = seq[::-1]
+
+
+        # The following lines are for the uniform timestepping case
+        self.dt = seq[0] - seq[1]
+        seq = seq * self.config.num_train_timesteps
+        seq[-1] = self.stopping_eps * self.config.num_train_timesteps
+        self._timesteps = seq
+        self.timesteps = torch.from_numpy(seq.copy()).to(device)
+
+
+        self._step_index = None
+        self._begin_index = None
+
+        self.noise_predictions = []
+
+
+    def set_timesteps_flow_matching(self,
+        num_inference_steps: Optional[int] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[float] = None,
+        timesteps: Optional[List[float]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the flow matching based models (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`, *optional*):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            sigmas (`List[float]`, *optional*):
+                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
+                automatically.
+            mu (`float`, *optional*):
+                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
+                shifting.
+            timesteps (`List[float]`, *optional*):
+                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
+                automatically.
+        """
+        self.num_inference_steps = num_inference_steps
+
+        # 1. Prepare default sigmas
+        is_timesteps_provided = timesteps is not None
+
+        if is_timesteps_provided:
+            timesteps = np.array(timesteps).astype(np.float32)
+
+        if sigmas is None:
+            if timesteps is None:
+                timesteps = np.linspace(
+                    self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+                )
+            sigmas = timesteps / self.config.num_train_timesteps
+        else:
+            sigmas = np.array(sigmas).astype(np.float32)
+            num_inference_steps = len(sigmas)
+
+        # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
+        #    "exponential" or "linear" type is applied
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
+
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
+        if self.config.shift_terminal:
+            sigmas = self.stretch_shift_to_terminal(sigmas)
+
+        # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
+        if self.config.use_karras_sigmas:
+            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_exponential_sigmas:
+            sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_beta_sigmas:
+            sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+        # 5. Convert sigmas and timesteps to tensors and move to specified device
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+        if not is_timesteps_provided:
+            timesteps = sigmas * self.config.num_train_timesteps
+        else:
+            timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)
+
+        # 6. Append the terminal sigma value.
+        #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
+        #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
+        if self.config.invert_sigmas:
+            sigmas = 1.0 - sigmas
+            timesteps = sigmas * self.config.num_train_timesteps
+            sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
+        else:
+            sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+        self.timesteps = timesteps
+        self.sigmas = sigmas
+
+
+        # Create the dt list
+        self.dt_list = self.sigmas[:-1] - self.sigmas[1:]
+        self.dt_list = self.dt_list.reshape(-1)
+
+        self.dt_list = self.dt_list.tolist()
+        self.dt_list = torch.tensor(self.dt_list).to(self.dtype)
+
+        self.velocity_predictions = []
+
+
+    @property
+    def shift(self):
+        """
+        The value used for shifting.
+        """
+        return self._shift
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+
+
+    def set_shift(self, shift: float):
+        self._shift = shift
+    
+    def set_begin_index(self, begin_index: int):
+        """
+        Set the begin index for the scheduler.
+
+        Args:
+            begin_index (`int`):
+                The begin index to set.
+        """
+        self._begin_index = begin_index
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward process in flow-matching
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=self.dtype)
+            timestep = timestep.to(sample.device, dtype=self.dtype)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+
+        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+
+        return sample
+
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+    
+    def index_for_timestep(self, timestep, schedule_timesteps):
+        """
+        Get the index for a given timestep in the schedule.
+
+        Args:
+            timestep (`torch.Tensor`):
+                The timestep to find the index for.
+            schedule_timesteps (`torch.Tensor`):
+                The schedule timesteps.
+
+        Returns:
+            `int`:
+                The index for the timestep.
+        """
+        # Find the closest timestep in the schedule
+        indices = torch.searchsorted(schedule_timesteps, timestep, right=True)
+        indices = torch.clamp(indices, 0, len(schedule_timesteps) - 1)
+        return indices.item()
+    
+
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor = None,
+        return_dict: bool = True,
+        **kwargs
+    ) -> torch.Tensor:
+        '''
+        One step of the STORK update for flow matching or noise-based diffusion models.
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a [`~schedulers.STORKSchedulerOutput`] instead of a plain tuple.
+                
+        Returns:
+            result (Union[Tuple, STORKSchedulerOutput]):
+                The next sample in the diffusion chain, either as a tuple or as a [`~schedulers.STORKSchedulerOutput`]. The value is converted back to the original dtype of `model_output` to avoid numerical issues.
+        '''
+        original_model_output_dtype = model_output.dtype
+        # Cast model_output and sample to "torch.float32" to avoid numerical issues
+        model_output = model_output.to(self.dtype)
+        sample = sample.to(self.dtype)
+        # Move sample to model_output's device
+        sample = sample.to(model_output.device)
+        
+        """
+        self.velocity_predictions always contain upcasted model_output in torch.float32 dtype.
+        """
+        
+        if self.prediction_type == "epsilon":
+            if self.solver_order == 2:
+                result = self.step_noise_2(model_output, timestep, sample, return_dict)
+            elif self.solver_order == 4:
+                result = self.step_noise_4(model_output, timestep, sample, return_dict)
+            else:
+                raise ValueError(f"Solver order {self.solver_order} is not yet supported for noise-based models")
+        elif self.prediction_type == "flow_prediction":
+            if self.solver_order == 1:
+                result = self.step_flow_matching_1(model_output, timestep, sample, return_dict)
+            elif self.solver_order == 2:
+                result = self.step_flow_matching_2(model_output, timestep, sample, return_dict)
+            elif self.solver_order == 4:
+                result = self.step_flow_matching_4(model_output, timestep, sample, return_dict)
+            else:
+                raise ValueError(f"Solver order {self.solver_order} is not yet supported for flow matching models")
+        else:
+            raise ValueError(f"Prediction type {self.prediction_type} is not yet supported")
+        
+        # Convert the result back to the original dtype of model_output, as this result will be used as the next input to the model
+        if return_dict:
+            result.prev_sample = result.prev_sample.to(original_model_output_dtype)
+        else:
+            result = (result[0].to(original_model_output_dtype),)
+        return result
+        
+
+    def step_flow_matching_1(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor = None,
+        return_dict: bool = False
+    ) -> torch.Tensor:
+        # Initialize the step index if it's the first step
+        if self._step_index is None:
+            self._step_index = 0
+
+
+        # Compute the startup phase or the derivative approximation for the main step
+        if self._step_index == 0:
+            img_next = sample - model_output * self.dt_list[self._step_index]
+            self._step_index += 1
+            self.velocity_predictions.append(model_output)
+
+            if not return_dict:
+                return (img_next,)
+            return STORKSchedulerOutput(prev_sample=img_next)
+        else:
+            t = self.sigmas[self._step_index]
+            t_start = torch.ones(model_output.shape, device=sample.device) * t
+            t_next = self.sigmas[self._step_index + 1]
+            
+            h1 = self.dt_list[self._step_index-1]
+
+            if self.derivative_order == 1:
+                # Ensure h1 is a tensor for proper broadcasting
+                h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                velocity_derivative = (self.velocity_predictions[-1] - model_output) / h1_tensor
+                velocity_second_derivative = None
+                velocity_third_derivative = None
+            elif self.derivative_order == 2:
+                # Ensure h1 and h2 are tensors for proper broadcasting
+                h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                if self._step_index == 1:
+                    img_next = sample - 1.5 * model_output * self.dt_list[self._step_index] + 0.5 * self.velocity_predictions[-1] * self.dt_list[self._step_index-1]
+                    self._step_index += 1
+                    self.velocity_predictions.append(model_output)
+
+                    if not return_dict:
+                        return (img_next,)
+                    return STORKSchedulerOutput(prev_sample=img_next)
+                else:
+                    h2 = self.dt_list[self._step_index-2]
+                    h2_tensor = torch.tensor(h2, device=model_output.device, dtype=model_output.dtype)
+                    velocity_derivative = (-self.velocity_predictions[-2] + 4 * self.velocity_predictions[-1] - 3 * model_output) / (2 * h1_tensor)
+                    velocity_second_derivative = 2 / (h1_tensor * h2_tensor * (h1_tensor + h2_tensor)) * (self.velocity_predictions[-2] * h1_tensor - self.velocity_predictions[-1] * (h1_tensor + h2_tensor) + model_output * h2_tensor)
+                    velocity_third_derivative = None
+            elif self.derivative_order == 3:
+
+                if self._step_index == 1 or self._step_index == 2:
+                    img_next = sample - 1.5 * model_output * self.dt_list[self._step_index] + 0.5 * self.velocity_predictions[-1] * self.dt_list[self._step_index-1]
+                    self._step_index += 1
+                    self.velocity_predictions.append(model_output)
+
+                    if not return_dict:
+                        return (img_next,)
+                    return STORKSchedulerOutput(prev_sample=img_next)
+                else:
+                    h2 = h1 + self.dt_list[self._step_index-2]
+                    h3 = h2 + self.dt_list[self._step_index-3]
+                    # Ensure h1, h2, and h3 are tensors for proper broadcasting
+                    h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                    h2_tensor = torch.tensor(h2, device=model_output.device, dtype=model_output.dtype)
+                    h3_tensor = torch.tensor(h3, device=model_output.device, dtype=model_output.dtype)
+                    velocity_derivative = ((h2_tensor * h3_tensor) * (self.velocity_predictions[-1] - model_output) - (h1_tensor * h3_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor * h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+                    velocity_second_derivative = 2 * ((h2_tensor + h3_tensor) * (self.velocity_predictions[-1] - model_output) - (h1_tensor + h3_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor + h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+                    velocity_third_derivative = 6 * ((h2_tensor - h3_tensor) * (self.velocity_predictions[-1] - model_output) + (h3_tensor - h1_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor - h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+            else:
+                print("The noise approximation order is not supported!")
+                exit()
+            
+            self.velocity_predictions.append(model_output)
+            self._step_index += 1
+
+
+
+        Y_j_2 = sample
+        Y_j_1 = sample
+        Y_j = sample
+
+        
+        # Implementation of our Runge-Kutta-Gegenbauer second order method
+        for j in range(1, self.s + 1):
+            # Calculate the corresponding \bar{alpha}_t and beta_t that aligns with the correct timestep
+            fraction = (j - 1) * (j + 2) / (self.s * (self.s + 3))
+            
+            if j == 1:
+                mu_tilde = 4 / (self.s * (self.s + 1))
+                dt = (t - t_next) * torch.ones(model_output.shape, device=sample.device)
+                Y_j = Y_j_1 - dt * mu_tilde * model_output
+            else:
+                mu = (2 * j + 1) * self.coeff_rock1(j) / (j * self.coeff_rock1(j - 1))
+                nu = -(j + 1) * self.coeff_rock1(j) / (j * self.coeff_rock1(j - 2))
+                mu_tilde = mu * 4 / (self.s * (self.s + 1))
+
+
+                # Probability flow ODE update
+                diff = -fraction * (t - t_next) * torch.ones(model_output.shape, device=sample.device)
+                velocity = self.taylor_approximation(self.derivative_order, diff, model_output, velocity_derivative, velocity_second_derivative, velocity_third_derivative)
+                Y_j = mu * Y_j_1 + nu * Y_j_2 - dt * mu_tilde * velocity
+                
+            Y_j_2 = Y_j_1
+            Y_j_1 = Y_j
+
+
+
+        img_next = Y_j
+        img_next = img_next.to(model_output.dtype)
+
+        return SchedulerOutput(prev_sample=img_next)
+    
+
+
+
+    def step_flow_matching_2(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor = None,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        '''
+        One step of the STORK2 update for flow matching based models.
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a [`~schedulers.STORKSchedulerOutput`] instead of a plain tuple.
+                
+        Returns:
+            result (Union[Tuple, STORKSchedulerOutput]):
+                The next sample in the diffusion chain, either as a tuple or as a [`~schedulers.STORKSchedulerOutput`]. The value is converted back to the original dtype of `model_output` to avoid numerical issues.
+        '''
+        # Initialize the step index if it's the first step
+        if self._step_index is None:
+            self._step_index = 0
+
+
+        # Compute the startup phase or the derivative approximation for the main step
+        if self._step_index == 0:
+            img_next = sample - model_output * self.dt_list[self._step_index]
+            self._step_index += 1
+            self.velocity_predictions.append(model_output)
+
+            if not return_dict:
+                return (img_next,)
+            return STORKSchedulerOutput(prev_sample=img_next)
+        else:
+            t = self.sigmas[self._step_index]
+            t_start = torch.ones(model_output.shape, device=sample.device) * t
+            t_next = self.sigmas[self._step_index + 1]
+            
+            h1 = self.dt_list[self._step_index-1]
+
+            if self.derivative_order == 1:
+                # Ensure h1 is a tensor for proper broadcasting
+                h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                velocity_derivative = (self.velocity_predictions[-1] - model_output) / h1_tensor
+                velocity_second_derivative = None
+                velocity_third_derivative = None
+            elif self.derivative_order == 2:
+                # Ensure h1 and h2 are tensors for proper broadcasting
+                h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                if self._step_index == 1:
+                    img_next = sample - 1.5 * model_output * self.dt_list[self._step_index] + 0.5 * self.velocity_predictions[-1] * self.dt_list[self._step_index-1]
+                    self._step_index += 1
+                    self.velocity_predictions.append(model_output)
+
+                    if not return_dict:
+                        return (img_next,)
+                    return STORKSchedulerOutput(prev_sample=img_next)
+                else:
+                    h2 = self.dt_list[self._step_index-2]
+                    h2_tensor = torch.tensor(h2, device=model_output.device, dtype=model_output.dtype)
+                    velocity_derivative = (-self.velocity_predictions[-2] + 4 * self.velocity_predictions[-1] - 3 * model_output) / (2 * h1_tensor)
+                    velocity_second_derivative = 2 / (h1_tensor * h2_tensor * (h1_tensor + h2_tensor)) * (self.velocity_predictions[-2] * h1_tensor - self.velocity_predictions[-1] * (h1_tensor + h2_tensor) + model_output * h2_tensor)
+                    velocity_third_derivative = None
+            elif self.derivative_order == 3:
+
+                if self._step_index == 1 or self._step_index == 2:
+                    img_next = sample - 1.5 * model_output * self.dt_list[self._step_index] + 0.5 * self.velocity_predictions[-1] * self.dt_list[self._step_index-1]
+                    self._step_index += 1
+                    self.velocity_predictions.append(model_output)
+
+                    if not return_dict:
+                        return (img_next,)
+                    return STORKSchedulerOutput(prev_sample=img_next)
+                else:
+                    h2 = h1 + self.dt_list[self._step_index-2]
+                    h3 = h2 + self.dt_list[self._step_index-3]
+                    # Ensure h1, h2, and h3 are tensors for proper broadcasting
+                    h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                    h2_tensor = torch.tensor(h2, device=model_output.device, dtype=model_output.dtype)
+                    h3_tensor = torch.tensor(h3, device=model_output.device, dtype=model_output.dtype)
+                    velocity_derivative = ((h2_tensor * h3_tensor) * (self.velocity_predictions[-1] - model_output) - (h1_tensor * h3_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor * h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+                    velocity_second_derivative = 2 * ((h2_tensor + h3_tensor) * (self.velocity_predictions[-1] - model_output) - (h1_tensor + h3_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor + h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+                    velocity_third_derivative = 6 * ((h2_tensor - h3_tensor) * (self.velocity_predictions[-1] - model_output) + (h3_tensor - h1_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor - h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+            else:
+                print("The noise approximation order is not supported!")
+                exit()
+            
+            self.velocity_predictions.append(model_output)
+            self._step_index += 1
+
+
+        Y_j_2 = sample
+        Y_j_1 = sample
+        Y_j = sample
+
+        
+        # Implementation of our Runge-Kutta-Gegenbauer second order method
+        for j in range(1, self.s + 1):
+            # Calculate the corresponding \bar{alpha}_t and beta_t that aligns with the correct timestep
+            if j > 1:
+                if j == 2:
+                    fraction = 4 / (3 * (self.s**2 + self.s - 2))
+                else:
+                    fraction = ((j - 1)**2 + (j - 1) - 2) / (self.s**2 + self.s - 2)
+            
+            if j == 1:
+                mu_tilde = 6 / ((self.s + 4) * (self.s - 1))
+                dt = (t - t_next) * torch.ones(model_output.shape, device=sample.device)
+                Y_j = Y_j_1 - dt * mu_tilde * model_output
+            else:
+                mu = (2 * j + 1) * self.b_coeff(j) / (j * self.b_coeff(j - 1))
+                nu = -(j + 1) * self.b_coeff(j) / (j * self.b_coeff(j - 2))
+                mu_tilde = mu * 6 / ((self.s + 4) * (self.s - 1))
+                gamma_tilde = -mu_tilde * (1 - j * (j + 1) * self.b_coeff(j-1)/ 2)
+
+
+                # Probability flow ODE update
+                diff = -fraction * (t - t_next) * torch.ones(model_output.shape, device=sample.device)
+                velocity = self.taylor_approximation(self.derivative_order, diff, model_output, velocity_derivative, velocity_second_derivative, velocity_third_derivative)
+                Y_j = mu * Y_j_1 + nu * Y_j_2 + (1 - mu - nu) * sample - dt * mu_tilde * velocity - dt * gamma_tilde * model_output
+                
+            Y_j_2 = Y_j_1
+            Y_j_1 = Y_j
+
+
+
+        img_next = Y_j
+        img_next = img_next.to(model_output.dtype)
+
+        if not return_dict:
+            return (img_next,) 
+        return STORKSchedulerOutput(prev_sample=img_next)
+
+
+    def step_flow_matching_4(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor = None,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        '''
+        One step of the STORK4 update for flow matching models
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.FloatTensor`: The next sample in the diffusion chain.
+        '''        
+        # Initialize the step index if it's the first step
+        if self._step_index is None:
+            self._step_index = 0
+
+
+        # Compute the startup phase or the derivative approximation for the main step
+        if self._step_index == 0:
+            img_next = sample - model_output * self.dt_list[self._step_index]
+            self._step_index += 1
+            self.velocity_predictions.append(model_output)
+
+            if not return_dict:
+                return (img_next,)
+            return STORKSchedulerOutput(prev_sample=img_next)
+        else:
+            t = self.sigmas[self._step_index]
+            t_start = torch.ones(model_output.shape, device=sample.device) * t
+            t_next = self.sigmas[self._step_index + 1]
+            
+            h1 = self.dt_list[self._step_index-1]
+
+            if self.derivative_order == 1:
+                # Ensure h1 is a tensor for proper broadcasting
+                h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                velocity_derivative = (self.velocity_predictions[-1] - model_output) / h1_tensor
+                velocity_second_derivative = None
+                velocity_third_derivative = None
+            elif self.derivative_order == 2:
+                # Ensure h1 and h2 are tensors for proper broadcasting
+                h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                if self._step_index == 1:
+                    img_next = sample - 1.5 * model_output * self.dt_list[self._step_index] + 0.5 * self.velocity_predictions[-1] * self.dt_list[self._step_index-1]
+                    self._step_index += 1
+                    self.velocity_predictions.append(model_output)
+
+                    if not return_dict:
+                        return (img_next,)
+                    return STORKSchedulerOutput(prev_sample=img_next)
+                else:
+                    h2 = self.dt_list[self._step_index-2]
+                    h2_tensor = torch.tensor(h2, device=model_output.device, dtype=model_output.dtype)
+                    velocity_derivative = (-self.velocity_predictions[-2] + 4 * self.velocity_predictions[-1] - 3 * model_output) / (2 * h1_tensor)
+                    velocity_second_derivative = 2 / (h1_tensor * h2_tensor * (h1_tensor + h2_tensor)) * (self.velocity_predictions[-2] * h1_tensor - self.velocity_predictions[-1] * (h1_tensor + h2_tensor) + model_output * h2_tensor)
+                    velocity_third_derivative = None
+            elif self.derivative_order == 3:
+
+                if self._step_index == 1 or self._step_index == 2:
+                    img_next = sample - 1.5 * model_output * self.dt_list[self._step_index] + 0.5 * self.velocity_predictions[-1] * self.dt_list[self._step_index-1]
+                    self._step_index += 1
+                    self.velocity_predictions.append(model_output)
+
+                    if not return_dict:
+                        return (img_next,)
+                    return STORKSchedulerOutput(prev_sample=img_next)
+                else:
+                    h2 = h1 + self.dt_list[self._step_index-2]
+                    h3 = h2 + self.dt_list[self._step_index-3]
+                    # Ensure h1, h2, and h3 are tensors for proper broadcasting
+                    h1_tensor = torch.tensor(h1, device=model_output.device, dtype=model_output.dtype)
+                    h2_tensor = torch.tensor(h2, device=model_output.device, dtype=model_output.dtype)
+                    h3_tensor = torch.tensor(h3, device=model_output.device, dtype=model_output.dtype)
+                    velocity_derivative = ((h2_tensor * h3_tensor) * (self.velocity_predictions[-1] - model_output) - (h1_tensor * h3_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor * h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+                    velocity_second_derivative = 2 * ((h2_tensor + h3_tensor) * (self.velocity_predictions[-1] - model_output) - (h1_tensor + h3_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor + h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+                    velocity_third_derivative = 6 * ((h2_tensor - h3_tensor) * (self.velocity_predictions[-1] - model_output) + (h3_tensor - h1_tensor) * (self.velocity_predictions[-2] - model_output) + (h1_tensor - h2_tensor) * (self.velocity_predictions[-3] - model_output)) / (h1_tensor * h2_tensor * h3_tensor)
+            else:
+                print("The noise approximation order is not supported!")
+                exit()
+            
+            self.velocity_predictions.append(model_output)
+            self._step_index += 1
+
+        Y_j_2 = sample
+        Y_j_1 = sample
+        Y_j = sample
+
+        ci1 = t_start
+        ci2 = t_start
+        ci3 = t_start
+
+        # Coefficients of ROCK4
+        ms, fpa, fpb, fpbe, recf = self.coeff_rock4()
+        # Choose the degree that's in the precomputed table
+        mdeg, mp = self.mdegr(self.s, ms)
+        mz = int(mp[0])
+        mr = int(mp[1])
+
+        '''
+        The first part of the STORK4 update
+        '''
+        for j in range(1, mdeg + 1):
+
+            # First sub-step in the first part of the STORK4 update
+            if j == 1:
+                temp1 = -(t - t_next) * recf[mr] * torch.ones(model_output.shape, device=sample.device)
+                ci1 = t_start + temp1
+                ci2 = ci1
+                Y_j_1 = sample + temp1 * model_output
+                # Y_j = sample + temp1 * model_output
+            # Second and the following sub-steps in the first part of the STORK4 update
+            else:
+                diff = ci1 - t_start
+                velocity = self.taylor_approximation(self.derivative_order, diff, model_output, velocity_derivative, velocity_second_derivative, velocity_third_derivative)
+
+                temp1 = -(t - t_next) * recf[mr + 2 * (j-2) + 1] * torch.ones(model_output.shape, device=sample.device)
+                temp3 = -recf[mr + 2 * (j-2) + 2] * torch.ones(model_output.shape, device=sample.device)
+                temp2 = torch.ones(model_output.shape, device=sample.device) - temp3
+
+                ci1 = temp1 + temp2 * ci2 + temp3 * ci3
+                Y_j = temp1 * velocity + temp2 * Y_j_1 + temp3 * Y_j_2
+
+            # Update the intermediate variables
+            Y_j_2 = Y_j_1
+            Y_j_1 = Y_j
+
+            ci3 = ci2
+            ci2 = ci1
+
+        '''
+        The finishing four-step procedure as a composition method
+        '''
+        # First finishing step
+        temp1 = -(t - t_next) * fpa[mz,0] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci1 - t_start
+        velocity = self.taylor_approximation(self.derivative_order, diff, model_output, velocity_derivative, velocity_second_derivative, velocity_third_derivative)
+        Y_j_1 = velocity
+        Y_j_3 = Y_j + temp1 * Y_j_1
+
+        # Second finishing step
+        ci2 = ci1 + temp1
+        temp1 = -(t - t_next) * fpa[mz,1] * torch.ones(model_output.shape, device=sample.device)
+        temp2 = -(t - t_next) * fpa[mz,2] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci2 - t_start
+        velocity = self.taylor_approximation(self.derivative_order, diff, model_output, velocity_derivative, velocity_second_derivative, velocity_third_derivative)
+        Y_j_2 = velocity
+        Y_j_4 = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2
+
+        # Third finishing step
+        ci2 = ci1 + temp1 + temp2
+        temp1 = -(t - t_next) * fpa[mz,3] * torch.ones(model_output.shape, device=sample.device)
+        temp2 = -(t - t_next) * fpa[mz,4] * torch.ones(model_output.shape, device=sample.device)
+        temp3 = -(t - t_next) * fpa[mz,5] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci2 - t_start
+        velocity = self.taylor_approximation(self.derivative_order, diff, model_output, velocity_derivative, velocity_second_derivative, velocity_third_derivative)
+        Y_j_3 = velocity
+        # This is the counterpart of the final step in the noise-based diffusion models STORK4
+        # fnt = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2 + temp3 * Y_j_3
+
+        # Fourth finishing step
+        ci2 = ci1 + temp1 + temp2 + temp3
+        temp1 = -(t - t_next) * fpb[mz,0] * torch.ones(model_output.shape, device=sample.device)
+        temp2 = -(t - t_next) * fpb[mz,1] * torch.ones(model_output.shape, device=sample.device)
+        temp3 = -(t - t_next) * fpb[mz,2] * torch.ones(model_output.shape, device=sample.device)
+        temp4 = -(t - t_next) * fpb[mz,3] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci2 - t_start
+        velocity = self.taylor_approximation(self.derivative_order, diff, model_output, velocity_derivative, velocity_second_derivative, velocity_third_derivative)
+        Y_j_4 = velocity
+        Y_j = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2 + temp3 * Y_j_3 + temp4 * Y_j_4
+        img_next = Y_j
+
+        if not return_dict:
+            return (img_next,)
+        return STORKSchedulerOutput(prev_sample=img_next)
+    
+
+    def step_noise_2(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor = None,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        '''
+        One step of the STORK2 update for noise-based diffusion models.
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a [`~schedulers.STORKSchedulerOutput`] instead of a plain tuple.
+
+        Returns:
+            `torch.FloatTensor`: The next sample in the diffusion chain.
+        '''
+        # Initialize the step index if it's the first step
+        if self._step_index is None:
+            self._step_index = 0
+            self.initial_noise = model_output
+
+
+        total_step = self.config.num_train_timesteps
+        t = self.timesteps[self._step_index] / total_step
+
+        beta_0, beta_1 = self.betas[0], self.betas[-1]
+        t_start = torch.ones(model_output.shape, device=sample.device) * t
+        beta_t = (beta_0 + t_start * (beta_1 - beta_0)) * total_step
+        log_mean_coeff = (-0.25 * t_start ** 2 * (beta_1 - beta_0) - 0.5 * t_start * beta_0) * total_step
+        std = torch.sqrt(1. - torch.exp(2. * log_mean_coeff))
+
+        # Tweedie's trick
+        if self._step_index == len(self.timesteps) - 1:
+            noise_last = model_output
+            img_next = sample - std * noise_last
+            if not return_dict:
+                return (img_next,)
+            return STORKSchedulerOutput(prev_sample=img_next)
+        
+        t_next = self.timesteps[self._step_index + 1] / total_step
+
+        # drift, diffusion -> f(x,t), g(t)
+        drift_initial, diffusion_initial = -0.5 * beta_t * sample, torch.sqrt(beta_t) * torch.ones(sample.shape, device=sample.device)
+        noise_initial = model_output
+        score = -noise_initial / std  # score -> noise
+        drift_initial = drift_initial - diffusion_initial ** 2 * score * 0.5 # drift -> dx/dt
+
+
+        dt = torch.ones(model_output.shape, device=sample.device) * self.dt
+
+        if self._step_index == 0:
+            # FIRST RUN
+            self.initial_sample = sample
+            img_next = sample - 0.5 * dt * drift_initial
+
+            self.noise_predictions.append(noise_initial)
+            self._step_index += 1
+
+            self.initial_sample = sample
+            self.initial_drift = drift_initial
+            self.initial_noise = model_output
+
+            return SchedulerOutput(prev_sample=img_next)
+        elif self._step_index == 1:
+            # SECOND RUN
+            t_previous = torch.ones(model_output.shape, device=sample.device) * self.timesteps[0] / 1000
+            drift_previous = self.drift_function(self.betas, self.config.num_train_timesteps, t_previous, self.initial_sample, self.noise_predictions[-1])
+
+            img_next = sample - 0.75 * dt * drift_initial + 0.25 * dt * drift_previous
+
+            self.noise_predictions.append(noise_initial)
+            self._step_index += 1
+
+            return SchedulerOutput(prev_sample=img_next)
+        elif self._step_index == 2:
+            h = 0.5 * dt
+    
+            noise_derivative = (3 * self.noise_predictions[0] - 4 * self.noise_predictions[1] + model_output) / (2 * h)
+            noise_second_derivative = (self.noise_predictions[0] - 2 * self.noise_predictions[1] + model_output) / (h ** 2)
+            noise_third_derivative = None
+
+            model_output = self.initial_noise
+            drift_initial = self.initial_drift
+            sample = self.initial_sample
+
+            t = self.timesteps[0] / total_step
+            t_start = torch.ones(model_output.shape, device=sample.device) * t
+            t_next = self.timesteps[2] / total_step
+        elif self._step_index == 3:
+            h = 0.5 * dt
+
+            noise_derivative = (-3 * noise_initial + 4 * self.noise_predictions[-1] - self.noise_predictions[-2]) / (2 * h)
+            noise_second_derivative = (noise_initial - 2 * self.noise_predictions[-1] + self.noise_predictions[-2]) / (h ** 2)
+            noise_third_derivative = None
+
+            self.noise_predictions.append(noise_initial)
+        elif self._step_index == 4:
+            h = dt
+
+            noise_derivative = (-3 * noise_initial + 4 * self.noise_predictions[-1] - self.noise_predictions[-2]) / (2 * h)
+            noise_second_derivative = (noise_initial - 2 * self.noise_predictions[-1] + self.noise_predictions[-2]) / (h ** 2)
+            noise_third_derivative = None
+            
+            self.noise_predictions.append(noise_initial)
+        else:
+            # ALL ELSE
+            h = dt
+            
+            noise_derivative = (2 * self.noise_predictions[-3] - 9 * self.noise_predictions[-2] + 18 * self.noise_predictions[-1] - 11 * noise_initial) / (6 * h)
+            noise_second_derivative = (-self.noise_predictions[-3] + 4 * self.noise_predictions[-2] -5 * self.noise_predictions[-1] + 2 * noise_initial) / (h**2)
+            noise_third_derivative = (self.noise_predictions[-3] - 3 * self.noise_predictions[-2] + 3 * self.noise_predictions[-1] - noise_initial) / (h**3)
+
+            self.noise_predictions.append(noise_initial)
+
+
+        Y_j_2 = sample
+        Y_j_1 = sample
+        Y_j = sample
+
+        # Implementation of our Runge-Kutta-Gegenbauer second order method
+        for j in range(1, self.s + 1):
+            # Calculate the corresponding \bar{alpha}_t and beta_t that aligns with the correct timestep
+            if j > 1:
+                if j == 2:
+                    fraction = 4 / (3 * (self.s**2 + self.s - 2))
+                else:
+                    fraction = ((j - 1)**2 + (j - 1) - 2) / (self.s**2 + self.s - 2)
+            
+            if j == 1:
+                mu_tilde = 6 / ((self.s + 4) * (self.s - 1))
+                dt = (t - t_next) * torch.ones(model_output.shape, device=sample.device)
+                Y_j = Y_j_1 - dt * mu_tilde * model_output
+            else:
+                mu = (2 * j + 1) * self.b_coeff(j) / (j * self.b_coeff(j - 1))
+                nu = -(j + 1) * self.b_coeff(j) / (j * self.b_coeff(j - 2))
+                mu_tilde = mu * 6 / ((self.s + 4) * (self.s - 1))
+                gamma_tilde = -mu_tilde * (1 - j * (j + 1) * self.b_coeff(j-1)/ 2)
+
+
+                # Probability flow ODE update
+                diff = -fraction * (t - t_next) * torch.ones(model_output.shape, device=sample.device)
+                velocity = self.taylor_approximation(self.derivative_order, diff, model_output, noise_derivative, noise_second_derivative, noise_third_derivative)
+                Y_j = mu * Y_j_1 + nu * Y_j_2 + (1 - mu - nu) * sample - dt * mu_tilde * velocity - dt * gamma_tilde * model_output
+                
+            Y_j_2 = Y_j_1
+            Y_j_1 = Y_j
+
+
+
+        img_next = Y_j
+        img_next = img_next.to(model_output.dtype)
+        self._step_index += 1
+
+        if not return_dict:
+            return (img_next,)
+        return STORKSchedulerOutput(prev_sample=img_next)
+
+
+    def step_noise_4(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor = None,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        '''
+        One step of the STORK4 update for noise-based diffusion models.
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a [`~schedulers.STORKSchedulerOutput`] instead of a plain tuple.
+
+        Returns:
+            `torch.FloatTensor`: The next sample in the diffusion chain.
+        '''
+
+
+
+        # Initialize the step index if it's the first step
+        if self._step_index is None:
+            self._step_index = 0
+            self.initial_noise = model_output
+
+
+        total_step = self.config.num_train_timesteps
+        t = self.timesteps[self._step_index] / total_step
+
+        beta_0, beta_1 = self.betas[0], self.betas[-1]
+        t_start = torch.ones(model_output.shape, device=sample.device) * t
+        beta_t = (beta_0 + t_start * (beta_1 - beta_0)) * total_step
+        log_mean_coeff = (-0.25 * t_start ** 2 * (beta_1 - beta_0) - 0.5 * t_start * beta_0) * total_step
+        std = torch.sqrt(1. - torch.exp(2. * log_mean_coeff))
+
+        # Tweedie's trick
+        if self._step_index == len(self.timesteps) - 1:
+            noise_last = model_output
+            img_next = sample - std * noise_last
+            if not return_dict:
+                return (img_next,)
+            return STORKSchedulerOutput(prev_sample=img_next)
+        
+        t_next = self.timesteps[self._step_index + 1] / total_step
+
+        # drift, diffusion -> f(x,t), g(t)
+        drift_initial, diffusion_initial = -0.5 * beta_t * sample, torch.sqrt(beta_t) * torch.ones(sample.shape, device=sample.device)
+        noise_initial = model_output
+        score = -noise_initial / std  # score -> noise
+        drift_initial = drift_initial - diffusion_initial ** 2 * score * 0.5 # drift -> dx/dt
+
+
+        dt = torch.ones(model_output.shape, device=sample.device) * self.dt
+
+
+        if self.derivative_order == 2:
+            if self._step_index == 0:
+                # Initial Euler update
+                self.initial_sample = sample
+                img_next = sample - dt * drift_initial
+
+                self.noise_predictions.append(noise_initial)
+                self._step_index += 1
+
+                self.initial_drift = drift_initial
+                
+                if not return_dict:
+                    return (img_next,)
+                return SchedulerOutput(prev_sample=img_next)
+            elif self._step_index == 1:
+                # Initial 2-step Adams-Bashforth update
+                drift_previous = self.initial_drift
+
+                img_next = sample - 1.5 * dt * drift_initial + 0.5 * dt * drift_previous
+
+                self.noise_predictions.append(noise_initial)
+                self._step_index += 1
+
+                if not return_dict:
+                    return (img_next,)
+                return SchedulerOutput(prev_sample=img_next)
+            else:
+                # STORK4 update
+                h = dt
+
+                # The first derivative is calculated using the three point approximation, 
+                # and the second derivative is calculated using the standardtwo point approximation.
+                noise_derivative = (-self.noise_predictions[-2] + 4 * self.noise_predictions[-1] - 3 * noise_initial) / (2 * h)
+                noise_second_derivative = (self.noise_predictions[-2] - 2 * self.noise_predictions[-1] + noise_initial) / h**2
+                noise_third_derivative = None
+
+                self.noise_predictions.append(noise_initial)
+                noise_approx_order = 2
+        elif self.derivative_order == 1:
+            if self._step_index == 0:
+                # Initial Euler update
+                self.initial_sample = sample
+                img_next = sample - dt * drift_initial
+
+                self.noise_predictions.append(noise_initial)
+                self._step_index += 1
+
+                self.initial_drift = drift_initial
+                
+                if not return_dict:
+                    return (img_next,)
+                return SchedulerOutput(prev_sample=img_next)
+            else:
+                # STORK4 update
+                h = dt
+
+                noise_derivative = (self.noise_predictions[-1] - noise_initial) / h
+                noise_second_derivative = None
+                noise_third_derivative = None
+
+                self.noise_predictions.append(noise_initial)
+                noise_approx_order = 1
+        else:
+            raise ValueError(f"Unknown derivative order: {self.derivative_order}")
+
+
+        Y_j_2 = sample
+        Y_j_1 = sample
+        Y_j = sample
+
+        ci1 = t_start
+        ci2 = t_start
+        ci3 = t_start
+
+        # Coefficients of ROCK4
+        ms, fpa, fpb, fpbe, recf = self.coeff_rock4()
+        # Choose the degree that's in the precomputed table
+        mdeg, mp = self.mdegr(self.s, ms)
+        mz = int(mp[0])
+        mr = int(mp[1])
+
+        '''
+        The first part of the STORK4 update
+        '''
+        for j in range(1, mdeg + 1):
+
+            # First sub-step in the first part of the STORK4 update
+            if j == 1:
+                temp1 = -(t - t_next) * recf[mr] * torch.ones(model_output.shape, device=sample.device)
+                ci1 = t_start + temp1
+                ci2 = ci1
+                Y_j_1 = sample + temp1 * model_output #subver
+                
+                # drift_approx = self.drift_function(self.betas, self.config.num_train_timesteps, t_start, Y_j, model_output)
+                # Y_j_1 = sample + temp1 * drift_approx
+                
+            # Second and the following sub-steps in the first part of the STORK4 update
+            else:
+                diff = ci1 - t_start
+                noise_approx = self.taylor_approximation(noise_approx_order, diff, model_output, noise_derivative, noise_second_derivative, noise_third_derivative)
+                drift_approx = self.drift_function(self.betas, self.config.num_train_timesteps, ci1, Y_j_1, noise_approx)
+
+                temp1 = -(t - t_next) * recf[mr + 2 * (j-2) + 1] * torch.ones(model_output.shape, device=sample.device)
+                temp3 = -recf[mr + 2 * (j-2) + 2] * torch.ones(model_output.shape, device=sample.device)
+                temp2 = torch.ones(model_output.shape, device=sample.device) - temp3
+
+                ci1 = temp1 + temp2 * ci2 + temp3 * ci3
+                Y_j = temp1 * drift_approx + temp2 * Y_j_1 + temp3 * Y_j_2
+
+            # Update the intermediate variables
+            Y_j_2 = Y_j_1
+            Y_j_1 = Y_j
+
+            ci3 = ci2
+            ci2 = ci1
+
+        '''
+        The finishing four-step procedure as a composition method
+        '''
+        # First finishing step
+        temp1 = -(t - t_next) * fpa[mz,0] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci1 - t_start
+        noise_approx = self.taylor_approximation(noise_approx_order, diff, model_output, noise_derivative, noise_second_derivative, noise_third_derivative)
+        drift_approx = self.drift_function(self.betas, self.config.num_train_timesteps, ci1, Y_j, noise_approx)
+        Y_j_1 = drift_approx
+        Y_j_3 = Y_j + temp1 * Y_j_1
+
+        # Second finishing step
+        ci2 = ci1 + temp1
+        temp1 = -(t - t_next) * fpa[mz,1] * torch.ones(model_output.shape, device=sample.device)
+        temp2 = -(t - t_next) * fpa[mz,2] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci2 - t_start
+        noise_approx = self.taylor_approximation(noise_approx_order, diff, model_output, noise_derivative, noise_second_derivative, noise_third_derivative)
+        drift_approx = self.drift_function(self.betas, self.config.num_train_timesteps, ci2, Y_j_3, noise_approx)
+        Y_j_2 = drift_approx
+        Y_j_4 = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2
+
+        # Third finishing step
+        ci2 = ci1 + temp1 + temp2
+        temp1 = -(t - t_next) * fpa[mz,3] * torch.ones(model_output.shape, device=sample.device)
+        temp2 = -(t - t_next) * fpa[mz,4] * torch.ones(model_output.shape, device=sample.device)
+        temp3 = -(t - t_next) * fpa[mz,5] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci2 - t_start
+        noise_approx = self.taylor_approximation(noise_approx_order, diff, model_output, noise_derivative, noise_second_derivative, noise_third_derivative)
+        drift_approx = self.drift_function(self.betas, self.config.num_train_timesteps, ci2, Y_j_4, noise_approx)
+        Y_j_3 = drift_approx
+        fnt = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2 + temp3 * Y_j_3
+
+        # Fourth finishing step
+        ci2 = ci1 + temp1 + temp2 + temp3
+        temp1 = -(t - t_next) * fpb[mz,0] * torch.ones(model_output.shape, device=sample.device)
+        temp2 = -(t - t_next) * fpb[mz,1] * torch.ones(model_output.shape, device=sample.device)
+        temp3 = -(t - t_next) * fpb[mz,2] * torch.ones(model_output.shape, device=sample.device)
+        temp4 = -(t - t_next) * fpb[mz,3] * torch.ones(model_output.shape, device=sample.device)
+        diff = ci2 - t_start
+        noise_approx = self.taylor_approximation(noise_approx_order, diff, model_output, noise_derivative, noise_second_derivative, noise_third_derivative)
+        drift_approx = self.drift_function(self.betas, self.config.num_train_timesteps, ci2, fnt, noise_approx)
+        Y_j_4 = drift_approx
+        Y_j = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2 + temp3 * Y_j_3 + temp4 * Y_j_4
+
+
+
+        img_next = Y_j
+        self._step_index += 1
+
+        if not return_dict:
+            return (img_next,)
+        return STORKSchedulerOutput(prev_sample=img_next)
+        
+
+
+
+    def __len__(self):
+        return self.config.num_train_timesteps
+    
+    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+    
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        """
+        Add noise to the original samples according to the noise magnitude at the given timestep.
+
+        Args:
+            original_samples (`torch.FloatTensor`):
+                The original samples.
+            noise (`torch.FloatTensor`):
+                The noise to add.
+            timesteps (`torch.IntTensor`):
+                The timesteps for which to add noise.
+
+        Returns:
+            `torch.FloatTensor`:
+                The noisy samples.
+        """
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+    
+    def get_velocity(
+        self,
+        sample: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        """
+        Get the velocity (score) for the given sample, noise, and timesteps.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The sample.
+            noise (`torch.FloatTensor`):
+                The noise.
+            timesteps (`torch.IntTensor`):
+                The timesteps.
+
+        Returns:
+            `torch.FloatTensor`:
+                The velocity.
+        """
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(sample.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+        return velocity
+    
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        if self.config.time_shift_type == "exponential":
+            return self._time_shift_exponential(mu, sigma, t)
+        elif self.config.time_shift_type == "linear":
+            return self._time_shift_linear(mu, sigma, t)
+
+    def _time_shift_exponential(self, mu, sigma, t):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def _time_shift_linear(self, mu, sigma, t):
+        return mu / (mu + (1 / t - 1) ** sigma)
+    
+    def taylor_approximation(self, taylor_approx_order, diff, model_output, derivative, second_derivative, third_derivative=None):
+        if taylor_approx_order == 1:
+            approx_value = model_output + diff * derivative
+        elif taylor_approx_order == 2:
+            if third_derivative is not None:
+                raise ValueError("The third derivative is computed but not used!")
+            approx_value = model_output + diff * derivative + 0.5 * diff**2 * second_derivative
+        elif taylor_approx_order == 3:
+            if third_derivative is None:
+                raise ValueError("The third derivative is not computed!")
+            approx_value = model_output + diff * derivative + 0.5 * diff**2 * second_derivative \
+                + diff**3 * third_derivative / 6
+        else:
+            print("The noise approximation order is not supported!")
+            exit()
+
+        return approx_value
+    
+
+    def drift_function(self, betas, total_step, t_eval, y_eval, noise):
+        '''
+        Drift function for the probability flow ODE in the noise-based diffusion model.
+
+        Args:
+            betas (`torch.FloatTensor`):
+                The betas of the diffusion model.
+            total_step (`int`):
+                The total number of steps in the diffusion chain.
+            t_eval (`torch.FloatTensor`):
+                The timestep to be evaluated at in the diffusion chain.
+            y_eval (`torch.FloatTensor`):
+                The sample to be evaluated at in the diffusion chain.
+            noise (`torch.FloatTensor`):
+                The noise used at the current timestep in the diffusion chain.
+
+        Returns:
+            `torch.FloatTensor`:
+                The drift term for the probability flow ODE in the diffusion model.
+        '''
+        beta_0, beta_1 = betas[0], betas[-1]
+        beta_t = (beta_0 + t_eval * (beta_1 - beta_0)) * total_step
+        beta_t = beta_t * torch.ones(y_eval.shape, device=y_eval.device)
+
+        log_mean_coeff = (-0.25 * t_eval ** 2 * (beta_1 - beta_0) - 0.5 * t_eval * beta_0) * total_step
+        std = torch.sqrt(1. - torch.exp(2. * log_mean_coeff))
+
+        # drift, diffusion -> f(x,t), g(t)
+        drift, diffusion = -0.5 * beta_t * y_eval, torch.sqrt(beta_t) * torch.ones(y_eval.shape, device=y_eval.device)
+        score = -noise / std  # score -> noise
+        drift = drift - diffusion ** 2 * score * 0.5 # drift -> dx/dt
+
+        return drift
+
+    def b_coeff(self, j):
+        '''
+        Coefficients of STORK2. The are based on the second order Runge-Kutta-Gegenbauer method.
+        Details of the coefficients can be found in https://www.sciencedirect.com/science/article/pii/S0021999120306537
+
+        Args:
+            j (`int`):
+                The sub-step index of the coefficient.
+
+        Returns:
+            `float`:
+                The coefficient of the STORK2.
+        '''
+        if j < 0:
+            print("The b_j coefficient in the RKG method can't have j negative")
+            return
+        if j == 0:
+            return 1
+        if j == 1:
+            return 1 / 3
+        
+        return 4 * (j - 1) * (j + 4) / (3 * j * (j + 1) * (j + 2) * (j + 3))
+
+    def coeff_rock1(self, j):
+        if j < 0:
+            print("The b_j coefficient in the RKG method can't have j negative")
+        return 2 / ((j + 1) * (j + 2))
+
+    def coeff_rock4(self):
+        '''
+        Load pre-computed coefficients of STORK4. The are based on the fourth order orthogonal Runge-Kutta-Chebyshev (ROCK4) method.
+        Details of the coefficients can be found in https://epubs.siam.org/doi/abs/10.1137/S1064827500379549.
+        The pre-computed coefficients are based on the implementation https://www.mathworks.com/matlabcentral/fileexchange/12129-rock4.
+
+        Args:
+            j (`int`):
+                The sub-step index of the coefficient.
+
+        Returns:
+            ms (`torch.FloatTensor`):
+                The degrees that coefficients were pre-computed for STORK4.
+            fpa, fpb, fpbe, recf (`torch.FloatTensor`):
+                The parameters for the finishing procedure.
+        '''
+        # Degrees
+        data = loadmat(f'{CONSTANTSFOLDER}/ms.mat')
+        ms = data['ms'][0]
+
+        # Parameters for the finishing procedure
+        data = loadmat(f'{CONSTANTSFOLDER}/fpa.mat')
+        fpa = data['fpa']
+
+        data = loadmat(f'{CONSTANTSFOLDER}/fpb.mat')
+        fpb = data['fpb']
+
+        data = loadmat(f'{CONSTANTSFOLDER}/fpbe.mat')
+        fpbe = data['fpbe']
+
+        # Parameters for the recurrence procedure
+        data = loadmat(f'{CONSTANTSFOLDER}/recf.mat')
+        recf = data['recf'][0]
+
+        return ms, fpa, fpb, fpbe, recf
+
+
+
+    def mdegr(self, mdeg1, ms):
+        '''
+        Find the optimal degree in the pre-computed degree coefficients table for the STORK4 method.
+
+        Args:
+            mdeg1 (`int`):
+                The degree to be evaluated.
+            ms (`torch.FloatTensor`):
+                The degrees that coefficients were pre-computed for STORK4.
+
+        Returns:
+            mdeg (`int`):
+                The optimal degree in the pre-computed degree coefficients table for the STORK4 method.
+            mp (`torch.FloatTensor`):
+                The pointer which select the degree in ms[i], such that mdeg<=ms[i].
+                mp[0] (`int`): The pointer which select the degree in ms[i], such that mdeg<=ms[i].
+                mp[1] (`int`): The pointer which gives the corresponding position of a_1 in the data recf for the selected degree.
+        '''           
+        mp = torch.zeros(2)
+        mp[1] = 1
+        mdeg = mdeg1
+        for i in range(len(ms)):
+            if (ms[i]/mdeg) >= 1:
+                mdeg = ms[i]
+                mp[0] = i
+                mp[1] = mp[1] - 1
+                break
+            else:   
+                mp[1] = mp[1] + ms[i] * 2 - 1
+
+        return mdeg, mp

model_index.json

@@ -0,0 +1,114 @@
+{
+  "_class_name": "BoomerPipeline",
+  "_diffusers_version": "0.33.0",
+  "_boomer_version": "1.0.0",
+  "transformer": {
+    "type": "BoomerFLADiT",
+    "config": "transformer/config.json",
+    "weights": "transformer/diffusion_pytorch_model.safetensors"
+  },
+  "vae": {
+    "repo_id": "mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers",
+    "scaling_factor": 0.41407
+  },
+  "text_encoder": {
+    "repo_id": "google/gemma-4-E2B-it",
+    "max_length": 300
+  },
+  "scheduler": "scheduler/scheduler_config.json",
+  "latent_normalization": {
+    "mean": [
+      -0.0492647976627326,
+      0.1468768895561574,
+      0.04348348892886716,
+      0.024480677381432905,
+      -0.36225658200088656,
+      -0.20211585707241028,
+      0.14117315317920875,
+      -0.10090931608978718,
+      0.03270434805644763,
+      0.025280784072185633,
+      0.42561436599740116,
+      0.07644308823255123,
+      -0.12361726209519984,
+      0.2738135117924219,
+      0.027258959344701725,
+      -0.15685215533194713,
+      0.1988568681778619,
+      0.07728443773748003,
+      0.21031734156716703,
+      0.10236920059442055,
+      -0.26953907125577387,
+      -0.1039037490181443,
+      -0.14040348514520445,
+      -0.050237464451169944,
+      0.21928026529320632,
+      0.05541749411464261,
+      0.15868418162302406,
+      0.09498460353222035,
+      0.07154154194705771,
+      -0.0980392861411281,
+      0.3445218162967998,
+      0.14452621160838316
+    ],
+    "std": [
+      0.7239755227946175,
+      0.7084603356016493,
+      0.7371127244335353,
+      0.7148677155404667,
+      0.7610612675902568,
+      0.7831300251777134,
+      1.241222644947736,
+      1.1914623118386434,
+      0.7064647426283694,
+      1.0233179582088132,
+      0.7671679694251226,
+      0.6818639786525276,
+      0.7394871026815577,
+      0.6749445490371844,
+      0.7961588737844489,
+      0.7955142324161893,
+      0.7545916153429181,
+      0.7799818111961734,
+      0.706798939521899,
+      0.7014546090493033,
+      0.9678039884252744,
+      0.7504288798344418,
+      0.7296232257036755,
+      0.7257654508983634,
+      2.2632974219950786,
+      0.8916002210501063,
+      0.8534945911823539,
+      0.7403039593986197,
+      0.7264856936752643,
+      0.6879722344092367,
+      0.7331094494058464,
+      0.6896992616885751
+    ]
+  },
+  "training_info": {
+    "dataset": "FineT2IcacheBF16_1024",
+    "base_model": "journeydb-pretrained-boomer-fla",
+    "image_size_px": 1024,
+    "patch_size": 1,
+    "latent_size": 32,
+    "latent_tokens": 1024,
+    "steps_finetune": 55000,
+    "steps_configured": 75000,
+    "batch_size": 24,
+    "lr": 0.0001,
+    "lr_scheduler": "linear-warmup-cosine",
+    "warmup_steps": 1000,
+    "min_lr_ratio": 0.1,
+    "flow_shift": 1.5,
+    "t_sampler": "plateau-logit-normal",
+    "logit_mean": 0.0,
+    "logit_std": 1.0,
+    "min_snr_gamma": 5.0,
+    "cond_dropout": 0.1,
+    "ema_decay": 0.999,
+    "ema_update_every": 8,
+    "grad_clip": 0.3,
+    "latent_stats_mode": "channel"
+  }
+}

modeling_boomer_fla.py

@@ -0,0 +1,1235 @@
+"""BoomerFLADiT model — self-contained for HuggingFace trust_remote_code distribution.
+
+All dependencies inlined: no boomer package import needed.
+External pip requirements: torch, flash-linear-attention (fla).
+"""
+# ── inlined from boomer/models/latent_dit.py ──────────────────────────────────
+from __future__ import annotations
+import math
+import sys
+import types
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint as _ckpt
+
+
+class AttentionRMSNorm(nn.Module):
+    def __init__(self, dim: int, scale_factor: float = 0.01, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim) * scale_factor)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        normed = x.float() * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        weight = self.weight.view(*([1] * (x.ndim - 2)), -1)
+        return (weight * normed).type_as(x)
+
+
+class CaptionEmbedder(nn.Module):
+    def __init__(self, in_channels: int, hidden_size: int, token_num: int) -> None:
+        super().__init__()
+        self.y_proj = nn.Sequential(
+            nn.Linear(in_channels, hidden_size),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(hidden_size, hidden_size),
+        )
+        null_init = torch.randn(token_num, in_channels) / math.sqrt(in_channels)
+        self.null_text_embedding = nn.Parameter(null_init.unsqueeze(0))
+
+    def forward(self, caption: torch.Tensor) -> torch.Tensor:
+        return self.y_proj(caption)
+
+    def null_condition(self, batch_size, *, device, dtype, mask_dtype=None, token_num=None):
+        text = self.null_text_embedding
+        if token_num is not None and token_num != text.shape[1]:
+            if token_num < text.shape[1]:
+                text = text[:, :token_num]
+            else:
+                pad = text.new_zeros(text.shape[0], token_num - text.shape[1], text.shape[2])
+                text = torch.cat([text, pad], dim=1)
+        text = text.expand(batch_size, -1, -1).to(device=device, dtype=dtype)
+        mask = torch.ones(batch_size, text.shape[1], device=device, dtype=mask_dtype or torch.long)
+        if token_num is not None and token_num > self.null_text_embedding.shape[1]:
+            mask[:, self.null_text_embedding.shape[1]:] = 0
+        return text, mask
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_dim: int) -> None:
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(1, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim))
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        dtype = self.net[0].weight.dtype
+        return self.net(timesteps.to(dtype=dtype).view(-1, 1))
+
+
+# ── rest of boomer_fla_dit.py below (unchanged except no boomer imports) ──────
+
+
+
+@dataclass(frozen=True)
+class BoomerFLADiTConfig:
+    model_type: str = "boomer_fla"
+    latent_channels: int = 32
+    latent_size: int = 16
+    text_dim: int = 1536
+    text_seq_len: int = 300
+    hidden_dim: int = 1152
+    depth: int = 28
+    num_heads: int = 16
+    mlp_ratio: float = 2.5
+    y_norm: bool = True
+    y_norm_scale_factor: float = 0.01
+    mixer_type: str = "fla_linear"
+    fla_mode: str = "chunk"
+    fla_feature_map: str = "relu"
+    fla_bidirectional: bool = False
+    use_short_conv: bool = False
+    conv_size: int = 4
+    image_attention_every: int = 0
+    image_attention_backend: str = "sdpa"
+    image_attention_rope: bool = False
+    image_rope_theta: float = 10000.0
+    cross_attention_backend: str = "sdpa"
+    cross_attention_qk_norm: bool = True
+    parallel_block: bool = False
+    dual_stream_depth: int = 0
+    multimodal_coord_ids: bool = False
+    use_abs_pos_embed: bool = True
+    patch_size: int = 1
+    gradient_checkpointing: bool = False
+
+
+def maybe_add_sibling_fla_repo() -> None:
+    candidates = [
+        Path(__file__).resolve().parents[3] / "flash-linear-attention",
+        Path("/content/flash-linear-attention"),
+        Path("/content/flame"),
+    ]
+    for path in candidates:
+        if (path / "fla").is_dir() and str(path) not in sys.path:
+            sys.path.insert(0, str(path))
+
+
+def maybe_add_sibling_flash_attention_repo() -> None:
+    candidates = [
+        Path(__file__).resolve().parents[3] / "flash-attention" / "hopper",
+        Path(__file__).resolve().parents[3] / "flash-attention",
+        Path("/work/flash-attention/hopper"),
+        Path("/work/flash-attention"),
+        Path("/home/jovyan/work/flash-attention"),
+        Path("/content/flash-attention/hopper"),
+        Path("/content/flash-attention"),
+    ]
+    for path in candidates:
+        if path.exists() and str(path) not in sys.path:
+            sys.path.insert(0, str(path))
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1.0 + scale) + shift
+
+
+class ConvLayer(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        kernel_size: int,
+        *,
+        groups: int = 1,
+        bias: bool = False,
+        act: str | None = None,
+    ) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_dim,
+            out_dim,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+            groups=groups,
+            bias=bias,
+        )
+        self.act = nn.SiLU() if act == "silu" else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.conv(x))
+
+
+class GLUMBConv(nn.Module):
+    """Sana GLUMBConv FFN: 1x1 expand, depthwise spatial conv, GLU, 1x1 project."""
+
+    def __init__(self, hidden_dim: int, mlp_ratio: float) -> None:
+        super().__init__()
+        inner_dim = int(hidden_dim * mlp_ratio)
+        self.inner_dim = inner_dim
+        self.inverted_conv = ConvLayer(hidden_dim, inner_dim * 2, 1, bias=True, act="silu")
+        self.depth_conv = ConvLayer(inner_dim * 2, inner_dim * 2, 3, groups=inner_dim * 2, bias=True)
+        self.point_conv = ConvLayer(inner_dim, hidden_dim, 1, bias=False)
+        nn.init.zeros_(self.point_conv.conv.weight)
+        self.glu_act = nn.SiLU()
+
+    def forward(self, x: torch.Tensor, *, height: int, width: int) -> torch.Tensor:
+        batch, tokens, channels = x.shape
+        if tokens != height * width:
+            raise ValueError(f"Expected {height * width} image tokens, got {tokens}")
+        x = x.reshape(batch, height, width, channels).permute(0, 3, 1, 2).contiguous()
+        x = self.inverted_conv(x)
+        x = self.depth_conv(x)
+        x, gate = x.chunk(2, dim=1)
+        x = x * self.glu_act(gate)
+        x = self.point_conv(x)
+        return x.reshape(batch, channels, tokens).transpose(1, 2).contiguous()
+
+
+class TorchSelfAttention(nn.Module):
+    def __init__(self, hidden_dim: int, num_heads: int) -> None:
+        super().__init__()
+        self.attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.attn(x, x, x, need_weights=False)[0]
+
+
+class TokenMLP(nn.Module):
+    def __init__(self, hidden_dim: int, mlp_ratio: float) -> None:
+        super().__init__()
+        inner_dim = int(hidden_dim * mlp_ratio)
+        self.net = nn.Sequential(
+            nn.Linear(hidden_dim, inner_dim),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(inner_dim, hidden_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class MultimodalCoordinateRoPE(nn.Module):
+    """FLUX-style coordinate-ID RoPE for joint text/image attention."""
+
+    def __init__(self, head_dim: int, *, image_size: int, text_seq_len: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        if head_dim < 6 or head_dim % 2 != 0:
+            raise ValueError(f"head_dim={head_dim} must be even and at least 6 for multimodal RoPE")
+        if theta <= 0.0:
+            raise ValueError(f"theta must be positive, got {theta}")
+        type_dim = max(2, (head_dim // 4) // 2 * 2)
+        while type_dim > 2 and (head_dim - type_dim) % 4 != 0:
+            type_dim -= 2
+        remaining = head_dim - type_dim
+        row_dim = max(2, (remaining // 2) // 2 * 2)
+        col_dim = remaining - row_dim
+        if col_dim < 2 or col_dim % 2 != 0:
+            raise ValueError(f"could not split head_dim={head_dim} into even multimodal RoPE axes")
+        self.axes_dim = (type_dim, row_dim, col_dim)
+        self.head_dim = head_dim
+        self.image_size = image_size
+        self.text_seq_len = text_seq_len
+
+        for index, dim in enumerate(self.axes_dim):
+            inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+            self.register_buffer(f"inv_freq_{index}", inv_freq, persistent=False)
+
+    @staticmethod
+    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+    def image_ids(self, batch_size: int, *, height: int, width: int, device: torch.device | str) -> torch.Tensor:
+        token_idx = torch.arange(height * width, device=device)
+        rows = token_idx // width
+        cols = token_idx % width
+        token_type = torch.ones_like(rows)
+        ids = torch.stack([token_type, rows, cols], dim=-1)
+        return ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+    def text_ids(self, batch_size: int, token_count: int, *, device: torch.device | str) -> torch.Tensor:
+        positions = torch.arange(token_count, device=device)
+        token_type = torch.zeros_like(positions)
+        zeros = torch.zeros_like(positions)
+        ids = torch.stack([token_type, positions, zeros], dim=-1)
+        return ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+    def _axis_apply(self, x: torch.Tensor, axis_ids: torch.Tensor, axis_index: int) -> torch.Tensor:
+        inv_freq = getattr(self, f"inv_freq_{axis_index}")
+        angles = axis_ids.float().unsqueeze(-1) * inv_freq.to(device=x.device).view(1, 1, -1)
+        cos = torch.cat([angles.cos(), angles.cos()], dim=-1).unsqueeze(2).to(dtype=x.dtype)
+        sin = torch.cat([angles.sin(), angles.sin()], dim=-1).unsqueeze(2).to(dtype=x.dtype)
+        return x * cos + self._rotate_half(x) * sin
+
+    def apply(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if q.shape[-1] != self.head_dim or k.shape[-1] != self.head_dim:
+            raise ValueError(f"expected head_dim={self.head_dim}, got q={q.shape[-1]} k={k.shape[-1]}")
+        if ids.shape[:2] != q.shape[:2] or ids.shape[-1] != len(self.axes_dim):
+            raise ValueError(f"expected ids shape (B, T, {len(self.axes_dim)}), got {tuple(ids.shape)}")
+        q_chunks = q.split(self.axes_dim, dim=-1)
+        k_chunks = k.split(self.axes_dim, dim=-1)
+        q_out = []
+        k_out = []
+        for index, (q_axis, k_axis) in enumerate(zip(q_chunks, k_chunks, strict=True)):
+            q_out.append(self._axis_apply(q_axis, ids[..., index], index))
+            k_out.append(self._axis_apply(k_axis, ids[..., index], index))
+        return torch.cat(q_out, dim=-1), torch.cat(k_out, dim=-1)
+
+
+class RoPE2D(nn.Module):
+    """2D RoPE for image tokens on a fixed H×W grid (row-major flattening).
+
+    Splits head_dim in half: the first half encodes height, the second width.
+    Each half uses standard 1D RoPE with shared cos/sin tables per axis.
+    """
+
+    def __init__(self, head_dim: int, grid_size: int, *, theta: float = 10000.0) -> None:
+        super().__init__()
+        if head_dim % 4 != 0:
+            raise ValueError(
+                f"head_dim={head_dim} must be divisible by 4 for 2D RoPE "
+                f"(half for H, half for W, each needing pairs)"
+            )
+        if grid_size <= 0:
+            raise ValueError(f"grid_size must be positive, got {grid_size}")
+        if theta <= 0.0:
+            raise ValueError(f"theta must be positive, got {theta}")
+        self.head_dim = head_dim
+        self.grid_size = grid_size
+        self.half_dim = head_dim // 2
+
+        freqs = 1.0 / (theta ** (torch.arange(0, self.half_dim, 2).float() / self.half_dim))
+        token_idx = torch.arange(grid_size * grid_size)
+        h_idx = token_idx // grid_size
+        w_idx = token_idx % grid_size
+
+        def axis_tables(pos_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            angles = torch.outer(pos_idx.float(), freqs)
+            cos = torch.cat([angles.cos(), angles.cos()], dim=-1)[None, :, None, :]
+            sin = torch.cat([angles.sin(), angles.sin()], dim=-1)[None, :, None, :]
+            return cos, sin
+
+        cos_h, sin_h = axis_tables(h_idx)
+        cos_w, sin_w = axis_tables(w_idx)
+        self.register_buffer("cos_h", cos_h, persistent=False)
+        self.register_buffer("sin_h", sin_h, persistent=False)
+        self.register_buffer("cos_w", cos_w, persistent=False)
+        self.register_buffer("sin_w", sin_w, persistent=False)
+
+    @staticmethod
+    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+    def _apply_axis_rope(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        return x * cos.to(dtype=x.dtype) + self._rotate_half(x) * sin.to(dtype=x.dtype)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch, tokens, num_heads, head_dim = q.shape
+        if head_dim != self.head_dim:
+            raise ValueError(f"expected head_dim={self.head_dim}, got {head_dim}")
+        expected_tokens = self.grid_size * self.grid_size
+        if tokens != expected_tokens:
+            raise ValueError(f"expected {expected_tokens} image tokens, got {tokens}")
+
+        q_h, q_w = q.chunk(2, dim=-1)
+        k_h, k_w = k.chunk(2, dim=-1)
+        q_h = self._apply_axis_rope(q_h, self.cos_h, self.sin_h)
+        q_w = self._apply_axis_rope(q_w, self.cos_w, self.sin_w)
+        k_h = self._apply_axis_rope(k_h, self.cos_h, self.sin_h)
+        k_w = self._apply_axis_rope(k_w, self.cos_w, self.sin_w)
+        return torch.cat([q_h, q_w], dim=-1), torch.cat([k_h, k_w], dim=-1)
+
+
+class FullImageSelfAttention(nn.Module):
+    """Full image-token attention for the small DC-AE latent grid."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int,
+        *,
+        backend: str = "sdpa",
+        grid_size: int | None = None,
+        rope: bool = False,
+        rope_theta: float = 10000.0,
+    ) -> None:
+        super().__init__()
+        if hidden_dim % num_heads != 0:
+            raise ValueError(f"hidden_dim={hidden_dim} must be divisible by num_heads={num_heads}")
+        if backend not in {"sdpa", "flash3", "flash4", "auto"}:
+            raise ValueError(f"Unsupported image_attention_backend: {backend}")
+        if rope and grid_size is None:
+            raise ValueError("grid_size is required when rope=True")
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+        self.backend = backend
+        self.qkv = nn.Linear(hidden_dim, hidden_dim * 3)
+        self.out_proj = nn.Linear(hidden_dim, hidden_dim)
+        nn.init.zeros_(self.out_proj.weight)
+        nn.init.zeros_(self.out_proj.bias)
+        self.rope = (
+            RoPE2D(self.head_dim, grid_size, theta=rope_theta)
+            if rope and grid_size is not None
+            else None
+        )
+        self._flash3_attn_func = None
+        self._flash3_import_attempted = False
+        self._flash4_attn_func = None
+        self._flash4_import_attempted = False
+
+    def _get_flash3_attn_func(self):
+        if self._flash3_import_attempted:
+            return self._flash3_attn_func
+        self._flash3_import_attempted = True
+        maybe_add_sibling_flash_attention_repo()
+        try:
+            from flash_attn_interface import flash_attn_func
+        except Exception:
+            try:
+                from flash_attn.flash_attn_interface import flash_attn_func
+            except Exception:
+                flash_attn_func = None
+        self._flash3_attn_func = flash_attn_func
+        return self._flash3_attn_func
+
+    def _get_flash4_attn_func(self):
+        if self._flash4_import_attempted:
+            return self._flash4_attn_func
+        self._flash4_import_attempted = True
+        maybe_add_sibling_flash_attention_repo()
+        try:
+            from flash_attn.cute.interface import flash_attn_func
+        except Exception:
+            flash4_paths = [
+                Path(__file__).resolve().parents[3] / "flash-attention" / "flash_attn",
+                Path("/work/flash-attention/flash_attn"),
+                Path("/home/jovyan/work/flash-attention/flash_attn"),
+                Path("/content/flash-attention/flash_attn"),
+            ]
+            existing_paths = [str(path) for path in flash4_paths if (path / "cute").is_dir()]
+            if existing_paths:
+                for name in list(sys.modules):
+                    if name == "flash_attn" or name.startswith("flash_attn."):
+                        del sys.modules[name]
+                flash_attn_pkg = types.ModuleType("flash_attn")
+                flash_attn_pkg.__path__ = existing_paths
+                sys.modules["flash_attn"] = flash_attn_pkg
+                try:
+                    from flash_attn.cute.interface import flash_attn_func
+                except Exception:
+                    flash_attn_func = None
+            else:
+                flash_attn_func = None
+        self._flash4_attn_func = flash_attn_func
+        return self._flash4_attn_func
+
+    def _flash3_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        flash_attn_func = self._get_flash3_attn_func()
+        if flash_attn_func is None:
+            raise ImportError(
+                "image_attention_backend='flash3' requires FlashAttention-3. "
+                "Install it or use --image-attn-backend sdpa."
+            )
+        out = flash_attn_func(q, k, v, causal=False)
+        if isinstance(out, tuple):
+            out = out[0]
+        return out
+
+    def _flash4_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        flash_attn_func = self._get_flash4_attn_func()
+        if flash_attn_func is None:
+            raise ImportError(
+                "image_attention_backend='flash4' requires FlashAttention-4/CuTe. "
+                "Install flash-attn-4 or use --image-attn-backend sdpa."
+            )
+        out = flash_attn_func(q, k, v, causal=False)
+        if isinstance(out, tuple):
+            out = out[0]
+        return out
+
+    @staticmethod
+    def _flash_compute_dtype(x: torch.Tensor) -> torch.dtype | None:
+        """FA kernels need fp16/bf16; fp32 master weights + compile may still pass fp32 activations."""
+        if not x.is_cuda:
+            return None
+        if x.dtype in {torch.float16, torch.bfloat16}:
+            return x.dtype
+        if torch.is_autocast_enabled():
+            return torch.get_autocast_dtype("cuda")
+        return torch.bfloat16
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, tokens, channels = x.shape
+        qkv = self.qkv(x).reshape(batch, tokens, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        if self.rope is not None:
+            q, k = self.rope(q, k)
+
+        flash_dtype = self._flash_compute_dtype(x)
+        use_flash = self.backend in {"flash3", "flash4", "auto"} and flash_dtype is not None
+        if use_flash and (q.dtype != flash_dtype or k.dtype != flash_dtype or v.dtype != flash_dtype):
+            q, k, v = q.to(flash_dtype), k.to(flash_dtype), v.to(flash_dtype)
+
+        if self.backend == "flash4" and use_flash:
+            out = self._flash4_attention(q, k, v)
+        elif self.backend == "flash3" and use_flash:
+            out = self._flash3_attention(q, k, v)
+        elif self.backend == "auto" and use_flash:
+            try:
+                out = self._flash4_attention(q, k, v)
+            except Exception:
+                try:
+                    out = self._flash3_attention(q, k, v)
+                except Exception:
+                    use_flash = False
+        if self.backend in {"flash3", "flash4"} and not use_flash:
+            raise RuntimeError(
+                f"image_attention_backend='{self.backend}' requires CUDA fp16/bf16 compute; got {x.device} {x.dtype}"
+            )
+        if use_flash and out.dtype != x.dtype:
+            out = out.to(dtype=x.dtype)
+        if not use_flash:
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
+            out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+            out = out.transpose(1, 2)
+
+        out = out.reshape(batch, tokens, channels)
+        return self.out_proj(out)
+
+
+class SanaMultiHeadCrossAttention(nn.Module):
+    """Sana-style cross-attention with optional q/k norm and SDPA/xformers kernels."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int,
+        *,
+        backend: str = "sdpa",
+        qk_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        if hidden_dim % num_heads != 0:
+            raise ValueError(f"hidden_dim={hidden_dim} must be divisible by num_heads={num_heads}")
+        if backend not in {"sdpa", "xformers", "auto"}:
+            raise ValueError(f"Unsupported cross_attention_backend: {backend}")
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+        self.backend = backend
+        self.q_linear = nn.Linear(hidden_dim, hidden_dim)
+        self.kv_linear = nn.Linear(hidden_dim, hidden_dim * 2)
+        self.q_norm = AttentionRMSNorm(hidden_dim, scale_factor=1.0, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = AttentionRMSNorm(hidden_dim, scale_factor=1.0, eps=1e-6) if qk_norm else nn.Identity()
+        self.proj = nn.Linear(hidden_dim, hidden_dim)
+        # adaLN-Zero style: cross-attn starts as a no-op so Gemma text cannot spike GDN states early.
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+        self._xformers_ops = None
+        self._xformers_import_attempted = False
+
+    def _get_xformers_ops(self):
+        if self._xformers_import_attempted:
+            return self._xformers_ops
+        self._xformers_import_attempted = True
+        try:
+            import xformers.ops as xops
+        except Exception:
+            xops = None
+        self._xformers_ops = xops
+        return self._xformers_ops
+
+    def _xformers_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        key_padding_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        xops = self._get_xformers_ops()
+        if xops is None:
+            raise ImportError(
+                "cross_attention_backend='xformers' requires xformers. "
+                "Install it or use --cross-attn-backend sdpa."
+            )
+
+        batch, image_tokens = q.shape[:2]
+        text_tokens = k.shape[1]
+        q_lens = [image_tokens] * batch
+        q_compact = q.reshape(1, batch * image_tokens, self.num_heads, self.head_dim)
+        if key_padding_mask is None:
+            kv_lens = [text_tokens] * batch
+            k_compact = k.reshape(1, batch * text_tokens, self.num_heads, self.head_dim)
+            v_compact = v.reshape(1, batch * text_tokens, self.num_heads, self.head_dim)
+        else:
+            valid_mask = ~key_padding_mask.bool()
+            kv_lens = valid_mask.sum(dim=1).tolist()
+            if any(length <= 0 for length in kv_lens):
+                raise ValueError("xformers cross-attention received a sample with zero valid text tokens")
+            k_compact = torch.cat([k[index, valid_mask[index]] for index in range(batch)], dim=0).unsqueeze(0)
+            v_compact = torch.cat([v[index, valid_mask[index]] for index in range(batch)], dim=0).unsqueeze(0)
+
+        attn_bias = xops.fmha.BlockDiagonalMask.from_seqlens(q_lens, kv_lens)
+        out = xops.memory_efficient_attention(q_compact, k_compact, v_compact, attn_bias=attn_bias, p=0.0)
+        return out.reshape(batch, image_tokens, self.num_heads, self.head_dim)
+
+    def _sdpa_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        key_padding_mask: torch.Tensor | None,
+        attn_bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        attn_mask = attn_bias
+        if attn_mask is None and key_padding_mask is not None:
+            attn_mask = key_padding_mask[:, None, None, :].to(dtype=q.dtype)
+            attn_mask = attn_mask.masked_fill(attn_mask > 0, -10000.0)
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        return out.transpose(1, 2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        key_padding_mask: torch.Tensor | None = None,
+        attn_bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch, image_tokens, channels = x.shape
+        # Sana order: linear projection first, then per-token q/k RMSNorm before head split.
+        # This caps dot-product growth when cond carries high-magnitude Gemma caption states.
+        q = self.q_linear(x)
+        q = self.q_norm(q).reshape(batch, image_tokens, self.num_heads, self.head_dim)
+        k, v = self.kv_linear(cond).chunk(2, dim=-1)
+        k = self.k_norm(k).reshape(batch, cond.shape[1], self.num_heads, self.head_dim)
+        v = v.reshape(batch, cond.shape[1], self.num_heads, self.head_dim)
+
+        use_xformers = self.backend in {"xformers", "auto"} and x.is_cuda and x.dtype in {
+            torch.float16,
+            torch.bfloat16,
+        }
+        if use_xformers:
+            try:
+                out = self._xformers_attention(q, k, v, key_padding_mask)
+            except Exception:
+                if self.backend == "xformers":
+                    raise
+                use_xformers = False
+        if self.backend == "xformers" and not use_xformers:
+            raise RuntimeError(
+                f"cross_attention_backend='xformers' requires CUDA fp16/bf16 tensors; got {x.device} {x.dtype}"
+            )
+        if not use_xformers:
+            out = self._sdpa_attention(q, k, v, key_padding_mask, attn_bias)
+
+        return self.proj(out.reshape(batch, image_tokens, channels))
+
+
+class FLASelfMixer(nn.Module):
+    def __init__(self, config: BoomerFLADiTConfig, *, layer_idx: int) -> None:
+        super().__init__()
+        try:
+            import fla.layers as fla_layers
+        except Exception:
+            maybe_add_sibling_fla_repo()
+            import fla.layers as fla_layers
+
+        hidden_dim = config.hidden_dim
+        self.bidirectional = config.fla_bidirectional
+
+        def make_mixer() -> nn.Module:
+            if config.mixer_type == "fla_linear":
+                return fla_layers.LinearAttention(
+                    hidden_size=hidden_dim,
+                    num_heads=config.num_heads,
+                    mode=config.fla_mode,
+                    feature_map=config.fla_feature_map,
+                    output_norm="rmsnorm",
+                    layer_idx=layer_idx,
+                )
+            if config.mixer_type == "fla_gated_deltanet":
+                return fla_layers.GatedDeltaNet(
+                    hidden_size=hidden_dim,
+                    num_heads=config.num_heads,
+                    head_dim=hidden_dim // config.num_heads,
+                    expand_v=1,
+                    mode=config.fla_mode,
+                    use_short_conv=config.use_short_conv,
+                    conv_size=config.conv_size,
+                    layer_idx=layer_idx,
+                )
+            if config.mixer_type == "fla_gla":
+                return fla_layers.GatedLinearAttention(
+                    hidden_size=hidden_dim,
+                    num_heads=config.num_heads,
+                    mode=config.fla_mode,
+                    feature_map=config.fla_feature_map,
+                    use_short_conv=config.use_short_conv,
+                    conv_size=config.conv_size,
+                    layer_idx=layer_idx,
+                )
+            raise ValueError(f"Unsupported FLA mixer_type: {config.mixer_type}")
+
+        self.mixer_fwd = make_mixer()
+        self.mixer_bwd = make_mixer() if self.bidirectional else None
+        if self.bidirectional:
+            self.out_proj = nn.Linear(hidden_dim * 2, hidden_dim, bias=False)
+            nn.init.zeros_(self.out_proj.weight)
+
+    @staticmethod
+    def _run_mixer(mixer: nn.Module, x: torch.Tensor) -> torch.Tensor:
+        y = mixer(x)
+        if isinstance(y, tuple):
+            y = y[0]
+        return y
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self._run_mixer(self.mixer_fwd, x)
+        if not self.bidirectional:
+            return y
+        if self.mixer_bwd is None:
+            raise RuntimeError("bidirectional FLASelfMixer is missing the backward mixer")
+        y_rev = self._run_mixer(self.mixer_bwd, x.flip(1)).flip(1)
+        return self.out_proj(torch.cat([y, y_rev], dim=-1))
+
+
+class BoomerFLABlock(nn.Module):
+    def __init__(self, config: BoomerFLADiTConfig, *, layer_idx: int) -> None:
+        super().__init__()
+        hidden_dim = config.hidden_dim
+        self.parallel_block = config.parallel_block
+        self.use_image_attention = (
+            config.image_attention_every > 0 and (layer_idx + 1) % config.image_attention_every == 0
+        )
+        self.norm1 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        if config.mixer_type in {"torch", "fallback"}:
+            self.self_attn = TorchSelfAttention(hidden_dim, config.num_heads)
+        else:
+            self.self_attn = FLASelfMixer(config, layer_idx=layer_idx)
+        if self.use_image_attention:
+            self.image_attn_norm = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+            self.image_attn_mod = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 3))
+            self.image_attn = FullImageSelfAttention(
+                hidden_dim,
+                config.num_heads,
+                backend=config.image_attention_backend,
+                grid_size=config.latent_size // config.patch_size,
+                rope=config.image_attention_rope,
+                rope_theta=config.image_rope_theta,
+            )
+            self.image_attn_scale_shift_table = nn.Parameter(torch.zeros(3, hidden_dim))
+        cross_backend = config.cross_attention_backend
+        if config.cross_attention_qk_norm and cross_backend == "mha":
+            raise ValueError(
+                "cross_attention_qk_norm requires SanaMultiHeadCrossAttention "
+                "(cross_attention_backend sdpa/xformers/auto), not mha"
+            )
+        if cross_backend == "mha":
+            self.cross_attn = nn.MultiheadAttention(hidden_dim, config.num_heads, batch_first=True)
+        else:
+            self.cross_attn = SanaMultiHeadCrossAttention(
+                hidden_dim,
+                config.num_heads,
+                backend=cross_backend,
+                qk_norm=config.cross_attention_qk_norm,
+            )
+        self.mod = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 9))
+        self.norm2 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.norm3 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.mlp = GLUMBConv(hidden_dim, config.mlp_ratio)
+        self.scale_shift_table = nn.Parameter(torch.zeros(9, hidden_dim))
+
+    def _cross_attention(
+        self,
+        x: torch.Tensor,
+        text_tokens: torch.Tensor,
+        text_key_padding_mask: torch.Tensor,
+        text_attn_bias: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if isinstance(self.cross_attn, nn.MultiheadAttention):
+            return self.cross_attn(
+                x,
+                text_tokens,
+                text_tokens,
+                key_padding_mask=text_key_padding_mask,
+                need_weights=False,
+            )[0]
+        return self.cross_attn(x, text_tokens, text_key_padding_mask, text_attn_bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text_tokens: torch.Tensor,
+        t_emb: torch.Tensor,
+        text_key_padding_mask: torch.Tensor,
+        text_attn_bias: torch.Tensor | None,
+        *,
+        height: int,
+        width: int,
+    ) -> torch.Tensor:
+        timestep_mod = self.mod(t_emb)
+        (
+            shift_msa,
+            scale_msa,
+            gate_msa,
+            shift_cross,
+            scale_cross,
+            gate_cross,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+        ) = (self.scale_shift_table[None] + timestep_mod.reshape(x.shape[0], 9, -1)).chunk(9, dim=1)
+        if self.parallel_block:
+            base = x
+            branches = [
+                gate_msa * self.self_attn(modulate(self.norm1(base), shift_msa, scale_msa)),
+                gate_cross
+                * self._cross_attention(
+                    modulate(self.norm3(base), shift_cross, scale_cross),
+                    text_tokens,
+                    text_key_padding_mask,
+                    text_attn_bias,
+                ),
+                gate_mlp * self.mlp(modulate(self.norm2(base), shift_mlp, scale_mlp), height=height, width=width),
+            ]
+            if self.use_image_attention:
+                image_attn_mod = self.image_attn_mod(t_emb)
+                shift_img, scale_img, gate_img = (
+                    self.image_attn_scale_shift_table[None] + image_attn_mod.reshape(x.shape[0], 3, -1)
+                ).chunk(3, dim=1)
+                branches.append(
+                    gate_img * self.image_attn(modulate(self.image_attn_norm(base), shift_img, scale_img))
+                )
+            return base + sum(branches)
+
+        x = x + gate_msa * self.self_attn(modulate(self.norm1(x), shift_msa, scale_msa))
+        if self.use_image_attention:
+            image_attn_mod = self.image_attn_mod(t_emb)
+            shift_img, scale_img, gate_img = (
+                self.image_attn_scale_shift_table[None] + image_attn_mod.reshape(x.shape[0], 3, -1)
+            ).chunk(3, dim=1)
+            x = x + gate_img * self.image_attn(modulate(self.image_attn_norm(x), shift_img, scale_img))
+        x = x + gate_cross * self._cross_attention(
+            modulate(self.norm3(x), shift_cross, scale_cross),
+            text_tokens,
+            text_key_padding_mask,
+            text_attn_bias,
+        )
+        x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp), height=height, width=width)
+        return x
+
+
+class BoomerFLADualStreamBlock(nn.Module):
+    """FLUX-style early block with one joint text+image attention operation."""
+
+    updates_text = True
+
+    def __init__(self, config: BoomerFLADiTConfig, *, layer_idx: int) -> None:
+        super().__init__()
+        hidden_dim = config.hidden_dim
+        if hidden_dim % config.num_heads != 0:
+            raise ValueError(f"hidden_dim={hidden_dim} must be divisible by num_heads={config.num_heads}")
+        self.num_heads = config.num_heads
+        self.head_dim = hidden_dim // config.num_heads
+        self.hidden_dim = hidden_dim
+        self.parallel_block = config.parallel_block
+
+        self.image_mod = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6))
+        self.image_norm1 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.image_qkv = nn.Linear(hidden_dim, hidden_dim * 3)
+        self.image_q_norm = AttentionRMSNorm(self.head_dim, scale_factor=1.0, eps=1e-6)
+        self.image_k_norm = AttentionRMSNorm(self.head_dim, scale_factor=1.0, eps=1e-6)
+        self.image_out_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.image_norm2 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.image_mlp = GLUMBConv(hidden_dim, config.mlp_ratio)
+        self.image_scale_shift_table = nn.Parameter(torch.zeros(6, hidden_dim))
+
+        self.text_mod = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6))
+        self.text_norm1 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.text_qkv = nn.Linear(hidden_dim, hidden_dim * 3)
+        self.text_q_norm = AttentionRMSNorm(self.head_dim, scale_factor=1.0, eps=1e-6)
+        self.text_k_norm = AttentionRMSNorm(self.head_dim, scale_factor=1.0, eps=1e-6)
+        self.text_out_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.text_norm2 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.text_mlp = TokenMLP(hidden_dim, config.mlp_ratio)
+        self.text_scale_shift_table = nn.Parameter(torch.zeros(6, hidden_dim))
+
+    def _qkv(
+        self,
+        x: torch.Tensor,
+        qkv: nn.Linear,
+        q_norm: AttentionRMSNorm,
+        k_norm: AttentionRMSNorm,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch, tokens, _ = x.shape
+        q, k, v = qkv(x).reshape(batch, tokens, 3, self.num_heads, self.head_dim).unbind(dim=2)
+        q = q_norm(q)
+        k = k_norm(k)
+        return q, k, v
+
+    def _joint_attention(
+        self,
+        image_tokens: torch.Tensor,
+        text_tokens: torch.Tensor,
+        text_key_padding_mask: torch.Tensor,
+        coord_rope: MultimodalCoordinateRoPE | None,
+        image_coord_ids: torch.Tensor | None,
+        text_coord_ids: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        image_q, image_k, image_v = self._qkv(
+            image_tokens,
+            self.image_qkv,
+            self.image_q_norm,
+            self.image_k_norm,
+        )
+        text_q, text_k, text_v = self._qkv(
+            text_tokens,
+            self.text_qkv,
+            self.text_q_norm,
+            self.text_k_norm,
+        )
+        q = torch.cat([text_q, image_q], dim=1)
+        k = torch.cat([text_k, image_k], dim=1)
+        v = torch.cat([text_v, image_v], dim=1)
+        if coord_rope is not None:
+            if image_coord_ids is None or text_coord_ids is None:
+                raise ValueError("coordinate ids are required when multimodal coord RoPE is enabled")
+            coord_ids = torch.cat([text_coord_ids, image_coord_ids], dim=1)
+            q, k = coord_rope.apply(q, k, coord_ids)
+
+        image_mask = torch.zeros(
+            image_tokens.shape[0],
+            image_tokens.shape[1],
+            device=image_tokens.device,
+            dtype=text_key_padding_mask.dtype,
+        )
+        key_padding_mask = torch.cat([text_key_padding_mask, image_mask], dim=1)
+        attn_bias = key_padding_mask[:, None, None, :].to(dtype=q.dtype)
+        attn_bias = attn_bias.masked_fill(attn_bias > 0, -10000.0)
+        out = F.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            attn_mask=attn_bias,
+            dropout_p=0.0,
+            is_causal=False,
+        )
+        out = out.transpose(1, 2).reshape(image_tokens.shape[0], text_tokens.shape[1] + image_tokens.shape[1], -1)
+        text_out, image_out = out.split([text_tokens.shape[1], image_tokens.shape[1]], dim=1)
+        return self.image_out_proj(image_out), self.text_out_proj(text_out)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text_tokens: torch.Tensor,
+        t_emb: torch.Tensor,
+        text_key_padding_mask: torch.Tensor,
+        text_attn_bias: torch.Tensor | None,
+        *,
+        height: int,
+        width: int,
+        coord_rope: MultimodalCoordinateRoPE | None = None,
+        image_coord_ids: torch.Tensor | None = None,
+        text_coord_ids: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del text_attn_bias
+        image_timestep_mod = self.image_mod(t_emb)
+        text_timestep_mod = self.text_mod(t_emb)
+        image_shift_attn, image_scale_attn, image_gate_attn, image_shift_mlp, image_scale_mlp, image_gate_mlp = (
+            self.image_scale_shift_table[None] + image_timestep_mod.reshape(x.shape[0], 6, -1)
+        ).chunk(6, dim=1)
+        text_shift_attn, text_scale_attn, text_gate_attn, text_shift_mlp, text_scale_mlp, text_gate_mlp = (
+            self.text_scale_shift_table[None] + text_timestep_mod.reshape(text_tokens.shape[0], 6, -1)
+        ).chunk(6, dim=1)
+
+        image_base = x
+        text_base = text_tokens
+        image_attn_in = modulate(self.image_norm1(image_base), image_shift_attn, image_scale_attn)
+        text_attn_in = modulate(self.text_norm1(text_base), text_shift_attn, text_scale_attn)
+        image_attn, text_attn = self._joint_attention(
+            image_attn_in,
+            text_attn_in,
+            text_key_padding_mask,
+            coord_rope,
+            image_coord_ids,
+            text_coord_ids,
+        )
+        if self.parallel_block:
+            x = image_base + image_gate_attn * image_attn + image_gate_mlp * self.image_mlp(
+                modulate(self.image_norm2(image_base), image_shift_mlp, image_scale_mlp),
+                height=height,
+                width=width,
+            )
+            text_tokens = text_base + text_gate_attn * text_attn + text_gate_mlp * self.text_mlp(
+                modulate(self.text_norm2(text_base), text_shift_mlp, text_scale_mlp)
+            )
+            return x, text_tokens
+
+        x = image_base + image_gate_attn * image_attn
+        text_tokens = text_base + text_gate_attn * text_attn
+        x = x + image_gate_mlp * self.image_mlp(
+            modulate(self.image_norm2(x), image_shift_mlp, image_scale_mlp),
+            height=height,
+            width=width,
+        )
+        text_tokens = text_tokens + text_gate_mlp * self.text_mlp(
+            modulate(self.text_norm2(text_tokens), text_shift_mlp, text_scale_mlp)
+        )
+        return x, text_tokens
+
+
+class BoomerFLADiT(nn.Module):
+    """Boomer DiT with FLA mixers, optional full image attention, and GLUMBConv FFNs."""
+
+    def __init__(self, config: BoomerFLADiTConfig = BoomerFLADiTConfig()) -> None:
+        super().__init__()
+        if config.patch_size <= 0:
+            raise ValueError(f"patch_size must be positive, got {config.patch_size}")
+        if config.latent_size % config.patch_size != 0:
+            raise ValueError(
+                f"latent_size={config.latent_size} must be divisible by patch_size={config.patch_size}"
+            )
+        if config.dual_stream_depth < 0:
+            raise ValueError(f"dual_stream_depth must be non-negative, got {config.dual_stream_depth}")
+        if config.dual_stream_depth > config.depth:
+            raise ValueError(f"dual_stream_depth={config.dual_stream_depth} exceeds depth={config.depth}")
+        self.config = config
+        hidden_dim = config.hidden_dim
+        self.patch_size = config.patch_size
+        self.token_grid_size = config.latent_size // config.patch_size
+        token_count = self.token_grid_size * self.token_grid_size
+        self.x_embedder = (
+            nn.Linear(config.latent_channels, hidden_dim)
+            if config.patch_size == 1
+            else nn.Conv2d(
+                config.latent_channels,
+                hidden_dim,
+                kernel_size=config.patch_size,
+                stride=config.patch_size,
+            )
+        )
+        self.pos_embed = nn.Parameter(torch.zeros(1, token_count, hidden_dim)) if config.use_abs_pos_embed else None
+        self.t_embedder = TimestepEmbedder(hidden_dim)
+        self.caption_embedder = CaptionEmbedder(config.text_dim, hidden_dim, config.text_seq_len)
+        self.attention_y_norm = (
+            AttentionRMSNorm(hidden_dim, scale_factor=config.y_norm_scale_factor) if config.y_norm else None
+        )
+        self.coord_embedder = (
+            MultimodalCoordinateRoPE(
+                hidden_dim // config.num_heads,
+                image_size=self.token_grid_size,
+                text_seq_len=config.text_seq_len,
+                theta=config.image_rope_theta,
+            )
+            if config.multimodal_coord_ids
+            else None
+        )
+        self.blocks = nn.ModuleList(
+            [
+                (
+                    BoomerFLADualStreamBlock(config, layer_idx=i)
+                    if i < config.dual_stream_depth
+                    else BoomerFLABlock(config, layer_idx=i)
+                )
+                for i in range(config.depth)
+            ]
+        )
+        self.final_norm = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
+        self.final_t_block = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 2))
+        self.out_proj = nn.Linear(hidden_dim, config.latent_channels * config.patch_size * config.patch_size)
+        self.initialize_weights()
+
+    def initialize_weights(self) -> None:
+        if self.pos_embed is not None:
+            nn.init.normal_(self.pos_embed, std=0.02)
+
+        for block in self.blocks:
+            if isinstance(block, BoomerFLADualStreamBlock):
+                nn.init.zeros_(block.image_mod[1].weight)
+                nn.init.zeros_(block.image_mod[1].bias)
+                nn.init.zeros_(block.text_mod[1].weight)
+                nn.init.zeros_(block.text_mod[1].bias)
+                nn.init.normal_(block.image_scale_shift_table, std=0.02)
+                nn.init.normal_(block.text_scale_shift_table, std=0.02)
+            else:
+                nn.init.zeros_(block.mod[1].weight)
+                nn.init.zeros_(block.mod[1].bias)
+                nn.init.normal_(block.scale_shift_table, std=0.02)
+                if block.use_image_attention:
+                    nn.init.zeros_(block.image_attn_mod[1].weight)
+                    nn.init.zeros_(block.image_attn_mod[1].bias)
+                    nn.init.normal_(block.image_attn_scale_shift_table, std=0.02)
+
+        nn.init.zeros_(self.final_t_block[1].weight)
+        nn.init.zeros_(self.final_t_block[1].bias)
+        nn.init.zeros_(self.out_proj.weight)
+        nn.init.zeros_(self.out_proj.bias)
+
+    def apply_y_norm(self, caption_tokens: torch.Tensor) -> torch.Tensor:
+        if self.attention_y_norm is None:
+            return caption_tokens
+        return self.attention_y_norm(caption_tokens)
+
+    def null_condition(
+        self,
+        batch_size: int,
+        *,
+        device: torch.device | str,
+        dtype: torch.dtype,
+        mask_dtype: torch.dtype | None = None,
+        token_num: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.caption_embedder.null_condition(
+            batch_size,
+            device=device,
+            dtype=dtype,
+            mask_dtype=mask_dtype,
+            token_num=token_num,
+        )
+
+    def apply_condition_dropout(
+        self,
+        text_embedding: torch.Tensor,
+        attention_mask: torch.Tensor,
+        probability: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if probability <= 0.0:
+            return text_embedding, attention_mask
+        batch_size = text_embedding.shape[0]
+        null_text, null_mask = self.null_condition(
+            batch_size,
+            device=text_embedding.device,
+            dtype=text_embedding.dtype,
+            mask_dtype=attention_mask.dtype,
+            token_num=text_embedding.shape[-2],
+        )
+        # torch.where over a per-sample bool. Avoids the bool(drop.any()) CUDA
+        # sync (which would defeat the training-loop sync removal) and skips
+        # the full-tensor .clone() that the previous in-place path required.
+        drop = torch.rand(batch_size, device=text_embedding.device) < probability
+        drop_text = drop.view(batch_size, *([1] * (text_embedding.dim() - 1)))
+        drop_mask = drop.view(batch_size, *([1] * (attention_mask.dim() - 1)))
+        text_embedding = torch.where(drop_text, null_text, text_embedding)
+        attention_mask = torch.where(drop_mask, null_mask, attention_mask)
+        return text_embedding, attention_mask
+
+    def forward(
+        self,
+        noisy_latent: torch.Tensor,
+        timesteps: torch.Tensor,
+        text_embedding: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        batch, channels, height, width = noisy_latent.shape
+        if channels != self.config.latent_channels:
+            raise ValueError(
+                f"Expected latent_channels={self.config.latent_channels}, got shape {tuple(noisy_latent.shape)}"
+            )
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"latent height/width must be divisible by patch_size={self.patch_size}, got {(height, width)}"
+            )
+        token_height = height // self.patch_size
+        token_width = width // self.patch_size
+        token_count = token_height * token_width
+        if self.pos_embed is not None and token_count != self.pos_embed.shape[1]:
+            raise ValueError(
+                f"absolute pos_embed expects {self.pos_embed.shape[1]} latent tokens, got {token_count}. "
+                "Disable it with --no-abs-pos-embed for variable latent sizes."
+            )
+        if text_embedding.shape[-1] != self.config.text_dim:
+            raise ValueError(f"text_embedding last dim must be {self.config.text_dim}, got {text_embedding.shape[-1]}")
+
+        text_tokens = self.caption_embedder(text_embedding)
+        text_tokens = self.apply_y_norm(text_tokens)
+        text_key_padding_mask = attention_mask == 0
+
+        if self.patch_size == 1:
+            x = noisy_latent.flatten(2).transpose(1, 2)
+            x = self.x_embedder(x)
+        else:
+            x = self.x_embedder(noisy_latent).flatten(2).transpose(1, 2)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        image_coord_ids = None
+        text_coord_ids = None
+        if self.coord_embedder is not None:
+            image_coord_ids = self.coord_embedder.image_ids(
+                batch,
+                height=token_height,
+                width=token_width,
+                device=x.device,
+            )
+            text_coord_ids = self.coord_embedder.text_ids(batch, text_tokens.shape[1], device=text_tokens.device)
+        text_attn_bias = text_key_padding_mask[:, None, None, :].to(dtype=x.dtype)
+        text_attn_bias = text_attn_bias.masked_fill(text_attn_bias > 0, -10000.0)
+        t_emb = self.t_embedder(timesteps)
+        use_ckpt = self.config.gradient_checkpointing and self.training
+        for block in self.blocks:
+            if getattr(block, "updates_text", False):
+                # Dual-stream block: returns (x, text_tokens).
+                # Non-tensor args (height, width, coord_rope, coord IDs) captured via closure.
+                _h, _w = token_height, token_width
+                _cr, _ii, _ti = self.coord_embedder, image_coord_ids, text_coord_ids
+                if use_ckpt:
+                    def _dual_fn(x, tt, te, mk, bi,
+                                 _blk=block, h=_h, w=_w, cr=_cr, ii=_ii, ti=_ti):
+                        return _blk(x, tt, te, mk, bi,
+                                    height=h, width=w, coord_rope=cr,
+                                    image_coord_ids=ii, text_coord_ids=ti)
+                    x, text_tokens = _ckpt(_dual_fn, x, text_tokens, t_emb,
+                                           text_key_padding_mask, text_attn_bias,
+                                           use_reentrant=False,
+                                           preserve_rng_state=False)
+                else:
+                    x, text_tokens = block(
+                        x, text_tokens, t_emb, text_key_padding_mask, text_attn_bias,
+                        height=token_height, width=token_width,
+                        coord_rope=self.coord_embedder,
+                        image_coord_ids=image_coord_ids, text_coord_ids=text_coord_ids,
+                    )
+            else:
+                # Single-stream block: returns x only.
+                _h, _w = token_height, token_width
+                if use_ckpt:
+                    def _single_fn(x, tt, te, mk, bi,
+                                   _blk=block, h=_h, w=_w):
+                        return _blk(x, tt, te, mk, bi, height=h, width=w)
+                    x = _ckpt(_single_fn, x, text_tokens, t_emb,
+                               text_key_padding_mask, text_attn_bias,
+                               use_reentrant=False,
+                               preserve_rng_state=False)
+                else:
+                    x = block(
+                        x, text_tokens, t_emb, text_key_padding_mask, text_attn_bias,
+                        height=token_height, width=token_width,
+                    )
+        final_mod = self.final_t_block(t_emb)
+        shift, scale = final_mod.reshape(batch, 2, -1).chunk(2, dim=1)
+        x = modulate(self.final_norm(x), shift, scale)
+        x = self.out_proj(x)
+        if self.patch_size == 1:
+            return x.transpose(1, 2).reshape(batch, channels, height, width)
+        patch = self.patch_size
+        x = x.reshape(batch, token_height, token_width, channels, patch, patch)
+        x = x.permute(0, 3, 1, 4, 2, 5).contiguous()
+        return x.reshape(batch, channels, height, width)

pipeline_boomer.py

@@ -0,0 +1,341 @@
+"""BoomerPipeline — HuggingFace DiffusionPipeline wrapper for Boomer FLA.
+
+Load with:
+    from diffusers import DiffusionPipeline
+    pipe = DiffusionPipeline.from_pretrained("akrao9/Boomer-T2I", trust_remote_code=True).to("cuda")
+    image = pipe("a photorealistic portrait of a woman with dark hair")[0]
+
+Requires:
+    pip install torch diffusers transformers accelerate safetensors
+    pip install git+https://github.com/Algomancer/STORK.git
+    pip install git+https://github.com/sustcsonglin/flash-linear-attention.git
+"""
+
+from __future__ import annotations
+
+import gc
+import json
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Optional, Union
+
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.utils import logging
+from PIL import Image
+
+logger = logging.get_logger(__name__)
+
+
+# ── pipeline output ────────────────────────────────────────────────────────────
+@dataclass
+class BoomerOutput:
+    """Return type of BoomerPipeline.__call__."""
+    images: List[Image.Image]
+
+    def __iter__(self):
+        return iter(self.images)
+
+    def __getitem__(self, idx):
+        return self.images[idx]
+
+    def __len__(self):
+        return len(self.images)
+
+
+# ── text encoding helpers ──────────────────────────────────────────────────────
+def _gemma_select_index(max_length: int) -> list[int]:
+    """Sana-style token selection: BOS + last (max_length-1) tokens."""
+    return [0] + list(range(-max_length + 1, 0))
+
+
+@torch.inference_mode()
+def _encode_prompts(
+    tokenizer: Any,
+    text_encoder: Any,
+    prompts: list[str],
+    max_length: int,
+    device: str,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Encode prompts with Gemma 4. Returns (embeddings, attention_mask)."""
+    token_kwargs = dict(
+        max_length=max_length,
+        padding="max_length",
+        truncation=True,
+        return_tensors="pt",
+    )
+    try:
+        tokens = tokenizer(text=prompts, **token_kwargs)
+    except TypeError:
+        tokens = tokenizer(prompts, **token_kwargs)
+
+    input_ids = tokens["input_ids"].to(device)
+    attention_mask = tokens["attention_mask"].to(device)
+
+    output = text_encoder(input_ids, attention_mask=attention_mask)
+    hidden = output[0] if isinstance(output, tuple) else output.last_hidden_state
+
+    select_idx = _gemma_select_index(max_length)
+    return (
+        hidden[:, select_idx, :].to(dtype=dtype),
+        attention_mask[:, select_idx],
+    )
+
+
+# ── latent helpers (inlined from latent_norm.py) ──────────────────────────────
+def _stat_tensor(value: Any, latent: torch.Tensor) -> torch.Tensor:
+    tensor = torch.as_tensor(value, device=latent.device, dtype=latent.dtype)
+    if tensor.ndim == 0:
+        return tensor
+    if tensor.ndim == 1:
+        if tensor.numel() != latent.shape[1]:
+            raise ValueError(f"latent stat has {tensor.numel()} channels, expected {latent.shape[1]}")
+        return tensor.view(1, tensor.numel(), 1, 1)
+    raise ValueError("latent stat must be scalar or 1-D channel list")
+
+
+def _denormalize(latent: torch.Tensor, mean: Any, std: Any) -> torch.Tensor:
+    return latent * _stat_tensor(std, latent) + _stat_tensor(mean, latent)
+
+
+# ── pipeline ───────────────────────────────────────────────────────────────────
+class BoomerPipeline(DiffusionPipeline):
+    """
+    Text-to-image generation with Boomer FLA (Flash Linear Attention DiT).
+
+    Components
+    ----------
+    transformer : BoomerFLADiT          657 M param FLA denoiser
+    vae         : AutoencoderDC          DC-AE f32c32 decoder
+    text_encoder: Gemma 4 2B decoder     1536-dim text embeddings
+    tokenizer   : AutoProcessor          Gemma tokenizer
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    def __init__(
+        self,
+        transformer,
+        vae,
+        text_encoder,
+        tokenizer,
+        *,
+        model_config,
+        latent_mean: Any = 0.0,
+        latent_std: Any = 1.0,
+        scaling_factor: float = 0.41407,
+        flow_shift: float = 1.5,
+        max_text_length: int = 300,
+    ) -> None:
+        super().__init__()
+        self.register_modules(
+            transformer=transformer,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+        )
+        self.model_config    = model_config
+        self.latent_mean     = latent_mean
+        self.latent_std      = latent_std
+        self.scaling_factor  = scaling_factor
+        self.flow_shift      = flow_shift
+        self.max_text_length = max_text_length
+
+    # ── loading ────────────────────────────────────────────────────────────────
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "BoomerPipeline":  # type: ignore[override]
+        """Load Boomer from a HuggingFace repo or local directory."""
+        from huggingface_hub import snapshot_download
+
+        token     = kwargs.pop("token", None) or kwargs.pop("use_auth_token", None)
+        dtype     = kwargs.pop("torch_dtype", torch.bfloat16)
+        cache_dir = kwargs.pop("cache_dir", None)
+
+        # ── 1. Resolve local snapshot path ────────────────────────────────────
+        if Path(pretrained_model_name_or_path).is_dir():
+            local = Path(pretrained_model_name_or_path)
+        else:
+            logger.info(f"Downloading Boomer snapshot from {pretrained_model_name_or_path} ...")
+            local = Path(snapshot_download(
+                pretrained_model_name_or_path,
+                token=token,
+                cache_dir=cache_dir,
+                ignore_patterns=["*.png", "*.jpg", "*.jpeg", "*.gif"],
+            ))
+
+        # Add snapshot dir to sys.path so local .py files import each other
+        if str(local) not in sys.path:
+            sys.path.insert(0, str(local))
+
+        # ── 2. Import local model/scheduler modules ────────────────────────────
+        from modeling_boomer_fla    import BoomerFLADiT, BoomerFLADiTConfig   # noqa: PLC0415
+        from scheduling_boomer_stork import make_stork_scheduler               # noqa: PLC0415
+        from safetensors.torch       import load_file                          # noqa: PLC0415
+
+        # ── 3. Load transformer config + weights ───────────────────────────────
+        transformer_dir = local / "transformer"
+        cfg_raw   = json.loads((transformer_dir / "config.json").read_text())
+        cfg_clean = {k: v for k, v in cfg_raw.items() if not k.startswith("_")}
+        model_config = BoomerFLADiTConfig(**cfg_clean)
+
+        logger.info("Loading Boomer FLA DiT weights ...")
+        state_dict  = load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors"))
+        transformer = BoomerFLADiT(model_config)
+        missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
+        if missing:
+            logger.warning(f"Missing keys in transformer state dict: {len(missing)}")
+        if unexpected:
+            logger.warning(f"Unexpected keys in transformer state dict: {len(unexpected)}")
+        transformer = transformer.to(dtype=dtype)
+
+        # ── 4. Read metadata from index / scheduler config ─────────────────────
+        model_index  = json.loads((local / "model_index.json").read_text())
+        sched_cfg    = json.loads((local / "scheduler" / "scheduler_config.json").read_text())
+        flow_shift   = float(sched_cfg.get("flow_shift", 1.5))
+        scaling_factor = float(sched_cfg.get("scaling_factor", 0.41407))
+        latent_norm  = model_index.get("latent_normalization", {})
+        latent_mean  = latent_norm.get("mean", 0.0)
+        latent_std   = latent_norm.get("std",  1.0)
+        te_info      = model_index.get("text_encoder", {})
+        te_repo      = te_info.get("repo_id", "google/gemma-4-E2B-it")
+        max_text_len = int(te_info.get("max_length", 300))
+        vae_info     = model_index.get("vae", {})
+        vae_repo     = vae_info.get("repo_id", "mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers")
+
+        # ── 5. Load VAE ─────────────────────────────────────────────────────────
+        from diffusers import AutoencoderDC  # noqa: PLC0415
+        logger.info(f"Loading VAE from {vae_repo} ...")
+        vae = AutoencoderDC.from_pretrained(vae_repo, torch_dtype=dtype, token=token)
+
+        # ── 6. Load text encoder ────────────────────────────────────────────────
+        from transformers import AutoModelForCausalLM, AutoProcessor  # noqa: PLC0415
+        logger.info(f"Loading text encoder from {te_repo} ...")
+        tokenizer = AutoProcessor.from_pretrained(te_repo, token=token)
+        text_encoder = AutoModelForCausalLM.from_pretrained(
+            te_repo, torch_dtype=dtype, token=token,
+        )
+        if hasattr(text_encoder, "get_decoder"):
+            text_encoder = text_encoder.get_decoder()
+
+        return cls(
+            transformer=transformer,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            model_config=model_config,
+            latent_mean=latent_mean,
+            latent_std=latent_std,
+            scaling_factor=scaling_factor,
+            flow_shift=flow_shift,
+            max_text_length=max_text_len,
+        )
+
+    # ── generation ─────────────────────────────────────────────────────────────
+    @torch.inference_mode()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        steps: int = 32,
+        seed: int = 42,
+        cfg_scale: float = 4.5,
+        cfg_rescale: float = 0.5,
+        substeps: int = 5,
+        offload_text_encoder: bool = True,
+        output_type: str = "pil",
+        **kwargs,
+    ) -> BoomerOutput:
+        """
+        Generate images from text prompts.
+
+        Parameters
+        ----------
+        prompt               : str or list[str]
+        steps                : denoising steps (default 32 with STORK-2)
+        seed                 : random seed
+        cfg_scale            : classifier-free guidance scale (4.0–5.0 recommended)
+        cfg_rescale          : CFG variance rescale (0.5 recommended)
+        substeps             : STORK-2 internal RK micro-steps (5 recommended)
+        offload_text_encoder : unload text encoder after encoding to free VRAM
+        output_type          : "pil" (default) or "latent"
+        """
+        from scheduling_boomer_stork import make_stork_scheduler  # noqa: PLC0415
+
+        prompts = [prompt] if isinstance(prompt, str) else prompt
+        batch   = len(prompts)
+        device  = self._execution_device
+        dtype   = next(self.transformer.parameters()).dtype
+
+        # ── Phase 1: encode text ───────────────────────────────────────────────
+        self.text_encoder.to(device)
+        self.text_encoder.eval()
+        text_emb, attn_mask = _encode_prompts(
+            self.tokenizer, self.text_encoder, prompts,
+            max_length=self.max_text_length,
+            device=device, dtype=dtype,
+        )
+        if offload_text_encoder:
+            self.text_encoder.to("cpu")
+            gc.collect()
+            if device != "cpu":
+                torch.cuda.empty_cache()
+
+        # ── Phase 2: denoise ───────────────────────────────────────────────────
+        self.transformer.to(device)
+        self.transformer.eval()
+
+        uncond_emb, uncond_mask = self.transformer.null_condition(
+            batch, device=device, dtype=dtype,
+            mask_dtype=attn_mask.dtype, token_num=text_emb.shape[-2],
+        )
+
+        sched = make_stork_scheduler(
+            steps=steps, device=device,
+            flow_shift=self.flow_shift,
+            solver_order=2, derivative_order=1, substeps=substeps,
+        )
+
+        gen    = torch.Generator(device=device).manual_seed(seed)
+        latent = torch.randn(
+            batch,
+            self.model_config.latent_channels,
+            self.model_config.latent_size,
+            self.model_config.latent_size,
+            generator=gen, device=device, dtype=dtype,
+        )
+
+        for step in range(steps):
+            sigma   = sched.sigmas[step].to(device=device, dtype=dtype)
+            lb      = latent.repeat(2, 1, 1, 1)
+            tb      = sigma.expand(batch * 2)
+            txt_b   = torch.cat([uncond_emb, text_emb], dim=0)
+            msk_b   = torch.cat([uncond_mask, attn_mask], dim=0)
+            uv, cv  = self.transformer(lb, tb, txt_b, msk_b).chunk(2)
+            guided  = uv + cfg_scale * (cv - uv)
+            if cfg_rescale > 0.0:
+                sc     = cv.std(dim=(1, 2, 3), keepdim=True)
+                sg     = guided.std(dim=(1, 2, 3), keepdim=True)
+                guided = cfg_rescale * guided * (sc / sg.clamp_min(1e-5)) + (1.0 - cfg_rescale) * guided
+            latent = sched.step(guided, sigma, latent, return_dict=True).prev_sample
+
+        if output_type == "latent":
+            return BoomerOutput(images=[latent])
+
+        # ── Phase 3: decode ────────────────────────────────────────────────────
+        self.vae.to(device)
+        self.vae.eval()
+
+        latent_dec = _denormalize(latent, self.latent_mean, self.latent_std)
+        # VAE expects latent / scaling_factor
+        latent_dec = latent_dec / self.scaling_factor
+        decoded    = self.vae.decode(latent_dec, return_dict=False)[0]
+
+        # [-1, 1] → [0, 1] → uint8 PIL
+        pixels = (decoded.float() / 2.0 + 0.5).clamp(0, 1)
+        images = []
+        for i in range(batch):
+            img_np = (pixels[i].permute(1, 2, 0).cpu().numpy() * 255).astype("uint8")
+            images.append(Image.fromarray(img_np))
+
+        return BoomerOutput(images=images)

scheduler/scheduler_config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "FlowMatchScheduler",
+  "flow_shift": 1.5,
+  "num_train_steps": 75000,
+  "sampler": "stork2",
+  "stork_substeps": 5,
+  "num_inference_steps": 32,
+  "scaling_factor": 0.41407
+}

scheduling_boomer_stork.py

@@ -0,0 +1,113 @@
+"""STORK flow-matching scheduler loader — self-contained for HuggingFace distribution.
+
+unshift_sigma inlined from boomer/sana_flow.py. No boomer package import needed.
+Requires: torch. STORK repo must be pip-installed or available on sys.path.
+  pip install git+https://github.com/Algomancer/STORK.git
+or clone alongside Boomer:
+  git clone https://github.com/Algomancer/STORK.git
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+
+
+# ── inlined from boomer/sana_flow.py ──────────────────────────────────────────
+def unshift_sigma(sigma: torch.Tensor, shift: float = 1.0) -> torch.Tensor:
+    """Invert Sana/FlowMatch sigma shift."""
+    if shift <= 0.0:
+        raise ValueError(f"flow shift must be positive, got {shift}")
+    if shift == 1.0:
+        return sigma
+    return sigma / (shift - (shift - 1.0) * sigma).clamp_min(1e-12)
+
+
+# ── STORK loader ───────────────────────────────────────────────────────────────
+def _stork_candidates() -> list[Path]:
+    """Common locations where STORKScheduler.py might live."""
+    here = Path(__file__).resolve().parent
+    return [
+        here,                    # bundled alongside this file (HF snapshot dir)
+        here.parent / "STORK",
+        here / "STORK",
+        Path("/content/STORK"),
+        Path("/tmp/STORK"),
+    ]
+
+
+def _find_stork_path() -> Path:
+    # 1. Already importable?
+    try:
+        import STORKScheduler  # noqa: F401
+        return Path(STORKScheduler.__file__)
+    except ImportError:
+        pass
+
+    # 2. Search known candidate directories
+    for candidate_dir in _stork_candidates():
+        p = candidate_dir / "STORKScheduler.py"
+        if p.is_file():
+            if str(candidate_dir) not in sys.path:
+                sys.path.insert(0, str(candidate_dir))
+            return p
+
+    # 3. Try pip-installed stork package
+    try:
+        import stork
+        p = Path(stork.__file__).parent / "STORKScheduler.py"
+        if p.is_file():
+            return p
+    except ImportError:
+        pass
+
+    searched = "\n  ".join(str(d / "STORKScheduler.py") for d in _stork_candidates())
+    raise FileNotFoundError(
+        "Could not find STORKScheduler.py. Options:\n"
+        "  pip install git+https://github.com/ZT220501/STORK.git\n"
+        "or clone it next to Boomer:\n"
+        "  git clone https://github.com/ZT220501/STORK.git\n"
+        f"Searched:\n  {searched}"
+    )
+
+
+def make_stork_scheduler(
+    *,
+    steps: int,
+    device: str | torch.device,
+    flow_shift: float,
+    solver_order: int,
+    derivative_order: int,
+    substeps: int,
+    start_sigma: float | None = None,
+) -> Any:
+    stork_path = _find_stork_path()
+    spec = importlib.util.spec_from_file_location("STORKScheduler", stork_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not load STORK scheduler from {stork_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    STORKScheduler = module.STORKScheduler
+
+    scheduler = STORKScheduler(
+        shift=flow_shift,
+        solver_order=solver_order,
+        prediction_type="flow_prediction",
+        derivative_order=derivative_order,
+        s=substeps,
+    )
+    if start_sigma is not None:
+        start = torch.tensor(float(start_sigma), dtype=torch.float32)
+        base_start = unshift_sigma(start, flow_shift)
+        base_sigmas = torch.linspace(float(base_start.item()), 0.0, steps + 1, dtype=torch.float32)[:-1].tolist()
+    else:
+        base_sigmas = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float32)[:-1].tolist()
+
+    scheduler.set_timesteps(num_inference_steps=steps, device=device, sigmas=base_sigmas)
+    scheduler.dt_list = scheduler.dt_list.to(device=device)
+    scheduler.sigmas = scheduler.sigmas.to(device=device)
+    return scheduler

transformer/config.json

@@ -0,0 +1,35 @@
+{
+  "model_type": "boomer_fla",
+  "latent_channels": 32,
+  "latent_size": 32,
+  "text_dim": 1536,
+  "text_seq_len": 384,
+  "hidden_dim": 896,
+  "depth": 24,
+  "num_heads": 14,
+  "mlp_ratio": 2.5,
+  "y_norm": true,
+  "y_norm_scale_factor": 0.01,
+  "mixer_type": "fla_gated_deltanet",
+  "fla_mode": "chunk",
+  "fla_feature_map": "relu",
+  "fla_bidirectional": true,
+  "use_short_conv": true,
+  "conv_size": 4,
+  "image_attention_every": 6,
+  "image_attention_backend": "sdpa",
+  "image_attention_rope": true,
+  "image_rope_theta": 10000.0,
+  "cross_attention_backend": "sdpa",
+  "cross_attention_qk_norm": true,
+  "parallel_block": true,
+  "dual_stream_depth": 2,
+  "multimodal_coord_ids": true,
+  "use_abs_pos_embed": false,
+  "patch_size": 1,
+  "gradient_checkpointing": false,
+  "_architecture_class": "BoomerFLADiT",
+  "_boomer_version": "1.0.0",
+  "_image_size_px": 1024,
+  "_latent_tokens": 1024
+}

transformer/diffusion_pytorch_model.safetensors

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