Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text

.meta.json

@@ -0,0 +1,19 @@
+{
+  "source_model": {
+    "name": "OpenMOSS-Team/MOSS-Image-Base",
+    "commit_id": "85975793d8b3d432f1b6e27f352abaa4c2a4c0d3"
+  },
+  "env": {
+    "python": "3.12.13",
+    "packages": {
+      "Pillow": "12.3.0",
+      "diffusers": "0.40.0.dev0",
+      "huggingface_hub": "1.27.0",
+      "safetensors": "0.8.0",
+      "tokenizers": "0.23.0rc0",
+      "torch": "2.13.0+cu126",
+      "transformers": "5.15.0.dev0"
+    },
+    "cuda": "12.6"
+  }
+}

README.md

@@ -0,0 +1,244 @@
+---
+library_name: diffusers
+base_model:
+- OpenMOSS-Team/MOSS-Image-Base
+---
+
+This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from [OpenMOSS-Team/MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base).
+
+Structural comparison with `OpenMOSS-Team/MOSS-Image-Base`:
+
+| Feature | Source | Tiny |
+|---|---:|---:|
+| Pipeline | custom MOSS Image | same |
+| Text vocabulary / precision | 248,320 / BF16 | same |
+| Qwen3.5 text layers and cadence | 32; 3 linear + 1 full | same |
+| Text width / attention head dimension | 4,096 / 256 | 32 / 32 |
+| Text vision tower | 27 layers | same depth; width 32 |
+| DiT blocks | 42 | same |
+| DiT width / heads / FFN | 3,072 / 24 / 8,192 | 64 / 2 / 128 |
+| DiT latent channels | 128 | same |
+| Flux2 VAE blocks / layers per block | 4 / 2 | same; widths reduced |
+| Checkpoint precision | BF16 text/DiT; FP32 VAE/RoPE | same |
+
+The source's custom scheduler, DiT, tokenizer and pipeline code are copied from
+the pinned revision.  The transformer keeps all 42 blocks, and the VAE stays
+FP32 because that is the source checkpoint representation and decoding path.
+
+| File path | Size |
+|------|------|
+| text_encoder/model.safetensors | 17.5MB |
+| transformer/diffusion_pytorch_model.safetensors | 3.8MB |
+| vae/diffusion_pytorch_model.safetensors | 5.8MB |
+
+
+### Example usage:
+
+- Transformers (✅ verified)
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+model_id = "tiny-random/moss-image"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = DiffusionPipeline.from_pretrained(
+    model_id,
+    custom_pipeline="pipeline_moss_image",
+    trust_remote_code=True,
+    torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
+)
+pipe.to(device)
+image = pipe(
+    "A red fox in a snowy pine forest",
+    height=32,
+    width=32,
+    num_inference_steps=1,
+    guidance_scale=1.0,
+    generator=torch.Generator("cpu").manual_seed(42),
+).images[0]
+print(image.size)
+```
+
+### Codes to create this repo (✅ verified):
+
+<details><summary>Click to expand</summary>
+
+```python
+import importlib.util
+import json
+from pathlib import Path
+from shutil import copyfile
+
+import torch
+from diffusers import AutoencoderKLFlux2
+from huggingface_hub import hf_hub_download, list_repo_files
+from transformers import AutoConfig, Qwen3_5Model, set_seed
+
+source_model_id = "OpenMOSS-Team/MOSS-Image-Base"
+save_folder = "/tmp/tiny-random/moss-image"
+source_revision = "85975793d8b3d432f1b6e27f352abaa4c2a4c0d3"
+set_seed(42)
+folder = Path(save_folder)
+folder.mkdir(parents=True, exist_ok=True)
+
+def load_source_module():
+    path = hf_hub_download(source_model_id, "pipeline_moss_image.py", revision=source_revision)
+    spec = importlib.util.spec_from_file_location("moss_image_source", path)
+    module = importlib.util.module_from_spec(spec)
+    if spec.loader is None:
+        raise RuntimeError("The pinned MOSS pipeline has no import loader")
+    spec.loader.exec_module(module)
+    return module
+
+def initialize(module):
+    with torch.no_grad():
+        for name, parameter in sorted(module.named_parameters()):
+            if not parameter.is_floating_point():
+                continue
+            if name.endswith("A_log") or name.endswith("dt_bias"):
+                continue
+            if name.startswith("language_model") and (
+                name.endswith("input_layernorm.weight")
+                or name.endswith("post_attention_layernorm.weight")
+                or name.endswith(".linear_attn.norm.weight")
+                or name.endswith("language_model.norm.weight")
+                or name.endswith(".self_attn.q_norm.weight")
+                or name.endswith(".self_attn.k_norm.weight")
+            ):
+                parameter.zero_()
+            elif ("norm" in name or ".bn." in name) and name.endswith("weight"):
+                parameter.fill_(1)
+            else:
+                torch.nn.init.normal_(parameter, mean=0.0, std=0.02)
+
+source = load_source_module()
+
+raw_text = json.loads(Path(hf_hub_download(
+    source_model_id, "text_encoder/config.json", revision=source_revision,
+)).read_text(encoding="utf-8"))
+text = raw_text["text_config"]
+text.update(
+    hidden_size=32,
+    head_dim=32,
+    num_attention_heads=1,
+    num_key_value_heads=1,
+    intermediate_size=64,
+    linear_key_head_dim=32,
+    linear_value_head_dim=32,
+    linear_num_key_heads=1,
+    linear_num_value_heads=1,
+    num_hidden_layers=32,
+    layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"] * 8,
+)
+text["rope_parameters"]["mrope_section"] = [1, 1, 2]
+raw_text["vision_config"].update(
+    hidden_size=32,
+    intermediate_size=64,
+    num_heads=1,
+    out_hidden_size=32,
+    depth=27,
+)
+text_folder = folder / "text_encoder"
+text_folder.mkdir(exist_ok=True)
+(text_folder / "config.json").write_text(json.dumps(raw_text, indent=2), encoding="utf-8")
+text_encoder = Qwen3_5Model(AutoConfig.from_pretrained(text_folder)).to(torch.bfloat16)
+initialize(text_encoder)
+text_encoder.save_pretrained(text_folder, safe_serialization=True)
+
+# The source stores a 42-block custom DiT in BF16.  Its VAE-facing 128
+# channels remain fixed: the pipeline unshuffles them to the VAE's 32
+# latent channels before decoding.
+transformer = source.MOSSImageTransformer(
+    dim=64,
+    ffn_dim=128,
+    num_heads=2,
+    num_layers=42,
+    text_dim=32,
+    freq_dim=32,
+    latent_channels=128,
+).to(torch.bfloat16)
+initialize(transformer)
+transformer.save_pretrained(folder / "transformer", safe_serialization=True)
+
+raw_vae = json.loads(Path(hf_hub_download(
+    source_model_id, "vae/config.json", revision=source_revision,
+)).read_text(encoding="utf-8"))
+raw_vae.update(
+    block_out_channels=[32, 32, 64, 64],
+    decoder_block_out_channels=[32, 32, 64, 64],
+    norm_num_groups=32,
+    sample_size=32,
+)
+vae = AutoencoderKLFlux2.from_config(raw_vae).to(torch.float32)
+initialize(vae)
+vae.save_pretrained(folder / "vae", safe_serialization=True)
+
+scheduler = source.MOSSImageScheduler()
+scheduler.save_pretrained(folder / "scheduler")
+
+# Preserve all executable source code plus tokenizer assets byte-for-byte.
+assets = {
+    "model_index.json", "pipeline_moss_image.py", "scheduler/pipeline_moss_image.py",
+    "scheduler/scheduler_config.json", "tokenizer/chat_template.jinja",
+    "tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json",
+    "transformer/pipeline_moss_image.py",
+}
+for filename in sorted(set(list_repo_files(source_model_id, revision=source_revision)) & assets):
+    destination = folder / filename
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    copyfile(hf_hub_download(source_model_id, filename, revision=source_revision), destination)
+
+pipe = source.MOSSImagePipeline(transformer, scheduler, text_encoder, None, vae)
+```
+
+</details>
+
+### Printing the model:
+
+<details><summary>Click to expand</summary>
+
+```text
+MOSSImagePipeline {
+  "_class_name": "MOSSImagePipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "scheduler": [
+    "moss_image_source",
+    "MOSSImageScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Qwen3_5Model"
+  ],
+  "tokenizer": [
+    null,
+    null
+  ],
+  "transformer": [
+    "moss_image_source",
+    "MOSSImageTransformer"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ]
+}
+```
+
+</details>
+
+### Test environment:
+
+- Pillow: 12.3.0
+- diffusers: 0.40.0.dev0
+- huggingface_hub: 1.27.0
+- safetensors: 0.8.0
+- tokenizers: 0.23.0rc0
+- torch: 2.13.0+cu126
+- transformers: 5.15.0.dev0
+
+### Change log:
+
+- 2026-10-23: Initial version.
+
+This repo is generated by GPT-6.1-Sol

model_index.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "MOSSImagePipeline",
+  "_diffusers_version": "0.38.0",
+  "scheduler": [
+    "pipeline_moss_image",
+    "MOSSImageScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Qwen3_5Model"
+  ],
+  "tokenizer": [
+    "transformers",
+    "Qwen2Tokenizer"
+  ],
+  "transformer": [
+    "pipeline_moss_image",
+    "MOSSImageTransformer"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ]
+}

pipeline_moss_image.py

@@ -0,0 +1,440 @@
+"""MOSSImage modeling and text-to-image pipeline (no Motor/TorchTitan imports).
+
+Dependencies: torch>=2.9, diffusers>=0.38, transformers>=5.3, safetensors,
+accelerate, Pillow. Copy this file anywhere and import it normally.
+
+    import torch
+    # Import MOSSImagePipeline from this file for component-level loading.
+    pipe = MOSSImagePipeline.from_components_pretrained(
+        transformer_path="exported_transformer",  # save_pretrained directory
+        text_encoder_path="Qwen/Qwen3.5-9B",
+        vae_path="black-forest-labs/FLUX.2-klein-4B",
+    ).to("cuda")
+    pipe("A cat", generator=torch.Generator("cuda").manual_seed(42)).images[0].save("cat.png")
+
+Default guidance uses CFG 3 and skips block 22 only in the negative branch
+(zero-based main DiT block indices; two forwards per sampling step):
+    pipe("A cat", guidance_scale=3.0, uncond_skip_layers=[22])
+    pipe("A cat", uncond_skip_layers=[])  # disable negative-branch skipping
+Optional additional positive-branch SLG adds a third forward:
+    pipe("A cat", slg_layers=[24], slg_scale=3.0)
+The additional SLG layer/scale above illustrates the API, not a quality-tuned setting.
+
+For a consolidated training state dict, create MOSSImageTransformer and
+call load_training_state_dict(state), then save_pretrained("exported_transformer").
+Select raw or EMA weights explicitly before calling it; distributed checkpoint
+shards/optimizer states are not model state dicts. Defaults match the finalized dense MOSSImage model,
+not MoE, CFG-distilled, video, or reference-image variants. Prompt strings are
+already-rendered captions (dataset-specific quality/domain tags are not added).
+
+First-release defaults: bf16 DiT/text-encoder weights, fp32 RoPE frequencies
+and phase/sin/cos calculation, then RoPE cast to the video activation dtype.
+Initial Gaussian noise is sampled directly in fp32, including with bf16 weights.
+CFG and Euler latent updates use fp32; VAE decoding defaults to fp32. Scheduler
+shift defaults to 5.0. Exact numerical parity with the original fp32-master
+inference pipeline is not a release requirement.
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from diffusers import AutoencoderKLFlux2, DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin  # Required by custom-component discovery.
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
+        return y.to(x.dtype) * self.weight
+
+
+def _situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))
+
+
+def _validate_skip_layers(layers, depth, name='skip_layers'):
+    layers = () if layers is None else tuple(layers)
+    if any(type(index) is not int or not 0 <= index < depth for index in layers):
+        raise ValueError(f'{name} must contain integer block indices in [0, {depth}).')
+    return frozenset(layers)
+
+
+def _rotate(x, rope):
+    sin, cos = rope.chunk(2, -1)
+    sin = torch.cat((sin, sin), -1).unsqueeze(-2)
+    cos = torch.cat((cos, cos), -1).unsqueeze(-2)
+    a, b = x.chunk(2, -1)
+    return x * cos + torch.cat((-b, a), -1) * sin
+
+
+def _attention(q, k, v, rope, heads):
+    shape = (*q.shape[:2], heads, q.shape[-1] // heads)
+    q = _rotate(q.reshape(shape), rope).transpose(1, 2)
+    k = _rotate(k.reshape(shape), rope).transpose(1, 2)
+    v = v.reshape(shape).transpose(1, 2)
+    return F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v).transpose(1, 2).flatten(2)
+
+
+class _TextBlock(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def forward(self, x, rope):
+        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj(_attention(q, k, v, rope, self.heads))
+        y = self.norm2(x)
+        return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class _TextProjection(nn.Module):
+    def __init__(self, text_dim, dim, ffn_dim, heads):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)])
+
+    def forward(self, x, rope):
+        x = self.proj(x.to(self.proj.weight.dtype))
+        for block in self.blocks:
+            x = block(x, rope)
+        return x
+
+
+class _Block(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.proj = nn.ModuleDict({
+            name: nn.Linear(dim, dim, bias=False) for name in ('q', 'k', 'v', 'o')
+        })
+        self.proj['ffn_up_gate'] = nn.Linear(dim, 2 * ffn_dim, bias=False)
+        self.proj['ffn_down'] = nn.Linear(ffn_dim, dim, bias=False)
+
+    def forward(self, x, rope, modulation):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
+        y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
+        dtype = self.proj['q'].weight.dtype
+        y = y.to(dtype)
+        q = self.q_norm(self.proj['q'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(self.proj['k'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj['o'](_attention(q, k, self.proj['v'](y), rope, self.heads)) * gate_a[:, None]
+        y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None]
+        hidden, gate = self.proj['ffn_up_gate'](y.to(dtype)).chunk(2, -1)
+        return x + self.proj['ffn_down'](_situ(hidden, gate)) * gate_f[:, None]
+
+
+class _RoPE(nn.Module):
+    def __init__(self, head_dim, base):
+        super().__init__()
+        self.num_bands = head_dim // 8
+        self.base = base
+        self.register_buffer('bands', self._default_bands(torch.device('cpu')))
+
+    def _default_bands(self, device):
+        # Match training's CPU fp32 construction, including under meta init.
+        exponent = torch.arange(self.num_bands, device='cpu', dtype=torch.float32) / self.num_bands
+        return (1 / (self.base ** exponent)).to(device)
+
+    def _apply(self, fn, recurse=True):
+        super()._apply(fn, recurse=recurse)
+        # .float() cannot recover frequencies already rounded by .to(bf16).
+        self.bands = self._default_bands(self.bands.device)
+        return self
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        super()._load_from_state_dict(*args, **kwargs)
+        # Frequencies are derived constants; retain the key for checkpoint compatibility.
+        self.bands = self._default_bands(self.bands.device)
+
+    def forward(self, height, width, text_length, device):
+        # Diffusers/Accelerate can install tensors directly, bypassing load hooks.
+        if self.bands.dtype != torch.float32:
+            self.bands = self._default_bands(device)
+        # Coordinates: (r, f, h, w); sequence order is image followed by text.
+        h, w = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing='ij')
+        image = torch.stack((torch.zeros_like(h), torch.full_like(h, text_length), h, w), -1).flatten(0, 1).float()
+        image[:, 2] -= (height - 1) / 2
+        image[:, 3] -= (width - 1) / 2
+        text = torch.zeros(text_length, 4, device=device, dtype=torch.float32)
+        text[:, 1] = torch.arange(text_length, device=device)
+        phase = torch.cat((image, text), 0).unsqueeze(-1) * self.bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2).unsqueeze(0)
+
+
+class MOSSImageTransformer(ModelMixin, ConfigMixin):
+    """MOSSImage DiT with training-compatible state-dict names and SDPA attention.
+
+    Each sample is evaluated at its true text length, including text refinement;
+    padded tokens cannot influence either text or joint attention.
+    """
+    @register_to_config
+    def __init__(self, dim=3072, ffn_dim=8192, num_heads=24, num_layers=42,
+                 text_dim=4096, freq_dim=256, latent_channels=128, rope_base=10000.0):
+        super().__init__()
+        if dim % num_heads or (dim // num_heads) % 8 or freq_dim % 2:
+            raise ValueError('Require dim divisible by heads, head_dim divisible by 8, and even freq_dim.')
+        self.video_embedder = nn.Linear(latent_channels, dim)
+        self.text_embedder = _TextProjection(text_dim, dim, ffn_dim, num_heads)
+        self.rope = _RoPE(dim // num_heads, rope_base)
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = nn.ModuleList([_Block(dim, ffn_dim, num_heads) for _ in range(num_layers)])
+        self.final_norm_video = _RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+
+    def load_training_state_dict(self, state_dict, *, prefix=''):
+        """Load a consolidated raw/EMA model mapping, strictly; no silent key drops."""
+        if prefix:
+            if any(not key.startswith(prefix) for key in state_dict):
+                raise ValueError('All keys must start with the explicit prefix.')
+            state_dict = {key[len(prefix):]: value for key, value in state_dict.items()}
+        return self.load_state_dict(state_dict, strict=True)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states,
+                encoder_attention_mask=None, return_dict=True, skip_layers=None):
+        skip_layers = _validate_skip_layers(skip_layers, len(self.layers))
+        if hidden_states.ndim != 4 or hidden_states.shape[1] != self.config.latent_channels:
+            raise ValueError('hidden_states must be [B, latent_channels, H, W].')
+        batch, _, height, width = hidden_states.shape
+        if encoder_hidden_states.ndim != 3 or encoder_hidden_states.shape[0] != batch:
+            raise ValueError('encoder_hidden_states must be [B, S, text_dim].')
+        mask = encoder_attention_mask
+        if mask is None:
+            mask = torch.ones(encoder_hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        if mask.shape != encoder_hidden_states.shape[:2] or not torch.all((mask == 0) | (mask == 1)):
+            raise ValueError('encoder_attention_mask must be a binary [B, S] mask.')
+        times = torch.as_tensor(timestep, device=hidden_states.device, dtype=torch.float32).reshape(-1)
+        times = times.expand(batch)
+        # Training computes the sinusoid in float64 and feeds float32 into its MLP.
+        half = self.config.freq_dim // 2
+        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)
+        phase = times.double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = self.time_embedding(embedding.to(self.time_embedding[0].weight.dtype)).float()
+        mod = self.shared_block_modulation(time.to(self.shared_block_modulation.weight.dtype)).view(batch, 6, -1)
+        head = self.video_head_modulation(time.to(self.video_head_modulation.weight.dtype)).view(batch, 2, -1)
+        outputs = []
+        for index in range(batch):
+            context = encoder_hidden_states[index:index + 1, mask[index].bool()]
+            if context.shape[1] == 0:
+                raise ValueError('Each sample needs at least one valid text token.')
+            video = hidden_states[index:index + 1].flatten(2).transpose(1, 2)
+            video = self.video_embedder(video.to(self.video_embedder.weight.dtype).contiguous())
+            # Match the reference activation dtype under fp32 weights + autocast.
+            rope = self.rope(height, width, context.shape[1], hidden_states.device).to(video.dtype)
+            text = self.text_embedder(context, rope[:, height * width:])
+            x = torch.cat((video, text), 1)
+            for layer_index, layer in enumerate(self.layers):
+                if layer_index not in skip_layers:
+                    x = layer(x, rope, mod[index:index + 1])
+            x = self.final_norm_video(x[:, :height * width].to(self.final_norm_video.weight.dtype))
+            x = x * (1 + head[index:index + 1, 1, None]) + head[index:index + 1, 0, None]
+            x = self.final_linear_video(x.to(self.final_linear_video.weight.dtype))
+            outputs.append(x.transpose(1, 2).reshape(1, -1, height, width))
+        sample = torch.cat(outputs)
+        return Transformer2DModelOutput(sample=sample) if return_dict else (sample,)
+
+
+class MOSSImageScheduler(FlowMatchEulerDiscreteScheduler):
+    """Euler integration using the training project's V2, nearest-grid schedule."""
+    @register_to_config
+    def __init__(self, num_train_timesteps=1000, shift=5.0, training_shift=1.0):
+        super().__init__(num_train_timesteps=num_train_timesteps, shift=1.0)
+        self.register_to_config(shift=shift, training_shift=training_shift)
+
+    def set_timesteps(self, num_inference_steps=50, device=None):
+        if num_inference_steps < 1 or self.config.shift <= 0 or self.config.training_shift <= 0:
+            raise ValueError('Steps and shifts must be positive.')
+        total = self.config.num_train_timesteps
+        def grid(count, shift):
+            lower = shift / (total + shift - 1)
+            values = torch.linspace(1, lower, count)
+            return shift * values / (1 + (shift - 1) * values)
+        sigmas = grid(num_inference_steps, self.config.shift)
+        training = grid(total, self.config.training_shift).flip(0)
+        upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+        lower = (upper - 1).clamp(min=0)
+        nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
+        sigmas = torch.unique_consecutive(training[nearest])
+        # Our sigmas have already been shifted and snapped. Avoid a second shift.
+        self._shift = 1.0
+        super().set_timesteps(device=device, sigmas=sigmas.tolist())
+
+
+class MOSSImagePipeline(DiffusionPipeline):
+    model_cpu_offload_seq = 'text_encoder->transformer->vae'
+    _optional_components = ['text_encoder', 'tokenizer', 'vae']
+
+    def __init__(self, transformer: MOSSImageTransformer,
+                 scheduler: MOSSImageScheduler, text_encoder=None, tokenizer=None, vae=None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler,
+                              text_encoder=text_encoder, tokenizer=tokenizer, vae=vae)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=16)
+
+    @classmethod
+    def from_components_pretrained(cls, transformer_path, text_encoder_path, vae_path,
+                                   torch_dtype=torch.bfloat16, vae_dtype=torch.float32,
+                                   local_files_only=False):
+        common = dict(local_files_only=local_files_only)
+        transformer = MOSSImageTransformer.from_pretrained(transformer_path, torch_dtype=torch_dtype, **common)
+        tokenizer = AutoTokenizer.from_pretrained(text_encoder_path, **common)
+        text_encoder = AutoModel.from_pretrained(text_encoder_path, dtype=torch_dtype, **common).eval()
+        subfolder = None if (Path(vae_path) / 'config.json').is_file() else 'vae'
+        vae = AutoencoderKLFlux2.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=vae_dtype, **common).eval()
+        return cls(transformer.eval(), MOSSImageScheduler(), text_encoder, tokenizer, vae)
+
+    @torch.no_grad()
+    def encode_prompt(self, prompt, max_sequence_length=4096):
+        if self.text_encoder is None or self.tokenizer is None:
+            raise ValueError('Text encoder/tokenizer required for string prompts.')
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        if not prompts or not all(isinstance(p, str) for p in prompts):
+            raise ValueError('prompt must be a string or a nonempty list of strings.')
+        rows = []
+        # Unpadded individual forwards also avoid padding effects in Qwen GDN.
+        for value in prompts:
+            formatted = self.tokenizer.apply_chat_template(
+                [{'role': 'user', 'content': value}], tokenize=False,
+                add_generation_prompt=True, enable_thinking=True)
+            inputs = self.tokenizer(formatted, return_tensors='pt', truncation=True,
+                                    max_length=max_sequence_length).to(self._execution_device)
+            output = self.text_encoder(**inputs, use_cache=False, return_dict=True)
+            rows.append(output.last_hidden_state[0])
+        length = max(row.shape[0] for row in rows)
+        embeds = rows[0].new_zeros(len(rows), length, rows[0].shape[-1])
+        mask = torch.zeros(len(rows), length, device=embeds.device, dtype=torch.bool)
+        for index, row in enumerate(rows):
+            embeds[index, :row.shape[0]] = row
+            mask[index, :row.shape[0]] = True
+        return embeds, mask
+
+    @torch.no_grad()
+    def __call__(self, prompt=None, *, height=1024, width=1024, num_inference_steps=50,
+                 guidance_scale=3.0, negative_prompt='', num_images_per_prompt=1,
+                 generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
+                 negative_prompt_embeds=None, negative_prompt_attention_mask=None,
+                 max_sequence_length=4096, output_type='pil', return_dict=True,
+                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 slg_start=0.0, slg_stop=1.0):
+        """Generate images with optional full-block skipping (zero-based indices).
+
+        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        prediction with a skipped-block prediction, retaining two forwards per
+        step. Pass None or an empty sequence to disable skipping. slg_layers
+        adds a skipped positive-prompt forward and the fp32 correction
+        slg_scale * (cond - cond_skip) to CFG. With guidance_scale=1, active
+        SLG uses cond_skip + slg_scale * (cond - cond_skip), matching Diffusers'
+        default SkipLayerGuidance formulation, with two forwards and no negative
+        branch. Both options can combine. slg_scale=1 then gives normal cond.
+        Both uncond skipping and SLG are active where
+        slg_start <= step_index / num_steps < slg_stop;
+        the interval follows denoising order, from initial noise to final image.
+        Outside it, inference uses normal CFG without skipping.
+        slg_scale=0 or empty slg_layers disables the additional forward.
+        """
+        uncond_skip_layers = _validate_skip_layers(
+            uncond_skip_layers, len(self.transformer.layers), 'uncond_skip_layers')
+        slg_layers = _validate_skip_layers(slg_layers, len(self.transformer.layers), 'slg_layers')
+        if not math.isfinite(slg_scale) or slg_scale < 0:
+            raise ValueError('slg_scale must be finite and nonnegative.')
+        if not 0 <= slg_start <= slg_stop <= 1:
+            raise ValueError('Require 0 <= slg_start <= slg_stop <= 1.')
+        if height <= 0 or width <= 0 or height % 16 or width % 16:
+            raise ValueError('height and width must be positive multiples of 16.')
+        if output_type not in {'pil', 'np', 'pt', 'latent'}:
+            raise ValueError('output_type must be pil, np, pt or latent.')
+        if num_images_per_prompt < 1 or not math.isfinite(guidance_scale) or guidance_scale < 0:
+            raise ValueError('Invalid image count or guidance scale.')
+        if not 1 <= max_sequence_length <= 4096:
+            raise ValueError('max_sequence_length must be between 1 and 4096.')
+        if (prompt is None) == (prompt_embeds is None):
+            raise ValueError('Provide exactly one of prompt and prompt_embeds.')
+        device = self._execution_device
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length)
+        batch = prompt_embeds.shape[0]
+        def prepare(embeds, mask):
+            if embeds.ndim != 3 or embeds.shape[0] != batch:
+                raise ValueError('Prompt embeddings must be [B, S, C] with matching batch size.')
+            if mask is None:
+                mask = torch.ones(embeds.shape[:2], dtype=torch.bool, device=embeds.device)
+            return (embeds.to(device).repeat_interleave(num_images_per_prompt, 0),
+                    mask.to(device).repeat_interleave(num_images_per_prompt, 0))
+        positive, positive_mask = prepare(prompt_embeds, prompt_attention_mask)
+        do_cfg = guidance_scale != 1.0
+        if do_cfg:
+            if negative_prompt_embeds is None:
+                negatives = [negative_prompt] * batch if isinstance(negative_prompt, str) else negative_prompt
+                if len(negatives) != batch:
+                    raise ValueError('negative_prompt must match prompt batch size.')
+                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(negatives, max_sequence_length)
+            negative, negative_mask = prepare(negative_prompt_embeds, negative_prompt_attention_mask)
+        shape = (batch * num_images_per_prompt, self.transformer.config.latent_channels, height // 16, width // 16)
+        if latents is None:
+            # Sample in fp32: promoting bf16 noise afterward cannot restore its precision.
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        elif tuple(latents.shape) != shape:
+            raise ValueError(f'latents must have shape {shape}.')
+        else:
+            latents = latents.to(device=device, dtype=torch.float32)
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        num_steps = len(self.scheduler.timesteps)
+        for step_index, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            skip_active = slg_start <= step_index / num_steps < slg_stop
+            # CFG must combine fp32 predictions, even when the DiT runs in bf16.
+            # Converting only the combined velocity would preserve rounding errors.
+            prediction = self.transformer(latents, timestep, positive, positive_mask).sample.float()
+            conditional = prediction
+            if do_cfg:
+                skip_kwargs = {'skip_layers': uncond_skip_layers} if uncond_skip_layers and skip_active else {}
+                uncond = self.transformer(latents, timestep, negative, negative_mask, **skip_kwargs).sample.float()
+                prediction = uncond + guidance_scale * (prediction - uncond)
+            if slg_layers and slg_scale > 0 and skip_active:
+                skipped = self.transformer(
+                    latents, timestep, positive, positive_mask, skip_layers=slg_layers).sample.float()
+                if not do_cfg:
+                    prediction = skipped
+                prediction = prediction + slg_scale * (conditional - skipped)
+            latents = self.scheduler.step(prediction, timestep, latents, return_dict=False)[0]
+        if output_type == 'latent':
+            images = latents
+        else:
+            if self.vae is None:
+                raise ValueError('A FLUX.2 VAE is required for decoded output.')
+            z = latents.to(dtype=self.vae.dtype)
+            bn = self.vae.bn
+            eps = self.vae.config.batch_norm_eps
+            z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1)
+            b, c, h, w = z.shape
+            z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+            images = self.vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+            images = self.image_processor.postprocess(images, output_type=output_type)
+        self.maybe_free_model_hooks()
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

scheduler/pipeline_moss_image.py

@@ -0,0 +1,440 @@
+"""MOSSImage modeling and text-to-image pipeline (no Motor/TorchTitan imports).
+
+Dependencies: torch>=2.9, diffusers>=0.38, transformers>=5.3, safetensors,
+accelerate, Pillow. Copy this file anywhere and import it normally.
+
+    import torch
+    # Import MOSSImagePipeline from this file for component-level loading.
+    pipe = MOSSImagePipeline.from_components_pretrained(
+        transformer_path="exported_transformer",  # save_pretrained directory
+        text_encoder_path="Qwen/Qwen3.5-9B",
+        vae_path="black-forest-labs/FLUX.2-klein-4B",
+    ).to("cuda")
+    pipe("A cat", generator=torch.Generator("cuda").manual_seed(42)).images[0].save("cat.png")
+
+Default guidance uses CFG 3 and skips block 22 only in the negative branch
+(zero-based main DiT block indices; two forwards per sampling step):
+    pipe("A cat", guidance_scale=3.0, uncond_skip_layers=[22])
+    pipe("A cat", uncond_skip_layers=[])  # disable negative-branch skipping
+Optional additional positive-branch SLG adds a third forward:
+    pipe("A cat", slg_layers=[24], slg_scale=3.0)
+The additional SLG layer/scale above illustrates the API, not a quality-tuned setting.
+
+For a consolidated training state dict, create MOSSImageTransformer and
+call load_training_state_dict(state), then save_pretrained("exported_transformer").
+Select raw or EMA weights explicitly before calling it; distributed checkpoint
+shards/optimizer states are not model state dicts. Defaults match the finalized dense MOSSImage model,
+not MoE, CFG-distilled, video, or reference-image variants. Prompt strings are
+already-rendered captions (dataset-specific quality/domain tags are not added).
+
+First-release defaults: bf16 DiT/text-encoder weights, fp32 RoPE frequencies
+and phase/sin/cos calculation, then RoPE cast to the video activation dtype.
+Initial Gaussian noise is sampled directly in fp32, including with bf16 weights.
+CFG and Euler latent updates use fp32; VAE decoding defaults to fp32. Scheduler
+shift defaults to 5.0. Exact numerical parity with the original fp32-master
+inference pipeline is not a release requirement.
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from diffusers import AutoencoderKLFlux2, DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin  # Required by custom-component discovery.
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
+        return y.to(x.dtype) * self.weight
+
+
+def _situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))
+
+
+def _validate_skip_layers(layers, depth, name='skip_layers'):
+    layers = () if layers is None else tuple(layers)
+    if any(type(index) is not int or not 0 <= index < depth for index in layers):
+        raise ValueError(f'{name} must contain integer block indices in [0, {depth}).')
+    return frozenset(layers)
+
+
+def _rotate(x, rope):
+    sin, cos = rope.chunk(2, -1)
+    sin = torch.cat((sin, sin), -1).unsqueeze(-2)
+    cos = torch.cat((cos, cos), -1).unsqueeze(-2)
+    a, b = x.chunk(2, -1)
+    return x * cos + torch.cat((-b, a), -1) * sin
+
+
+def _attention(q, k, v, rope, heads):
+    shape = (*q.shape[:2], heads, q.shape[-1] // heads)
+    q = _rotate(q.reshape(shape), rope).transpose(1, 2)
+    k = _rotate(k.reshape(shape), rope).transpose(1, 2)
+    v = v.reshape(shape).transpose(1, 2)
+    return F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v).transpose(1, 2).flatten(2)
+
+
+class _TextBlock(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def forward(self, x, rope):
+        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj(_attention(q, k, v, rope, self.heads))
+        y = self.norm2(x)
+        return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class _TextProjection(nn.Module):
+    def __init__(self, text_dim, dim, ffn_dim, heads):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)])
+
+    def forward(self, x, rope):
+        x = self.proj(x.to(self.proj.weight.dtype))
+        for block in self.blocks:
+            x = block(x, rope)
+        return x
+
+
+class _Block(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.proj = nn.ModuleDict({
+            name: nn.Linear(dim, dim, bias=False) for name in ('q', 'k', 'v', 'o')
+        })
+        self.proj['ffn_up_gate'] = nn.Linear(dim, 2 * ffn_dim, bias=False)
+        self.proj['ffn_down'] = nn.Linear(ffn_dim, dim, bias=False)
+
+    def forward(self, x, rope, modulation):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
+        y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
+        dtype = self.proj['q'].weight.dtype
+        y = y.to(dtype)
+        q = self.q_norm(self.proj['q'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(self.proj['k'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj['o'](_attention(q, k, self.proj['v'](y), rope, self.heads)) * gate_a[:, None]
+        y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None]
+        hidden, gate = self.proj['ffn_up_gate'](y.to(dtype)).chunk(2, -1)
+        return x + self.proj['ffn_down'](_situ(hidden, gate)) * gate_f[:, None]
+
+
+class _RoPE(nn.Module):
+    def __init__(self, head_dim, base):
+        super().__init__()
+        self.num_bands = head_dim // 8
+        self.base = base
+        self.register_buffer('bands', self._default_bands(torch.device('cpu')))
+
+    def _default_bands(self, device):
+        # Match training's CPU fp32 construction, including under meta init.
+        exponent = torch.arange(self.num_bands, device='cpu', dtype=torch.float32) / self.num_bands
+        return (1 / (self.base ** exponent)).to(device)
+
+    def _apply(self, fn, recurse=True):
+        super()._apply(fn, recurse=recurse)
+        # .float() cannot recover frequencies already rounded by .to(bf16).
+        self.bands = self._default_bands(self.bands.device)
+        return self
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        super()._load_from_state_dict(*args, **kwargs)
+        # Frequencies are derived constants; retain the key for checkpoint compatibility.
+        self.bands = self._default_bands(self.bands.device)
+
+    def forward(self, height, width, text_length, device):
+        # Diffusers/Accelerate can install tensors directly, bypassing load hooks.
+        if self.bands.dtype != torch.float32:
+            self.bands = self._default_bands(device)
+        # Coordinates: (r, f, h, w); sequence order is image followed by text.
+        h, w = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing='ij')
+        image = torch.stack((torch.zeros_like(h), torch.full_like(h, text_length), h, w), -1).flatten(0, 1).float()
+        image[:, 2] -= (height - 1) / 2
+        image[:, 3] -= (width - 1) / 2
+        text = torch.zeros(text_length, 4, device=device, dtype=torch.float32)
+        text[:, 1] = torch.arange(text_length, device=device)
+        phase = torch.cat((image, text), 0).unsqueeze(-1) * self.bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2).unsqueeze(0)
+
+
+class MOSSImageTransformer(ModelMixin, ConfigMixin):
+    """MOSSImage DiT with training-compatible state-dict names and SDPA attention.
+
+    Each sample is evaluated at its true text length, including text refinement;
+    padded tokens cannot influence either text or joint attention.
+    """
+    @register_to_config
+    def __init__(self, dim=3072, ffn_dim=8192, num_heads=24, num_layers=42,
+                 text_dim=4096, freq_dim=256, latent_channels=128, rope_base=10000.0):
+        super().__init__()
+        if dim % num_heads or (dim // num_heads) % 8 or freq_dim % 2:
+            raise ValueError('Require dim divisible by heads, head_dim divisible by 8, and even freq_dim.')
+        self.video_embedder = nn.Linear(latent_channels, dim)
+        self.text_embedder = _TextProjection(text_dim, dim, ffn_dim, num_heads)
+        self.rope = _RoPE(dim // num_heads, rope_base)
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = nn.ModuleList([_Block(dim, ffn_dim, num_heads) for _ in range(num_layers)])
+        self.final_norm_video = _RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+
+    def load_training_state_dict(self, state_dict, *, prefix=''):
+        """Load a consolidated raw/EMA model mapping, strictly; no silent key drops."""
+        if prefix:
+            if any(not key.startswith(prefix) for key in state_dict):
+                raise ValueError('All keys must start with the explicit prefix.')
+            state_dict = {key[len(prefix):]: value for key, value in state_dict.items()}
+        return self.load_state_dict(state_dict, strict=True)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states,
+                encoder_attention_mask=None, return_dict=True, skip_layers=None):
+        skip_layers = _validate_skip_layers(skip_layers, len(self.layers))
+        if hidden_states.ndim != 4 or hidden_states.shape[1] != self.config.latent_channels:
+            raise ValueError('hidden_states must be [B, latent_channels, H, W].')
+        batch, _, height, width = hidden_states.shape
+        if encoder_hidden_states.ndim != 3 or encoder_hidden_states.shape[0] != batch:
+            raise ValueError('encoder_hidden_states must be [B, S, text_dim].')
+        mask = encoder_attention_mask
+        if mask is None:
+            mask = torch.ones(encoder_hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        if mask.shape != encoder_hidden_states.shape[:2] or not torch.all((mask == 0) | (mask == 1)):
+            raise ValueError('encoder_attention_mask must be a binary [B, S] mask.')
+        times = torch.as_tensor(timestep, device=hidden_states.device, dtype=torch.float32).reshape(-1)
+        times = times.expand(batch)
+        # Training computes the sinusoid in float64 and feeds float32 into its MLP.
+        half = self.config.freq_dim // 2
+        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)
+        phase = times.double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = self.time_embedding(embedding.to(self.time_embedding[0].weight.dtype)).float()
+        mod = self.shared_block_modulation(time.to(self.shared_block_modulation.weight.dtype)).view(batch, 6, -1)
+        head = self.video_head_modulation(time.to(self.video_head_modulation.weight.dtype)).view(batch, 2, -1)
+        outputs = []
+        for index in range(batch):
+            context = encoder_hidden_states[index:index + 1, mask[index].bool()]
+            if context.shape[1] == 0:
+                raise ValueError('Each sample needs at least one valid text token.')
+            video = hidden_states[index:index + 1].flatten(2).transpose(1, 2)
+            video = self.video_embedder(video.to(self.video_embedder.weight.dtype).contiguous())
+            # Match the reference activation dtype under fp32 weights + autocast.
+            rope = self.rope(height, width, context.shape[1], hidden_states.device).to(video.dtype)
+            text = self.text_embedder(context, rope[:, height * width:])
+            x = torch.cat((video, text), 1)
+            for layer_index, layer in enumerate(self.layers):
+                if layer_index not in skip_layers:
+                    x = layer(x, rope, mod[index:index + 1])
+            x = self.final_norm_video(x[:, :height * width].to(self.final_norm_video.weight.dtype))
+            x = x * (1 + head[index:index + 1, 1, None]) + head[index:index + 1, 0, None]
+            x = self.final_linear_video(x.to(self.final_linear_video.weight.dtype))
+            outputs.append(x.transpose(1, 2).reshape(1, -1, height, width))
+        sample = torch.cat(outputs)
+        return Transformer2DModelOutput(sample=sample) if return_dict else (sample,)
+
+
+class MOSSImageScheduler(FlowMatchEulerDiscreteScheduler):
+    """Euler integration using the training project's V2, nearest-grid schedule."""
+    @register_to_config
+    def __init__(self, num_train_timesteps=1000, shift=5.0, training_shift=1.0):
+        super().__init__(num_train_timesteps=num_train_timesteps, shift=1.0)
+        self.register_to_config(shift=shift, training_shift=training_shift)
+
+    def set_timesteps(self, num_inference_steps=50, device=None):
+        if num_inference_steps < 1 or self.config.shift <= 0 or self.config.training_shift <= 0:
+            raise ValueError('Steps and shifts must be positive.')
+        total = self.config.num_train_timesteps
+        def grid(count, shift):
+            lower = shift / (total + shift - 1)
+            values = torch.linspace(1, lower, count)
+            return shift * values / (1 + (shift - 1) * values)
+        sigmas = grid(num_inference_steps, self.config.shift)
+        training = grid(total, self.config.training_shift).flip(0)
+        upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+        lower = (upper - 1).clamp(min=0)
+        nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
+        sigmas = torch.unique_consecutive(training[nearest])
+        # Our sigmas have already been shifted and snapped. Avoid a second shift.
+        self._shift = 1.0
+        super().set_timesteps(device=device, sigmas=sigmas.tolist())
+
+
+class MOSSImagePipeline(DiffusionPipeline):
+    model_cpu_offload_seq = 'text_encoder->transformer->vae'
+    _optional_components = ['text_encoder', 'tokenizer', 'vae']
+
+    def __init__(self, transformer: MOSSImageTransformer,
+                 scheduler: MOSSImageScheduler, text_encoder=None, tokenizer=None, vae=None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler,
+                              text_encoder=text_encoder, tokenizer=tokenizer, vae=vae)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=16)
+
+    @classmethod
+    def from_components_pretrained(cls, transformer_path, text_encoder_path, vae_path,
+                                   torch_dtype=torch.bfloat16, vae_dtype=torch.float32,
+                                   local_files_only=False):
+        common = dict(local_files_only=local_files_only)
+        transformer = MOSSImageTransformer.from_pretrained(transformer_path, torch_dtype=torch_dtype, **common)
+        tokenizer = AutoTokenizer.from_pretrained(text_encoder_path, **common)
+        text_encoder = AutoModel.from_pretrained(text_encoder_path, dtype=torch_dtype, **common).eval()
+        subfolder = None if (Path(vae_path) / 'config.json').is_file() else 'vae'
+        vae = AutoencoderKLFlux2.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=vae_dtype, **common).eval()
+        return cls(transformer.eval(), MOSSImageScheduler(), text_encoder, tokenizer, vae)
+
+    @torch.no_grad()
+    def encode_prompt(self, prompt, max_sequence_length=4096):
+        if self.text_encoder is None or self.tokenizer is None:
+            raise ValueError('Text encoder/tokenizer required for string prompts.')
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        if not prompts or not all(isinstance(p, str) for p in prompts):
+            raise ValueError('prompt must be a string or a nonempty list of strings.')
+        rows = []
+        # Unpadded individual forwards also avoid padding effects in Qwen GDN.
+        for value in prompts:
+            formatted = self.tokenizer.apply_chat_template(
+                [{'role': 'user', 'content': value}], tokenize=False,
+                add_generation_prompt=True, enable_thinking=True)
+            inputs = self.tokenizer(formatted, return_tensors='pt', truncation=True,
+                                    max_length=max_sequence_length).to(self._execution_device)
+            output = self.text_encoder(**inputs, use_cache=False, return_dict=True)
+            rows.append(output.last_hidden_state[0])
+        length = max(row.shape[0] for row in rows)
+        embeds = rows[0].new_zeros(len(rows), length, rows[0].shape[-1])
+        mask = torch.zeros(len(rows), length, device=embeds.device, dtype=torch.bool)
+        for index, row in enumerate(rows):
+            embeds[index, :row.shape[0]] = row
+            mask[index, :row.shape[0]] = True
+        return embeds, mask
+
+    @torch.no_grad()
+    def __call__(self, prompt=None, *, height=1024, width=1024, num_inference_steps=50,
+                 guidance_scale=3.0, negative_prompt='', num_images_per_prompt=1,
+                 generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
+                 negative_prompt_embeds=None, negative_prompt_attention_mask=None,
+                 max_sequence_length=4096, output_type='pil', return_dict=True,
+                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 slg_start=0.0, slg_stop=1.0):
+        """Generate images with optional full-block skipping (zero-based indices).
+
+        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        prediction with a skipped-block prediction, retaining two forwards per
+        step. Pass None or an empty sequence to disable skipping. slg_layers
+        adds a skipped positive-prompt forward and the fp32 correction
+        slg_scale * (cond - cond_skip) to CFG. With guidance_scale=1, active
+        SLG uses cond_skip + slg_scale * (cond - cond_skip), matching Diffusers'
+        default SkipLayerGuidance formulation, with two forwards and no negative
+        branch. Both options can combine. slg_scale=1 then gives normal cond.
+        Both uncond skipping and SLG are active where
+        slg_start <= step_index / num_steps < slg_stop;
+        the interval follows denoising order, from initial noise to final image.
+        Outside it, inference uses normal CFG without skipping.
+        slg_scale=0 or empty slg_layers disables the additional forward.
+        """
+        uncond_skip_layers = _validate_skip_layers(
+            uncond_skip_layers, len(self.transformer.layers), 'uncond_skip_layers')
+        slg_layers = _validate_skip_layers(slg_layers, len(self.transformer.layers), 'slg_layers')
+        if not math.isfinite(slg_scale) or slg_scale < 0:
+            raise ValueError('slg_scale must be finite and nonnegative.')
+        if not 0 <= slg_start <= slg_stop <= 1:
+            raise ValueError('Require 0 <= slg_start <= slg_stop <= 1.')
+        if height <= 0 or width <= 0 or height % 16 or width % 16:
+            raise ValueError('height and width must be positive multiples of 16.')
+        if output_type not in {'pil', 'np', 'pt', 'latent'}:
+            raise ValueError('output_type must be pil, np, pt or latent.')
+        if num_images_per_prompt < 1 or not math.isfinite(guidance_scale) or guidance_scale < 0:
+            raise ValueError('Invalid image count or guidance scale.')
+        if not 1 <= max_sequence_length <= 4096:
+            raise ValueError('max_sequence_length must be between 1 and 4096.')
+        if (prompt is None) == (prompt_embeds is None):
+            raise ValueError('Provide exactly one of prompt and prompt_embeds.')
+        device = self._execution_device
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length)
+        batch = prompt_embeds.shape[0]
+        def prepare(embeds, mask):
+            if embeds.ndim != 3 or embeds.shape[0] != batch:
+                raise ValueError('Prompt embeddings must be [B, S, C] with matching batch size.')
+            if mask is None:
+                mask = torch.ones(embeds.shape[:2], dtype=torch.bool, device=embeds.device)
+            return (embeds.to(device).repeat_interleave(num_images_per_prompt, 0),
+                    mask.to(device).repeat_interleave(num_images_per_prompt, 0))
+        positive, positive_mask = prepare(prompt_embeds, prompt_attention_mask)
+        do_cfg = guidance_scale != 1.0
+        if do_cfg:
+            if negative_prompt_embeds is None:
+                negatives = [negative_prompt] * batch if isinstance(negative_prompt, str) else negative_prompt
+                if len(negatives) != batch:
+                    raise ValueError('negative_prompt must match prompt batch size.')
+                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(negatives, max_sequence_length)
+            negative, negative_mask = prepare(negative_prompt_embeds, negative_prompt_attention_mask)
+        shape = (batch * num_images_per_prompt, self.transformer.config.latent_channels, height // 16, width // 16)
+        if latents is None:
+            # Sample in fp32: promoting bf16 noise afterward cannot restore its precision.
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        elif tuple(latents.shape) != shape:
+            raise ValueError(f'latents must have shape {shape}.')
+        else:
+            latents = latents.to(device=device, dtype=torch.float32)
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        num_steps = len(self.scheduler.timesteps)
+        for step_index, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            skip_active = slg_start <= step_index / num_steps < slg_stop
+            # CFG must combine fp32 predictions, even when the DiT runs in bf16.
+            # Converting only the combined velocity would preserve rounding errors.
+            prediction = self.transformer(latents, timestep, positive, positive_mask).sample.float()
+            conditional = prediction
+            if do_cfg:
+                skip_kwargs = {'skip_layers': uncond_skip_layers} if uncond_skip_layers and skip_active else {}
+                uncond = self.transformer(latents, timestep, negative, negative_mask, **skip_kwargs).sample.float()
+                prediction = uncond + guidance_scale * (prediction - uncond)
+            if slg_layers and slg_scale > 0 and skip_active:
+                skipped = self.transformer(
+                    latents, timestep, positive, positive_mask, skip_layers=slg_layers).sample.float()
+                if not do_cfg:
+                    prediction = skipped
+                prediction = prediction + slg_scale * (conditional - skipped)
+            latents = self.scheduler.step(prediction, timestep, latents, return_dict=False)[0]
+        if output_type == 'latent':
+            images = latents
+        else:
+            if self.vae is None:
+                raise ValueError('A FLUX.2 VAE is required for decoded output.')
+            z = latents.to(dtype=self.vae.dtype)
+            bn = self.vae.bn
+            eps = self.vae.config.batch_norm_eps
+            z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1)
+            b, c, h, w = z.shape
+            z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+            images = self.vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+            images = self.image_processor.postprocess(images, output_type=output_type)
+        self.maybe_free_model_hooks()
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

scheduler/scheduler_config.json

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

text_encoder/config.json

@@ -0,0 +1,109 @@
+{
+  "architectures": [
+    "Qwen3_5Model"
+  ],
+  "dtype": "bfloat16",
+  "image_token_id": 248056,
+  "model_type": "qwen3_5",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 32,
+    "hidden_act": "silu",
+    "hidden_size": 32,
+    "initializer_range": 0.02,
+    "intermediate_size": 64,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 32,
+    "linear_num_key_heads": 1,
+    "linear_num_value_heads": 1,
+    "linear_value_head_dim": 32,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "mlp_only_layers": [],
+    "model_type": "qwen3_5_text",
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "num_attention_heads": 1,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 1,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        1,
+        1,
+        2
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 248320
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.15.0.dev0",
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 27,
+    "dtype": "bfloat16",
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 32,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 64,
+    "model_type": "qwen3_5_vision",
+    "num_heads": 1,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 32,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053
+}

text_encoder/model.safetensors

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

tokenizer/chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

tokenizer/tokenizer.json

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

tokenizer/tokenizer_config.json

@@ -0,0 +1,32 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "local_files_only": true,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}

transformer/config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "MOSSImageTransformer",
+  "_diffusers_version": "0.40.0.dev0",
+  "dim": 64,
+  "ffn_dim": 128,
+  "freq_dim": 32,
+  "latent_channels": 128,
+  "num_heads": 2,
+  "num_layers": 42,
+  "rope_base": 10000.0,
+  "text_dim": 32
+}

transformer/diffusion_pytorch_model.safetensors

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

transformer/pipeline_moss_image.py

@@ -0,0 +1,440 @@
+"""MOSSImage modeling and text-to-image pipeline (no Motor/TorchTitan imports).
+
+Dependencies: torch>=2.9, diffusers>=0.38, transformers>=5.3, safetensors,
+accelerate, Pillow. Copy this file anywhere and import it normally.
+
+    import torch
+    # Import MOSSImagePipeline from this file for component-level loading.
+    pipe = MOSSImagePipeline.from_components_pretrained(
+        transformer_path="exported_transformer",  # save_pretrained directory
+        text_encoder_path="Qwen/Qwen3.5-9B",
+        vae_path="black-forest-labs/FLUX.2-klein-4B",
+    ).to("cuda")
+    pipe("A cat", generator=torch.Generator("cuda").manual_seed(42)).images[0].save("cat.png")
+
+Default guidance uses CFG 3 and skips block 22 only in the negative branch
+(zero-based main DiT block indices; two forwards per sampling step):
+    pipe("A cat", guidance_scale=3.0, uncond_skip_layers=[22])
+    pipe("A cat", uncond_skip_layers=[])  # disable negative-branch skipping
+Optional additional positive-branch SLG adds a third forward:
+    pipe("A cat", slg_layers=[24], slg_scale=3.0)
+The additional SLG layer/scale above illustrates the API, not a quality-tuned setting.
+
+For a consolidated training state dict, create MOSSImageTransformer and
+call load_training_state_dict(state), then save_pretrained("exported_transformer").
+Select raw or EMA weights explicitly before calling it; distributed checkpoint
+shards/optimizer states are not model state dicts. Defaults match the finalized dense MOSSImage model,
+not MoE, CFG-distilled, video, or reference-image variants. Prompt strings are
+already-rendered captions (dataset-specific quality/domain tags are not added).
+
+First-release defaults: bf16 DiT/text-encoder weights, fp32 RoPE frequencies
+and phase/sin/cos calculation, then RoPE cast to the video activation dtype.
+Initial Gaussian noise is sampled directly in fp32, including with bf16 weights.
+CFG and Euler latent updates use fp32; VAE decoding defaults to fp32. Scheduler
+shift defaults to 5.0. Exact numerical parity with the original fp32-master
+inference pipeline is not a release requirement.
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from diffusers import AutoencoderKLFlux2, DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin  # Required by custom-component discovery.
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
+        return y.to(x.dtype) * self.weight
+
+
+def _situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))
+
+
+def _validate_skip_layers(layers, depth, name='skip_layers'):
+    layers = () if layers is None else tuple(layers)
+    if any(type(index) is not int or not 0 <= index < depth for index in layers):
+        raise ValueError(f'{name} must contain integer block indices in [0, {depth}).')
+    return frozenset(layers)
+
+
+def _rotate(x, rope):
+    sin, cos = rope.chunk(2, -1)
+    sin = torch.cat((sin, sin), -1).unsqueeze(-2)
+    cos = torch.cat((cos, cos), -1).unsqueeze(-2)
+    a, b = x.chunk(2, -1)
+    return x * cos + torch.cat((-b, a), -1) * sin
+
+
+def _attention(q, k, v, rope, heads):
+    shape = (*q.shape[:2], heads, q.shape[-1] // heads)
+    q = _rotate(q.reshape(shape), rope).transpose(1, 2)
+    k = _rotate(k.reshape(shape), rope).transpose(1, 2)
+    v = v.reshape(shape).transpose(1, 2)
+    return F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v).transpose(1, 2).flatten(2)
+
+
+class _TextBlock(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def forward(self, x, rope):
+        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj(_attention(q, k, v, rope, self.heads))
+        y = self.norm2(x)
+        return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class _TextProjection(nn.Module):
+    def __init__(self, text_dim, dim, ffn_dim, heads):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)])
+
+    def forward(self, x, rope):
+        x = self.proj(x.to(self.proj.weight.dtype))
+        for block in self.blocks:
+            x = block(x, rope)
+        return x
+
+
+class _Block(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.proj = nn.ModuleDict({
+            name: nn.Linear(dim, dim, bias=False) for name in ('q', 'k', 'v', 'o')
+        })
+        self.proj['ffn_up_gate'] = nn.Linear(dim, 2 * ffn_dim, bias=False)
+        self.proj['ffn_down'] = nn.Linear(ffn_dim, dim, bias=False)
+
+    def forward(self, x, rope, modulation):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
+        y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
+        dtype = self.proj['q'].weight.dtype
+        y = y.to(dtype)
+        q = self.q_norm(self.proj['q'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(self.proj['k'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj['o'](_attention(q, k, self.proj['v'](y), rope, self.heads)) * gate_a[:, None]
+        y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None]
+        hidden, gate = self.proj['ffn_up_gate'](y.to(dtype)).chunk(2, -1)
+        return x + self.proj['ffn_down'](_situ(hidden, gate)) * gate_f[:, None]
+
+
+class _RoPE(nn.Module):
+    def __init__(self, head_dim, base):
+        super().__init__()
+        self.num_bands = head_dim // 8
+        self.base = base
+        self.register_buffer('bands', self._default_bands(torch.device('cpu')))
+
+    def _default_bands(self, device):
+        # Match training's CPU fp32 construction, including under meta init.
+        exponent = torch.arange(self.num_bands, device='cpu', dtype=torch.float32) / self.num_bands
+        return (1 / (self.base ** exponent)).to(device)
+
+    def _apply(self, fn, recurse=True):
+        super()._apply(fn, recurse=recurse)
+        # .float() cannot recover frequencies already rounded by .to(bf16).
+        self.bands = self._default_bands(self.bands.device)
+        return self
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        super()._load_from_state_dict(*args, **kwargs)
+        # Frequencies are derived constants; retain the key for checkpoint compatibility.
+        self.bands = self._default_bands(self.bands.device)
+
+    def forward(self, height, width, text_length, device):
+        # Diffusers/Accelerate can install tensors directly, bypassing load hooks.
+        if self.bands.dtype != torch.float32:
+            self.bands = self._default_bands(device)
+        # Coordinates: (r, f, h, w); sequence order is image followed by text.
+        h, w = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing='ij')
+        image = torch.stack((torch.zeros_like(h), torch.full_like(h, text_length), h, w), -1).flatten(0, 1).float()
+        image[:, 2] -= (height - 1) / 2
+        image[:, 3] -= (width - 1) / 2
+        text = torch.zeros(text_length, 4, device=device, dtype=torch.float32)
+        text[:, 1] = torch.arange(text_length, device=device)
+        phase = torch.cat((image, text), 0).unsqueeze(-1) * self.bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2).unsqueeze(0)
+
+
+class MOSSImageTransformer(ModelMixin, ConfigMixin):
+    """MOSSImage DiT with training-compatible state-dict names and SDPA attention.
+
+    Each sample is evaluated at its true text length, including text refinement;
+    padded tokens cannot influence either text or joint attention.
+    """
+    @register_to_config
+    def __init__(self, dim=3072, ffn_dim=8192, num_heads=24, num_layers=42,
+                 text_dim=4096, freq_dim=256, latent_channels=128, rope_base=10000.0):
+        super().__init__()
+        if dim % num_heads or (dim // num_heads) % 8 or freq_dim % 2:
+            raise ValueError('Require dim divisible by heads, head_dim divisible by 8, and even freq_dim.')
+        self.video_embedder = nn.Linear(latent_channels, dim)
+        self.text_embedder = _TextProjection(text_dim, dim, ffn_dim, num_heads)
+        self.rope = _RoPE(dim // num_heads, rope_base)
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = nn.ModuleList([_Block(dim, ffn_dim, num_heads) for _ in range(num_layers)])
+        self.final_norm_video = _RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+
+    def load_training_state_dict(self, state_dict, *, prefix=''):
+        """Load a consolidated raw/EMA model mapping, strictly; no silent key drops."""
+        if prefix:
+            if any(not key.startswith(prefix) for key in state_dict):
+                raise ValueError('All keys must start with the explicit prefix.')
+            state_dict = {key[len(prefix):]: value for key, value in state_dict.items()}
+        return self.load_state_dict(state_dict, strict=True)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states,
+                encoder_attention_mask=None, return_dict=True, skip_layers=None):
+        skip_layers = _validate_skip_layers(skip_layers, len(self.layers))
+        if hidden_states.ndim != 4 or hidden_states.shape[1] != self.config.latent_channels:
+            raise ValueError('hidden_states must be [B, latent_channels, H, W].')
+        batch, _, height, width = hidden_states.shape
+        if encoder_hidden_states.ndim != 3 or encoder_hidden_states.shape[0] != batch:
+            raise ValueError('encoder_hidden_states must be [B, S, text_dim].')
+        mask = encoder_attention_mask
+        if mask is None:
+            mask = torch.ones(encoder_hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        if mask.shape != encoder_hidden_states.shape[:2] or not torch.all((mask == 0) | (mask == 1)):
+            raise ValueError('encoder_attention_mask must be a binary [B, S] mask.')
+        times = torch.as_tensor(timestep, device=hidden_states.device, dtype=torch.float32).reshape(-1)
+        times = times.expand(batch)
+        # Training computes the sinusoid in float64 and feeds float32 into its MLP.
+        half = self.config.freq_dim // 2
+        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)
+        phase = times.double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = self.time_embedding(embedding.to(self.time_embedding[0].weight.dtype)).float()
+        mod = self.shared_block_modulation(time.to(self.shared_block_modulation.weight.dtype)).view(batch, 6, -1)
+        head = self.video_head_modulation(time.to(self.video_head_modulation.weight.dtype)).view(batch, 2, -1)
+        outputs = []
+        for index in range(batch):
+            context = encoder_hidden_states[index:index + 1, mask[index].bool()]
+            if context.shape[1] == 0:
+                raise ValueError('Each sample needs at least one valid text token.')
+            video = hidden_states[index:index + 1].flatten(2).transpose(1, 2)
+            video = self.video_embedder(video.to(self.video_embedder.weight.dtype).contiguous())
+            # Match the reference activation dtype under fp32 weights + autocast.
+            rope = self.rope(height, width, context.shape[1], hidden_states.device).to(video.dtype)
+            text = self.text_embedder(context, rope[:, height * width:])
+            x = torch.cat((video, text), 1)
+            for layer_index, layer in enumerate(self.layers):
+                if layer_index not in skip_layers:
+                    x = layer(x, rope, mod[index:index + 1])
+            x = self.final_norm_video(x[:, :height * width].to(self.final_norm_video.weight.dtype))
+            x = x * (1 + head[index:index + 1, 1, None]) + head[index:index + 1, 0, None]
+            x = self.final_linear_video(x.to(self.final_linear_video.weight.dtype))
+            outputs.append(x.transpose(1, 2).reshape(1, -1, height, width))
+        sample = torch.cat(outputs)
+        return Transformer2DModelOutput(sample=sample) if return_dict else (sample,)
+
+
+class MOSSImageScheduler(FlowMatchEulerDiscreteScheduler):
+    """Euler integration using the training project's V2, nearest-grid schedule."""
+    @register_to_config
+    def __init__(self, num_train_timesteps=1000, shift=5.0, training_shift=1.0):
+        super().__init__(num_train_timesteps=num_train_timesteps, shift=1.0)
+        self.register_to_config(shift=shift, training_shift=training_shift)
+
+    def set_timesteps(self, num_inference_steps=50, device=None):
+        if num_inference_steps < 1 or self.config.shift <= 0 or self.config.training_shift <= 0:
+            raise ValueError('Steps and shifts must be positive.')
+        total = self.config.num_train_timesteps
+        def grid(count, shift):
+            lower = shift / (total + shift - 1)
+            values = torch.linspace(1, lower, count)
+            return shift * values / (1 + (shift - 1) * values)
+        sigmas = grid(num_inference_steps, self.config.shift)
+        training = grid(total, self.config.training_shift).flip(0)
+        upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+        lower = (upper - 1).clamp(min=0)
+        nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
+        sigmas = torch.unique_consecutive(training[nearest])
+        # Our sigmas have already been shifted and snapped. Avoid a second shift.
+        self._shift = 1.0
+        super().set_timesteps(device=device, sigmas=sigmas.tolist())
+
+
+class MOSSImagePipeline(DiffusionPipeline):
+    model_cpu_offload_seq = 'text_encoder->transformer->vae'
+    _optional_components = ['text_encoder', 'tokenizer', 'vae']
+
+    def __init__(self, transformer: MOSSImageTransformer,
+                 scheduler: MOSSImageScheduler, text_encoder=None, tokenizer=None, vae=None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler,
+                              text_encoder=text_encoder, tokenizer=tokenizer, vae=vae)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=16)
+
+    @classmethod
+    def from_components_pretrained(cls, transformer_path, text_encoder_path, vae_path,
+                                   torch_dtype=torch.bfloat16, vae_dtype=torch.float32,
+                                   local_files_only=False):
+        common = dict(local_files_only=local_files_only)
+        transformer = MOSSImageTransformer.from_pretrained(transformer_path, torch_dtype=torch_dtype, **common)
+        tokenizer = AutoTokenizer.from_pretrained(text_encoder_path, **common)
+        text_encoder = AutoModel.from_pretrained(text_encoder_path, dtype=torch_dtype, **common).eval()
+        subfolder = None if (Path(vae_path) / 'config.json').is_file() else 'vae'
+        vae = AutoencoderKLFlux2.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=vae_dtype, **common).eval()
+        return cls(transformer.eval(), MOSSImageScheduler(), text_encoder, tokenizer, vae)
+
+    @torch.no_grad()
+    def encode_prompt(self, prompt, max_sequence_length=4096):
+        if self.text_encoder is None or self.tokenizer is None:
+            raise ValueError('Text encoder/tokenizer required for string prompts.')
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        if not prompts or not all(isinstance(p, str) for p in prompts):
+            raise ValueError('prompt must be a string or a nonempty list of strings.')
+        rows = []
+        # Unpadded individual forwards also avoid padding effects in Qwen GDN.
+        for value in prompts:
+            formatted = self.tokenizer.apply_chat_template(
+                [{'role': 'user', 'content': value}], tokenize=False,
+                add_generation_prompt=True, enable_thinking=True)
+            inputs = self.tokenizer(formatted, return_tensors='pt', truncation=True,
+                                    max_length=max_sequence_length).to(self._execution_device)
+            output = self.text_encoder(**inputs, use_cache=False, return_dict=True)
+            rows.append(output.last_hidden_state[0])
+        length = max(row.shape[0] for row in rows)
+        embeds = rows[0].new_zeros(len(rows), length, rows[0].shape[-1])
+        mask = torch.zeros(len(rows), length, device=embeds.device, dtype=torch.bool)
+        for index, row in enumerate(rows):
+            embeds[index, :row.shape[0]] = row
+            mask[index, :row.shape[0]] = True
+        return embeds, mask
+
+    @torch.no_grad()
+    def __call__(self, prompt=None, *, height=1024, width=1024, num_inference_steps=50,
+                 guidance_scale=3.0, negative_prompt='', num_images_per_prompt=1,
+                 generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
+                 negative_prompt_embeds=None, negative_prompt_attention_mask=None,
+                 max_sequence_length=4096, output_type='pil', return_dict=True,
+                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 slg_start=0.0, slg_stop=1.0):
+        """Generate images with optional full-block skipping (zero-based indices).
+
+        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        prediction with a skipped-block prediction, retaining two forwards per
+        step. Pass None or an empty sequence to disable skipping. slg_layers
+        adds a skipped positive-prompt forward and the fp32 correction
+        slg_scale * (cond - cond_skip) to CFG. With guidance_scale=1, active
+        SLG uses cond_skip + slg_scale * (cond - cond_skip), matching Diffusers'
+        default SkipLayerGuidance formulation, with two forwards and no negative
+        branch. Both options can combine. slg_scale=1 then gives normal cond.
+        Both uncond skipping and SLG are active where
+        slg_start <= step_index / num_steps < slg_stop;
+        the interval follows denoising order, from initial noise to final image.
+        Outside it, inference uses normal CFG without skipping.
+        slg_scale=0 or empty slg_layers disables the additional forward.
+        """
+        uncond_skip_layers = _validate_skip_layers(
+            uncond_skip_layers, len(self.transformer.layers), 'uncond_skip_layers')
+        slg_layers = _validate_skip_layers(slg_layers, len(self.transformer.layers), 'slg_layers')
+        if not math.isfinite(slg_scale) or slg_scale < 0:
+            raise ValueError('slg_scale must be finite and nonnegative.')
+        if not 0 <= slg_start <= slg_stop <= 1:
+            raise ValueError('Require 0 <= slg_start <= slg_stop <= 1.')
+        if height <= 0 or width <= 0 or height % 16 or width % 16:
+            raise ValueError('height and width must be positive multiples of 16.')
+        if output_type not in {'pil', 'np', 'pt', 'latent'}:
+            raise ValueError('output_type must be pil, np, pt or latent.')
+        if num_images_per_prompt < 1 or not math.isfinite(guidance_scale) or guidance_scale < 0:
+            raise ValueError('Invalid image count or guidance scale.')
+        if not 1 <= max_sequence_length <= 4096:
+            raise ValueError('max_sequence_length must be between 1 and 4096.')
+        if (prompt is None) == (prompt_embeds is None):
+            raise ValueError('Provide exactly one of prompt and prompt_embeds.')
+        device = self._execution_device
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length)
+        batch = prompt_embeds.shape[0]
+        def prepare(embeds, mask):
+            if embeds.ndim != 3 or embeds.shape[0] != batch:
+                raise ValueError('Prompt embeddings must be [B, S, C] with matching batch size.')
+            if mask is None:
+                mask = torch.ones(embeds.shape[:2], dtype=torch.bool, device=embeds.device)
+            return (embeds.to(device).repeat_interleave(num_images_per_prompt, 0),
+                    mask.to(device).repeat_interleave(num_images_per_prompt, 0))
+        positive, positive_mask = prepare(prompt_embeds, prompt_attention_mask)
+        do_cfg = guidance_scale != 1.0
+        if do_cfg:
+            if negative_prompt_embeds is None:
+                negatives = [negative_prompt] * batch if isinstance(negative_prompt, str) else negative_prompt
+                if len(negatives) != batch:
+                    raise ValueError('negative_prompt must match prompt batch size.')
+                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(negatives, max_sequence_length)
+            negative, negative_mask = prepare(negative_prompt_embeds, negative_prompt_attention_mask)
+        shape = (batch * num_images_per_prompt, self.transformer.config.latent_channels, height // 16, width // 16)
+        if latents is None:
+            # Sample in fp32: promoting bf16 noise afterward cannot restore its precision.
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        elif tuple(latents.shape) != shape:
+            raise ValueError(f'latents must have shape {shape}.')
+        else:
+            latents = latents.to(device=device, dtype=torch.float32)
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        num_steps = len(self.scheduler.timesteps)
+        for step_index, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            skip_active = slg_start <= step_index / num_steps < slg_stop
+            # CFG must combine fp32 predictions, even when the DiT runs in bf16.
+            # Converting only the combined velocity would preserve rounding errors.
+            prediction = self.transformer(latents, timestep, positive, positive_mask).sample.float()
+            conditional = prediction
+            if do_cfg:
+                skip_kwargs = {'skip_layers': uncond_skip_layers} if uncond_skip_layers and skip_active else {}
+                uncond = self.transformer(latents, timestep, negative, negative_mask, **skip_kwargs).sample.float()
+                prediction = uncond + guidance_scale * (prediction - uncond)
+            if slg_layers and slg_scale > 0 and skip_active:
+                skipped = self.transformer(
+                    latents, timestep, positive, positive_mask, skip_layers=slg_layers).sample.float()
+                if not do_cfg:
+                    prediction = skipped
+                prediction = prediction + slg_scale * (conditional - skipped)
+            latents = self.scheduler.step(prediction, timestep, latents, return_dict=False)[0]
+        if output_type == 'latent':
+            images = latents
+        else:
+            if self.vae is None:
+                raise ValueError('A FLUX.2 VAE is required for decoded output.')
+            z = latents.to(dtype=self.vae.dtype)
+            bn = self.vae.bn
+            eps = self.vae.config.batch_norm_eps
+            z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1)
+            b, c, h, w = z.shape
+            z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+            images = self.vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+            images = self.image_processor.postprocess(images, output_type=output_type)
+        self.maybe_free_model_hooks()
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

vae/config.json

@@ -0,0 +1,45 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.40.0.dev0",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    32,
+    32,
+    64,
+    64
+  ],
+  "decoder_block_out_channels": [
+    32,
+    32,
+    64,
+    64
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 32,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 32,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

vae/diffusion_pytorch_model.safetensors

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