diff --git a/python/sglang/multimodal_gen/configs/models/dits/bria_fibo2.py b/python/sglang/multimodal_gen/configs/models/dits/bria_fibo2.py deleted file mode 100644 index 598561433c1a..000000000000 --- a/python/sglang/multimodal_gen/configs/models/dits/bria_fibo2.py +++ /dev/null @@ -1,100 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Architecture configuration for the BRIA Fibo 2.0 diffusion transformer. - -Fibo 2.0 keeps the Z-Image transformer block and adds three things on top of it: - -* a Perceiver resampler that squeezes the text encoder features into a variable - number of "gist" tokens, -* five injection blocks that also cross-attend to their own text bundle, -* per-token (noisy/clean) adaLN modulation, which is how editing conditions the - generation on the clean latents of the images to edit. - -The parameter names emitted here are the checkpoint's own diffusers names -(``attention.to_q``, ``attention.to_out.0``, ``cross_attn_gate``, -``perceiver.latent_queries``, ...). The runtime model mirrors them one to one, so -``param_names_mapping`` stays empty and no weight reshuffle happens at load time. -In particular q/k/v are kept separate: a cross-attention block reads a -``cap_feat_dim`` (7680) wide text bundle while its query comes from a ``dim`` -(3584) wide stream, so the two cannot be fused into one stacked projection. -""" - -from dataclasses import dataclass, field -from typing import Any, Tuple - -from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig - - -@dataclass -class BriaFibo2ArchConfig(DiTArchConfig): - # 2x2 latent patches become one token, features ordered (channel, dy, dx) - all_patch_size: Tuple[int, ...] = (2,) - patch_size: int = 2 - in_channels: int = 32 - out_channels: int | None = None - - dim: int = 3584 - num_layers: int = 44 - n_refiner_layers: int = 2 - num_attention_heads: int = 28 - norm_eps: float = 1e-5 - - # Text conditioning: 3 Qwen3-VL layers of 2560 channels concatenated - cap_feat_dim: int = 7680 - # Blocks that cross-attend to their own text bundle. Bundle 0 feeds the - # Perceiver, bundle k + 1 feeds injection_layer_ids[k]. - injection_layer_ids: Tuple[int, ...] = (6, 14, 22, 30, 38) - perceiver_num_layers: int = 2 - - # Gist token count: interpolated from the text length, quantized to a step - min_num_gist_tokens: int = 256 - max_num_gist_tokens: int = 512 - gist_step: int = 64 - gist_min_text_len: int = 256 - gist_max_text_len: int = 2048 - - rope_theta: float = 256.0 - t_scale: float = 1000.0 - # The diffusion stages hand the transformer the scheduler's own timesteps - # (0..num_train_timesteps), while the reference transformer takes the noise - # level in [0, 1] and multiplies by `t_scale` itself. Dividing by the - # scheduler's unit count bridges the two conventions without touching the - # embedding domain the checkpoint was trained in. - scheduler_timesteps: float = 1000.0 - # Three axes (plane, row, column); the sum is the attention head dim - axes_dims: Tuple[int, int, int] = (32, 48, 48) - axes_lens: Tuple[int, int, int] = (2048, 512, 512) - - # Fibo keeps q/k/v separate (see module docstring), so nothing is stacked - stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) - - # diffusers names are used as-is by the runtime model - param_names_mapping: dict = field(default_factory=dict) - - def __post_init__(self) -> None: - super().__post_init__() - self.out_channels = self.out_channels or self.in_channels - self.num_channels_latents = self.in_channels - self.hidden_size = self.dim - self.num_attention_heads = self.num_attention_heads - if sum(self.axes_dims) != self.dim // self.num_attention_heads: - raise ValueError( - "axes_dims must sum to the attention head dim: " - f"sum({self.axes_dims}) != {self.dim} // {self.num_attention_heads}" - ) - - -@dataclass -class BriaFibo2DitConfig(DiTConfig): - arch_config: BriaFibo2ArchConfig = field(default_factory=BriaFibo2ArchConfig) - - prefix: str = "bria_fibo2" - - # One block serves the gist tokens (hundreds) and the image plus context - # tokens (thousands); the gist branch is the short one, so its GEMMs should - # not be autotuned for the long shape alone. - torch_compile_inductor_config: dict[str, Any] = field( - default_factory=lambda: {"multi_kernel_hints": [512, 4096]} - ) - - -__all__ = ["BriaFibo2ArchConfig", "BriaFibo2DitConfig"] diff --git a/python/sglang/multimodal_gen/configs/models/encoders/bria_fibo2.py b/python/sglang/multimodal_gen/configs/models/encoders/bria_fibo2.py deleted file mode 100644 index 8822d8df0bc4..000000000000 --- a/python/sglang/multimodal_gen/configs/models/encoders/bria_fibo2.py +++ /dev/null @@ -1,250 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Qwen3-VL text encoder configuration for BRIA Fibo 2.0. - -Fibo 2.0 conditions its transformer on six "caption bundles", each made of three -Qwen3-VL hidden states concatenated along the channel axis after a per-token -float32 RMS normalization. Bundle 0 feeds the Perceiver resampler, bundles 1..5 -feed the injection blocks in order. - -Two things here are easy to get silently wrong, so both are asserted: - -* **Layer index semantics.** The reference indexes the ``output_hidden_states`` - tuple the way ``transformers`` builds it: index 0 is the embedding output and - index ``i`` is the input of decoder layer ``i``, for ``n_layers + 1`` entries. - The in-tree Qwen3-VL text tower instead appends *after* each layer and never - appends the final norm, so it yields ``n_layers`` entries where entry ``i`` is - the output of layer ``i``. Hence ``sglang[i] == hf[i + 1]`` and the reference - indices must be shifted down by one. ``SGLANG_HIDDEN_STATE_OFFSET`` records - that, and the length assertion turns any upstream change into an error instead - of wrong weights. -* **Architecture defaults.** The inherited defaults describe Qwen3-VL-8B, not the - 4B variant Fibo ships. ``post_diffusers_config_update`` flattens the real - per-component values out of ``text_config`` / ``vision_config`` rather than - relying on any default. -""" - -from dataclasses import dataclass, field -from typing import Any, Sequence - -import torch - -from sglang.multimodal_gen.configs.models.encoders.base import ( - TextEncoderArchConfig, - TextEncoderConfig, -) -from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLArchConfig -from sglang.multimodal_gen.configs.models.fsdp import ( - is_block, - is_embed_tokens, - is_final_norm, - is_layer, -) - -# Qwen3-VL layers behind the transformer's six text inputs, in `transformers` -# index space: the Perceiver's bundle first, then one bundle per injection block -# (injection_layer_ids = (6, 14, 22, 30, 38)). -TEXT_ENCODER_OUT_LAYERS: tuple[tuple[int, ...], ...] = ( - (9, 20, 31), - (5, 16, 27), - (7, 18, 29), - (9, 20, 31), - (11, 22, 33), - (13, 24, 35), -) - -# The in-tree text tower drains decoder-layer outputs only, so its tuple is one -# entry shorter and every reference index maps one lower. Never drop this shift. -SGLANG_HIDDEN_STATE_OFFSET = 1 - -PROMPT_TEMPLATE_ENCODE = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" - -VISION_MARKER = "<|vision_start|><|image_pad|><|vision_end|>" - -CAP_FEAT_DIM = 7680 - - -def unique_hf_layers( - out_layers: Sequence[Sequence[int]] = TEXT_ENCODER_OUT_LAYERS, -) -> list[int]: - """The distinct reference layer indices, sorted.""" - return sorted({layer for bundle in out_layers for layer in bundle}) - - -def sglang_layer_indices( - out_layers: Sequence[Sequence[int]] = TEXT_ENCODER_OUT_LAYERS, -) -> tuple[tuple[int, ...], ...]: - """Reference layer indices translated into the in-tree tower's tuple.""" - return tuple( - tuple(layer - SGLANG_HIDDEN_STATE_OFFSET for layer in bundle) - for bundle in out_layers - ) - - -def build_caption_bundles( - hidden_states: Sequence[torch.Tensor], - num_hidden_layers: int, - dtype: torch.dtype | None = None, - out_layers: Sequence[Sequence[int]] = TEXT_ENCODER_OUT_LAYERS, -) -> torch.Tensor: - """Stack the six caption bundles from the text tower's hidden states. - - Args: - hidden_states: the text tower's ``output_hidden_states`` tuple, in the - in-tree convention (outputs of every decoder layer, no embedding or - final-norm entry), length ``num_hidden_layers``. - num_hidden_layers: layer count the tower was built with, used to check - the tuple length so an upstream collection change fails loudly. - dtype: dtype the normalized states are cast back to; defaults to the - dtype of the input states. - - Returns: - ``(batch, 6, text_len, 3 * hidden_size)``, i.e. ``(B, 6, T, 7680)`` for - the released checkpoints. - """ - if len(hidden_states) != num_hidden_layers: - raise ValueError( - "BriaFibo2 expects the text tower to return one hidden state per " - f"decoder layer ({num_hidden_layers}), got {len(hidden_states)}. " - "Changing how Qwen3-VL collects `all_hidden_states` changes the " - f"layer mapping (currently offset {SGLANG_HIDDEN_STATE_OFFSET})." - ) - - target_dtype = dtype if dtype is not None else hidden_states[0].dtype - normalized: dict[int, torch.Tensor] = {} - for hf_layer in unique_hf_layers(out_layers): - index = hf_layer - SGLANG_HIDDEN_STATE_OFFSET - states = hidden_states[index].float() - # unit RMS per token, matching the reference - normalized[hf_layer] = ( - states * torch.rsqrt(states.pow(2).mean(dim=-1, keepdim=True) + 1e-6) - ).to(target_dtype) - - bundles = [ - torch.cat([normalized[layer] for layer in bundle], dim=-1) - for bundle in out_layers - ] - return torch.stack(bundles, dim=1) - - -def _read(container: Any, key: str, default: Any = None) -> Any: - """Read a field from either a config object or a plain dict.""" - if container is None: - return default - if isinstance(container, dict): - return container.get(key, default) - return getattr(container, key, default) - - -@dataclass -class BriaFibo2Qwen3VLArchConfig(Qwen3VLArchConfig): - """Qwen3-VL-4B as Fibo 2.0 uses it. - - The inherited values describe the 8B encoder used by JoyImage; everything - that matters is overwritten from the checkpoint config in - ``post_diffusers_config_update``, so these defaults only have to be harmless. - """ - - # Fibo reads three hidden states per bundle, six bundles in total - out_layers: tuple[tuple[int, ...], ...] = TEXT_ENCODER_OUT_LAYERS - cap_feat_dim: int = CAP_FEAT_DIM - - # Reference prompt handling - prompt_template_encode: str = PROMPT_TEMPLATE_ENCODE - max_sequence_length: int = 4096 - max_num_images: int = 5 - text_encoder_image_size: int = 768 - text_len: int = 4096 - - # The Fibo encoder has no pooled output and needs every layer - output_hidden_states: bool = True - - _fsdp_shard_conditions: list = field( - default_factory=lambda: [is_layer, is_block, is_embed_tokens, is_final_norm] - ) - - def __post_init__(self) -> None: - # Deliberately not calling super().__post_init__(): `Qwen3VLArchConfig` - # installs JoyImage's tokenizer kwargs there, whose `max_length` is - # `text_len + prompt_template_encode_start_idx["image"]` (34 extra - # tokens). Fibo's template is a plain chat turn with no such offset, so - # the inherited value would silently mis-size every prompt. - self.tokenizer_kwargs = { - "padding": True, - "truncation": True, - "max_length": self.max_sequence_length, - "return_tensors": "pt", - } - - -@dataclass -class BriaFibo2Qwen3VLConfig(TextEncoderConfig): - arch_config: TextEncoderArchConfig = field( - default_factory=BriaFibo2Qwen3VLArchConfig - ) - - def post_diffusers_config_update(self) -> None: - """Flatten the real encoder shape out of the checkpoint config. - - `update_model_arch` copies only top-level keys, so the nested - `text_config` / `vision_config` never reach the flat fields and the 8B - defaults would survive. Reading them here keeps the 4B variant correct - without hardcoding a layer count. - """ - arch = self.arch_config - - text_config = _read(arch, "text_config") - vision_config = _read(arch, "vision_config") - - arch.hidden_size = int(_read(text_config, "hidden_size", arch.hidden_size)) - arch.intermediate_size = int( - _read(text_config, "intermediate_size", arch.intermediate_size) - ) - arch.num_hidden_layers = int( - _read(text_config, "num_hidden_layers", arch.num_hidden_layers) - ) - arch.num_attention_heads = int( - _read(text_config, "num_attention_heads", arch.num_attention_heads) - ) - arch.num_key_value_heads = int( - _read(text_config, "num_key_value_heads", arch.num_key_value_heads or 0) - or arch.num_key_value_heads - ) - arch.head_dim = int(_read(text_config, "head_dim", arch.head_dim or 0) or 0) - arch.vocab_size = int(_read(text_config, "vocab_size", arch.vocab_size)) - arch.rms_norm_eps = float(_read(text_config, "rms_norm_eps", arch.rms_norm_eps)) - arch.tie_word_embeddings = bool( - _read(text_config, "tie_word_embeddings", arch.tie_word_embeddings) - ) - rope_theta = _read(text_config, "rope_theta") - if rope_theta is None: - rope_theta = _read(_read(text_config, "rope_parameters"), "rope_theta") - if rope_theta is not None: - arch.rope_theta = float(rope_theta) - - if vision_config is not None: - arch.vision_config = vision_config - - if arch.head_dim in (None, 0): - arch.head_dim = arch.hidden_size // arch.num_attention_heads - - # A three-layer bundle times the encoder width must match cap_feat_dim, - # which is also the cross-attention key/value width of the transformer - expected = 3 * arch.hidden_size - if expected != arch.cap_feat_dim: - raise ValueError( - "cap_feat_dim must be 3 text-encoder hidden sizes: " - f"3 * {arch.hidden_size} != {arch.cap_feat_dim}" - ) - - -__all__ = [ - "BriaFibo2Qwen3VLArchConfig", - "BriaFibo2Qwen3VLConfig", - "PROMPT_TEMPLATE_ENCODE", - "SGLANG_HIDDEN_STATE_OFFSET", - "TEXT_ENCODER_OUT_LAYERS", - "VISION_MARKER", - "build_caption_bundles", - "sglang_layer_indices", - "unique_hf_layers", -] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/bria_fibo2.py b/python/sglang/multimodal_gen/runtime/models/dits/bria_fibo2.py deleted file mode 100644 index c9895d19d970..000000000000 --- a/python/sglang/multimodal_gen/runtime/models/dits/bria_fibo2.py +++ /dev/null @@ -1,813 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""BRIA Fibo 2.0 diffusion transformer. - -A single-stream DiT in the Z-Image family, extended in three ways: - -* a Perceiver resampler turns the first text bundle into ``num_gist_tokens`` - gist tokens, a count derived per prompt from the text length, -* the blocks listed in ``injection_layer_ids`` also cross-attend to their own - text bundle, -* adaLN modulation is per token: tokens marked noisy by ``noise_mask`` take the - modulation of the real timestep, the others (gist and context tokens, which - are clean) take the modulation of t = 0. That is how editing conditions a - generation on the clean latents of the images to edit. - -The whole image plus every context image share one sequence and one patch -embedding. Their positions come from a three-axis RoPE, so the token order only -affects float rounding, not the result. - -Parameter names are the checkpoint's diffusers names, so weights load without a -remap. q/k/v and w1/w3 are deliberately kept as separate projections: a -cross-attention reads a ``cap_feat_dim`` wide text bundle with a ``dim`` wide -query, so those cannot share a stacked projection. - -This file targets the single-GPU path. Tensor/sequence parallelism, an -attention-backend switch (the reference uses ``USPAttention`` elsewhere in the -tree) and the batched variable-gist path are follow-ups; the batched case is -rejected explicitly rather than half-supported. -""" - -from __future__ import annotations - -import math -from typing import Any, Optional, Sequence, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from sglang.multimodal_gen.configs.models.dits.bria_fibo2 import BriaFibo2DitConfig -from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( - QuantizationConfig, -) -from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.models.dits.zimage import RopeEmbedder -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) - -ADALN_EMBED_DIM = 256 - - -def compute_gist_count( - text_lengths: Sequence[int], - min_gist: int, - max_gist: int, - gist_step: int, - min_text_len: int, - max_text_len: int, -) -> list[int]: - """Number of gist tokens per prompt, from its text length in tokens. - - Lengths between ``min_text_len`` and ``max_text_len`` map linearly onto - ``[min_gist, max_gist]``, rounded up to a multiple of ``gist_step``. - """ - span = max_text_len - min_text_len - counts = [] - for n in text_lengths: - ratio = min(1.0, max(0.0, (float(n) - min_text_len) / span)) - raw = min_gist + ratio * (max_gist - min_gist) - quantized = int(math.ceil(raw / gist_step)) * gist_step - counts.append(max(min_gist, min(max_gist, quantized))) - return counts - - -class BriaFibo2RMSNorm(nn.Module): - """RMSNorm matching the reference implementation. - - The variance is taken in float32. ``zimage.ZImageRMSNorm`` deliberately keeps - Z-Image's native bf16 behaviour instead, so the two are not interchangeable. - """ - - def __init__(self, dim: int, eps: float = 1e-5) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(dim)) - self.variance_epsilon = eps - - def forward(self, x: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.to(torch.float32) - variance = x_float.pow(2).mean(dim=-1, keepdim=True) - normalized = x_float * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * normalized.to(input_dtype) - - -def _mask_to_bool(mask: torch.Tensor) -> torch.Tensor: - """SDPA wants a bool (or float) mask; the stages carry it as long/bool.""" - if mask.dtype == torch.bool: - return mask - if mask.is_floating_point(): - return mask - return mask != 0 - - -def normalize_encoder_attention_mask( - mask: Any, device: torch.device -) -> Optional[torch.Tensor]: - """Normalize the text mask to one ``(batch, text_len)`` tensor. - - The diffusion stages thread prompt masks **per sample** (`Req.prompt_attention_mask` - is a list), which is convenient for the encoders but not for the transformer; - accept either form so the caller does not have to convert. - """ - if mask is None: - return None - if isinstance(mask, torch.Tensor): - mask = mask if mask.device == device else mask.to(device) - return _mask_to_bool(mask) - if isinstance(mask, (list, tuple)): - items = [m for m in mask if m is not None] - if not items or len(items) != len(mask): - return None - tensors = [ - _mask_to_bool( - (m if isinstance(m, torch.Tensor) else torch.as_tensor(m)).to(device) - ) - for m in items - ] - if len(tensors) == 1: - single = tensors[0] - return single if single.dim() == 2 else single.reshape(1, -1) - if all(t.dim() == 2 and t.shape[0] == 1 for t in tensors): - return torch.cat(tensors, dim=0) - return torch.stack(tensors, dim=0) - return None - - -def select_per_token( - value_noisy: torch.Tensor, - value_clean: torch.Tensor, - noise_mask: torch.Tensor, - seq_len: int, -) -> torch.Tensor: - """Broadcast the noisy or the clean value onto each token of the sequence.""" - noise_mask_expanded = noise_mask.unsqueeze(-1) # (batch, seq_len, 1) - return torch.where( - noise_mask_expanded == 1, - value_noisy.unsqueeze(1).expand(-1, seq_len, -1), - value_clean.unsqueeze(1).expand(-1, seq_len, -1), - ) - - -def apply_rotary_emb( - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> torch.Tensor: - """Rotate adjacent channel pairs, the complex multiply of the reference. - - ``x`` is ``[batch, seq_len, heads, head_dim]``; ``cos`` and ``sin`` are - ``[1, seq_len, head_dim // 2]``. The rotation is computed in float32 and cast - back, as the reference does. - """ - dtype = x.dtype - x_float = x.float() - x_even = x_float[..., 0::2] - x_odd = x_float[..., 1::2] - cos = cos.unsqueeze(2).float() - sin = sin.unsqueeze(2).float() - out_even = x_even * cos - x_odd * sin - out_odd = x_even * sin + x_odd * cos - return torch.stack((out_even, out_odd), dim=-1).flatten(-2).to(dtype) - - -class BriaFibo2TimestepEmbedder(nn.Module): - def __init__( - self, - out_size: int, - mid_size: Optional[int] = None, - frequency_embedding_size: int = 256, - ) -> None: - super().__init__() - if mid_size is None: - mid_size = out_size - self.mlp = nn.Sequential( - nn.Linear(frequency_embedding_size, mid_size, bias=True), - nn.SiLU(), - nn.Linear(mid_size, out_size, bias=True), - ) - self.frequency_embedding_size = frequency_embedding_size - - @staticmethod - def timestep_embedding( - t: torch.Tensor, dim: int, max_period: int = 10000 - ) -> torch.Tensor: - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) - * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) - / half - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat( - [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 - ) - return embedding - - def forward(self, t: torch.Tensor) -> torch.Tensor: - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) - return self.mlp(t_freq.to(self.mlp[0].weight.dtype)) - - -class BriaFibo2FeedForward(nn.Module): - def __init__(self, dim: int, hidden_dim: int) -> None: - super().__init__() - self.w1 = nn.Linear(dim, hidden_dim, bias=False) - self.w2 = nn.Linear(hidden_dim, dim, bias=False) - self.w3 = nn.Linear(dim, hidden_dim, bias=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.w2(F.silu(self.w1(x)) * self.w3(x)) - - -class BriaFibo2Attention(nn.Module): - """Self-attention (with RoPE) or cross-attention to ``encoder_hidden_states``.""" - - def __init__( - self, - query_dim: int, - heads: int, - dim_head: int, - kv_dim: Optional[int] = None, - eps: float = 1e-5, - ) -> None: - super().__init__() - self.heads = heads - self.head_dim = dim_head - kv_dim = kv_dim or query_dim - self.to_q = nn.Linear(query_dim, heads * dim_head, bias=False) - self.to_k = nn.Linear(kv_dim, heads * dim_head, bias=False) - self.to_v = nn.Linear(kv_dim, heads * dim_head, bias=False) - # The reference keeps the output projection in a ModuleList, which the - # checkpoint reflects as `to_out.0` - self.to_out = nn.ModuleList( - [nn.Linear(heads * dim_head, query_dim, bias=False)] - ) - self.norm_q = BriaFibo2RMSNorm(dim_head, eps=eps) - self.norm_k = BriaFibo2RMSNorm(dim_head, eps=eps) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor: - batch_size, seq_len = hidden_states.shape[0], hidden_states.shape[1] - query = self.to_q(hidden_states) - kv_source = ( - hidden_states if encoder_hidden_states is None else encoder_hidden_states - ) - key = self.to_k(kv_source) - value = self.to_v(kv_source) - - query = query.view(batch_size, seq_len, self.heads, self.head_dim) - key = key.view(batch_size, -1, self.heads, self.head_dim) - value = value.view(batch_size, -1, self.heads, self.head_dim) - - query = self.norm_q(query) - key = self.norm_k(key) - - if freqs_cis is not None: - cos, sin = freqs_cis - query = apply_rotary_emb(query, cos, sin) - key = apply_rotary_emb(key, cos, sin) - - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - - # A 2D mask marks the real keys, so it broadcasts over query and heads - if attention_mask is not None and attention_mask.ndim == 2: - attention_mask = attention_mask[:, None, None, :] - - out = F.scaled_dot_product_attention( - query, - key, - value, - attn_mask=attention_mask, - dropout_p=0.0, - is_causal=False, - ) - out = out.transpose(1, 2).reshape(batch_size, seq_len, -1) - return self.to_out[0](out) - - -class BriaFibo2TransformerBlock(nn.Module): - """A Z-Image block, plus optional text cross-attention and per-token modulation.""" - - def __init__( - self, - dim: int, - n_heads: int, - norm_eps: float, - cross_attention_dim: Optional[int] = None, - ) -> None: - super().__init__() - self.attention = BriaFibo2Attention(dim, n_heads, dim // n_heads, eps=norm_eps) - self.feed_forward = BriaFibo2FeedForward(dim, (dim * 8 // 3) // 64 * 64) - - self.attention_norm1 = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.ffn_norm1 = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.attention_norm2 = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.ffn_norm2 = BriaFibo2RMSNorm(dim, eps=norm_eps) - - self.adaLN_modulation = nn.Sequential( - nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True) - ) - - if cross_attention_dim is not None: - self.cross_attn_norm1 = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.cross_attention = BriaFibo2Attention( - dim, n_heads, dim // n_heads, kv_dim=cross_attention_dim, eps=norm_eps - ) - self.cross_attn_norm2 = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.cross_attn_gate = nn.Parameter(torch.zeros(1)) - else: - self.cross_attention = None - - def forward( - self, - x: torch.Tensor, - attn_mask: Optional[torch.Tensor], - freqs_cis: Tuple[torch.Tensor, torch.Tensor], - adaln_input: Optional[torch.Tensor] = None, - noise_mask: Optional[torch.Tensor] = None, - adaln_noisy: Optional[torch.Tensor] = None, - adaln_clean: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - seq_len = x.shape[1] - - if noise_mask is not None: - # Per-token modulation: noisy and clean tokens get different values - mod_noisy = self.adaLN_modulation(adaln_noisy) - mod_clean = self.adaLN_modulation(adaln_clean) - ( - scale_msa_noisy, - gate_msa_noisy, - scale_mlp_noisy, - gate_mlp_noisy, - ) = mod_noisy.chunk(4, dim=1) - ( - scale_msa_clean, - gate_msa_clean, - scale_mlp_clean, - gate_mlp_clean, - ) = mod_clean.chunk(4, dim=1) - - gate_msa_noisy, gate_mlp_noisy = ( - gate_msa_noisy.tanh(), - gate_mlp_noisy.tanh(), - ) - gate_msa_clean, gate_mlp_clean = ( - gate_msa_clean.tanh(), - gate_mlp_clean.tanh(), - ) - scale_msa_noisy, scale_mlp_noisy = ( - 1.0 + scale_msa_noisy, - 1.0 + scale_mlp_noisy, - ) - scale_msa_clean, scale_mlp_clean = ( - 1.0 + scale_msa_clean, - 1.0 + scale_mlp_clean, - ) - - scale_msa = select_per_token( - scale_msa_noisy, scale_msa_clean, noise_mask, seq_len - ) - scale_mlp = select_per_token( - scale_mlp_noisy, scale_mlp_clean, noise_mask, seq_len - ) - gate_msa = select_per_token( - gate_msa_noisy, gate_msa_clean, noise_mask, seq_len - ) - gate_mlp = select_per_token( - gate_mlp_noisy, gate_mlp_clean, noise_mask, seq_len - ) - else: - # Global modulation: one value broadcast over the whole sequence - mod = self.adaLN_modulation(adaln_input) - scale_msa, gate_msa, scale_mlp, gate_mlp = mod.unsqueeze(1).chunk(4, dim=2) - gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() - scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp - - attn_out = self.attention( - self.attention_norm1(x) * scale_msa, - attention_mask=attn_mask, - freqs_cis=freqs_cis, - ) - x = x + gate_msa * self.attention_norm2(attn_out) - - if self.cross_attention is not None: - cross_out = self.cross_attention( - self.cross_attn_norm1(x), - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - ) - x = x + self.cross_attn_gate.tanh() * self.cross_attn_norm2(cross_out) - - x = x + gate_mlp * self.ffn_norm2( - self.feed_forward(self.ffn_norm1(x) * scale_mlp) - ) - return x - - -class BriaFibo2PerceiverLayer(nn.Module): - """Gist tokens read the text bundle, then each other, then an FFN.""" - - def __init__( - self, dim: int, n_heads: int, cross_attention_dim: int, norm_eps: float - ) -> None: - super().__init__() - self.cross_attn_norm = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.cross_attention = BriaFibo2Attention( - dim, n_heads, dim // n_heads, kv_dim=cross_attention_dim, eps=norm_eps - ) - self.self_attn_norm = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.self_attention = BriaFibo2Attention( - dim, n_heads, dim // n_heads, eps=norm_eps - ) - self.ffn_norm = BriaFibo2RMSNorm(dim, eps=norm_eps) - self.ffn = BriaFibo2FeedForward(dim, (dim * 8 // 3) // 64 * 64) - - def forward( - self, - latents: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - latents = latents + self.cross_attention( - self.cross_attn_norm(latents), - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - ) - latents = latents + self.self_attention( - self.self_attn_norm(latents), attention_mask=attention_mask - ) - latents = latents + self.ffn(self.ffn_norm(latents)) - return latents - - -class BriaFibo2PerceiverResampler(nn.Module): - """Squeezes the text bundle into ``num_gist_tokens`` gist tokens.""" - - def __init__( - self, - dim: int, - n_heads: int, - cross_attention_dim: int, - num_queries: int, - num_layers: int, - norm_eps: float, - ) -> None: - super().__init__() - self.latent_queries = nn.Parameter(torch.randn(num_queries, dim) * 0.02) - self.layers = nn.ModuleList( - [ - BriaFibo2PerceiverLayer(dim, n_heads, cross_attention_dim, norm_eps) - for _ in range(num_layers) - ] - ) - self.norm_out = BriaFibo2RMSNorm(dim, eps=norm_eps) - - def forward( - self, - encoder_hidden_states: torch.Tensor, - num_gist_tokens: int, - encoder_attention_mask: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - batch_size = encoder_hidden_states.shape[0] - latents = ( - self.latent_queries[:num_gist_tokens] - .unsqueeze(0) - .expand(batch_size, -1, -1) - ) - for layer in self.layers: - latents = layer( - latents, encoder_hidden_states, encoder_attention_mask, attention_mask - ) - return self.norm_out(latents) - - -class BriaFibo2FinalLayer(nn.Module): - def __init__(self, hidden_size: int, out_channels: int) -> None: - super().__init__() - self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.linear = nn.Linear(hidden_size, out_channels, bias=True) - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), - nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), - ) - - def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: - scale = 1.0 + self.adaLN_modulation(c) - return self.linear(self.norm_final(x) * scale.unsqueeze(1)) - - -class BriaFibo2Transformer2DModel(CachableDiT): - param_names_mapping = BriaFibo2DitConfig().arch_config.param_names_mapping - reverse_param_names_mapping = ( - BriaFibo2DitConfig().arch_config.reverse_param_names_mapping - ) - # Single-GPU implementation: nothing is sharded, so FSDP has no layer groups - _fsdp_shard_conditions: list = [] - _compile_conditions: list = [] - packed_modules_mapping: dict = {} - - def __init__( - self, - config: BriaFibo2DitConfig, - hf_config: dict[str, Any], - quant_config: Optional[QuantizationConfig] = None, - ) -> None: - super().__init__(config=config, hf_config=hf_config) - - arch_config = self.config - self.dim = arch_config.dim - self.in_channels = arch_config.in_channels - self.out_channels = arch_config.out_channels - self.patch_size = arch_config.patch_size - self.n_heads = arch_config.num_attention_heads - self.hidden_size = arch_config.dim - self.num_attention_heads = arch_config.num_attention_heads - self.num_channels_latents = arch_config.in_channels - self.gradient_checkpointing = False - - self.x_embedder = nn.Linear( - self.patch_size * self.patch_size * self.in_channels, self.dim, bias=True - ) - self.t_embedder = BriaFibo2TimestepEmbedder( - min(self.dim, ADALN_EMBED_DIM), mid_size=1024 - ) - self.rope_embedder = RopeEmbedder( - theta=arch_config.rope_theta, - axes_dims=list(arch_config.axes_dims), - axes_lens=list(arch_config.axes_lens), - ) - - self.noise_refiner = nn.ModuleList( - [ - BriaFibo2TransformerBlock(self.dim, self.n_heads, arch_config.norm_eps) - for _ in range(arch_config.n_refiner_layers) - ] - ) - self.perceiver = BriaFibo2PerceiverResampler( - self.dim, - self.n_heads, - arch_config.cap_feat_dim, - arch_config.max_num_gist_tokens, - arch_config.perceiver_num_layers, - arch_config.norm_eps, - ) - injection_layer_ids = set(arch_config.injection_layer_ids) - self.layers = nn.ModuleList( - [ - BriaFibo2TransformerBlock( - self.dim, - self.n_heads, - arch_config.norm_eps, - cross_attention_dim=( - arch_config.cap_feat_dim - if layer_id in injection_layer_ids - else None - ), - ) - for layer_id in range(arch_config.num_layers) - ] - ) - self.final_layer = BriaFibo2FinalLayer( - self.dim, self.patch_size * self.patch_size * self.out_channels - ) - - @staticmethod - def create_coordinate_grid( - size: Sequence[int], - start: Optional[Sequence[int]] = None, - device: Optional[torch.device] = None, - ) -> torch.Tensor: - if start is None: - start = (0 for _ in size) - axes = [ - torch.arange(x0, x0 + span, dtype=torch.int32, device=device) - for x0, span in zip(start, size) - ] - grids = torch.meshgrid(axes, indexing="ij") - return torch.stack(grids, dim=-1) - - def _patchify(self, latents: torch.Tensor) -> torch.Tensor: - """2x2 latent patches become tokens, features ordered (channel, dy, dx).""" - batch_size, channels, height, width = latents.shape - patch_size = self.patch_size - latents = latents.view( - batch_size, - channels, - height // patch_size, - patch_size, - width // patch_size, - patch_size, - ) - return latents.permute(0, 2, 4, 1, 3, 5).reshape( - batch_size, (height // patch_size) * (width // patch_size), -1 - ) - - def _unpatchify( - self, tokens: torch.Tensor, height: int, width: int - ) -> torch.Tensor: - batch_size = tokens.shape[0] - patch_size = self.patch_size - height_tokens, width_tokens = height // patch_size, width // patch_size - tokens = tokens.view( - batch_size, - height_tokens, - width_tokens, - self.out_channels, - patch_size, - patch_size, - ) - return tokens.permute(0, 3, 1, 4, 2, 5).reshape( - batch_size, self.out_channels, height, width - ) - - @staticmethod - def _unsqueeze_freqs( - freqs: Tuple[torch.Tensor, torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - cos, sin = freqs - return cos.unsqueeze(0), sin.unsqueeze(0) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, - guidance: Any = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - source_latents: Optional[Sequence[torch.Tensor]] = None, - **kwargs: Any, - ) -> torch.Tensor: - """ - Args: - hidden_states: clean-image latents being denoised, ``(B, in_channels, H, W)``. - encoder_hidden_states: ``(B, 1 + len(injection_layer_ids), text_len, cap_feat_dim)``. - Bundle 0 feeds the Perceiver, bundle k + 1 feeds injection block k. - timestep: noise level in ``[0, 1]``, 1 is pure noise. - encoder_attention_mask: bool mask ``(B, text_len)`` of the real text tokens. - source_latents: normalized latents of the images to edit, one - ``(B, in_channels, h, w)`` tensor per image, in order. - """ - arch_config = self.config - encoder_attention_mask = normalize_encoder_attention_mask( - encoder_attention_mask, hidden_states.device - ) - batch_size, _, height, width = hidden_states.shape - patch_size = self.patch_size - height_tokens, width_tokens = height // patch_size, width // patch_size - num_image_tokens = height_tokens * width_tokens - device = hidden_states.device - - # 1. Timestep embeddings. Image tokens take the real timestep, the clean - # tokens (gist and context) take t = 0. `timestep` arrives in the - # scheduler's units, so rescale it into the [0, t_scale] domain the - # checkpoint's embedder was trained on. - timestep = timestep.to(hidden_states.dtype) * ( - arch_config.t_scale / arch_config.scheduler_timesteps - ) - adaln_noisy = self.t_embedder(timestep) - adaln_clean = self.t_embedder(torch.zeros_like(timestep)) - - # 2. Patchify the image and the context images, which share the patch - # embedding. The image sits on RoPE plane `max_num_gist_tokens + 1` - # whatever the gist count is; context image k sits on the k-th plane - # after it with its own rows and columns. - image_plane = arch_config.max_num_gist_tokens + 1 - tokens = [self._patchify(hidden_states)] - ids = [ - self.create_coordinate_grid( - (1, height_tokens, width_tokens), (image_plane, 0, 0), device - ) - ] - for k, latents in enumerate(source_latents or []): - tokens.append(self._patchify(latents)) - context_size = ( - 1, - latents.shape[2] // patch_size, - latents.shape[3] // patch_size, - ) - ids.append( - self.create_coordinate_grid( - context_size, (image_plane + k + 1, 0, 0), device - ) - ) - hidden_states = self.x_embedder(torch.cat(tokens, dim=1)) - image_freqs = self._unsqueeze_freqs( - self.rope_embedder(torch.cat([grid.flatten(0, 2) for grid in ids])) - ) - - # 3. Noise refiner. With context images, only the image tokens are noisy. - refiner_noise_mask = None - if source_latents: - refiner_noise_mask = torch.zeros( - batch_size, hidden_states.shape[1], dtype=torch.long, device=device - ) - refiner_noise_mask[:, :num_image_tokens] = 1 - for block in self.noise_refiner: - hidden_states = block( - hidden_states, - None, - image_freqs, - adaln_input=adaln_noisy, - noise_mask=refiner_noise_mask, - adaln_noisy=adaln_noisy, - adaln_clean=adaln_clean, - ) - - # 4. Gist count per prompt from its text length, then the Perceiver turns - # the first bundle into gist tokens. - if encoder_attention_mask is not None: - text_lengths = encoder_attention_mask.sum(dim=1).tolist() - else: - text_lengths = [encoder_hidden_states.shape[2]] * batch_size - gist_counts = compute_gist_count( - text_lengths, - min_gist=arch_config.min_num_gist_tokens, - max_gist=arch_config.max_num_gist_tokens, - gist_step=arch_config.gist_step, - min_text_len=arch_config.gist_min_text_len, - max_text_len=arch_config.gist_max_text_len, - ) - if min(gist_counts) != max(gist_counts): - # The per-prompt gist padding path reorders the sequence and needs a - # varlen-aware attention backend. Reject it instead of silently - # producing wrong tokens. - raise ValueError( - "BriaFibo2 does not support a batch whose prompts need different " - f"numbers of gist tokens yet (got {sorted(set(gist_counts))}). " - "Send these prompts as separate requests." - ) - num_gist_tokens = gist_counts[0] - gist = self.perceiver( - encoder_hidden_states[:, 0], num_gist_tokens, encoder_attention_mask - ) - gist_ids = self.create_coordinate_grid( - (num_gist_tokens, 1, 1), (1, 0, 0), device - ).flatten(0, 2) - - # 5. One stream: gist tokens first, then the image, then the context, - # the order the model was trained with. Positions come from RoPE, so the - # order only affects float rounding. - gist_freqs = self._unsqueeze_freqs(self.rope_embedder(gist_ids)) - image_start = num_gist_tokens - hidden_states = torch.cat([gist, hidden_states], dim=1) - freqs_cis = ( - torch.cat([gist_freqs[0], image_freqs[0]], dim=1), - torch.cat([gist_freqs[1], image_freqs[1]], dim=1), - ) - noise_mask = torch.zeros( - batch_size, hidden_states.shape[1], dtype=torch.long, device=device - ) - noise_mask[:, image_start : image_start + num_image_tokens] = 1 - - # 6. The main blocks. Injection blocks read their own bundle. - bundle_index = { - layer_id: index + 1 - for index, layer_id in enumerate(arch_config.injection_layer_ids) - } - for layer_id, block in enumerate(self.layers): - text = ( - encoder_hidden_states[:, bundle_index[layer_id]] - if layer_id in bundle_index - else None - ) - hidden_states = block( - hidden_states, - None, - freqs_cis, - noise_mask=noise_mask, - adaln_noisy=adaln_noisy, - adaln_clean=adaln_clean, - encoder_hidden_states=text, - encoder_attention_mask=encoder_attention_mask, - ) - - # 7. Final layer on the image tokens only, then unpatchify. - hidden_states = self.final_layer( - hidden_states[:, image_start : image_start + num_image_tokens], adaln_noisy - ) - return self._unpatchify(hidden_states, height, width) - - -EntryClass = BriaFibo2Transformer2DModel - -__all__ = [ - "BriaFibo2Attention", - "BriaFibo2PerceiverResampler", - "BriaFibo2Transformer2DModel", - "BriaFibo2TransformerBlock", - "compute_gist_count", - "select_per_token", -] diff --git a/python/sglang/multimodal_gen/test/unit/test_bria_fibo2.py b/python/sglang/multimodal_gen/test/unit/test_bria_fibo2.py deleted file mode 100644 index 7e007fab2fbc..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_bria_fibo2.py +++ /dev/null @@ -1,507 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""CPU-only tests for the BRIA Fibo 2.0 support. - -These pin down the semantics that silently produce wrong images when they drift: -the gist count mapping, the patch feature order, the three-axis RoPE (including -the fixed image plane), the per-token noisy/clean modulation, and the parameter -names/shapes the checkpoint ships. -""" - -import os - -import pytest -import torch - -from sglang.multimodal_gen.configs.models.dits.bria_fibo2 import ( - BriaFibo2ArchConfig, - BriaFibo2DitConfig, -) -from sglang.multimodal_gen.runtime.models.dits.bria_fibo2 import ( - ADALN_EMBED_DIM, - BriaFibo2Attention, - BriaFibo2Transformer2DModel, - apply_rotary_emb, - compute_gist_count, - select_per_token, -) - -# The real model is 8B; building it on the meta device costs no memory -META_CONFIG_DIR = os.environ.get( - "FIBO2_MODEL_CFG_DIR", "/workspace/dsh/fibo2/model_cfg" -) - - -def tiny_arch_config(**overrides) -> BriaFibo2ArchConfig: - kwargs = dict( - dim=64, - num_attention_heads=2, - num_layers=3, - n_refiner_layers=1, - in_channels=32, - patch_size=2, - cap_feat_dim=48, - injection_layer_ids=(1,), - perceiver_num_layers=1, - max_num_gist_tokens=8, - min_num_gist_tokens=4, - gist_step=2, - gist_min_text_len=2, - gist_max_text_len=8, - # head_dim is dim // n_heads = 32, and the axes must sum to it - axes_dims=(8, 12, 12), - axes_lens=(64, 16, 16), - # The tiny models are driven with noise levels in [0, 1], the reference - # convention, so no rescaling happens on the way to the embedder. - scheduler_timesteps=1.0, - ) - kwargs.update(overrides) - return BriaFibo2ArchConfig(**kwargs) - - -def tiny_model() -> BriaFibo2Transformer2DModel: - config = BriaFibo2DitConfig(arch_config=tiny_arch_config()) - model = BriaFibo2Transformer2DModel(config, hf_config={}) - return model.eval() - - -# --------------------------------------------------------------------------- -# T1 gist count -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "text_len,expected", - [ - (0, 256), - (1, 256), - (256, 256), - (500, 320), - (1000, 384), - (2048, 512), - (100000, 512), - ], -) -def test_gist_count_matches_reference_boundaries(text_len, expected): - got = compute_gist_count( - [text_len], - min_gist=256, - max_gist=512, - gist_step=64, - min_text_len=256, - max_text_len=2048, - ) - assert got == [expected] - - -def test_gist_count_is_always_quantized_and_clamped(): - lengths = list(range(0, 2600, 7)) - counts = compute_gist_count(lengths, 256, 512, 64, 256, 2048) - for n, c in zip(lengths, counts): - assert 256 <= c <= 512, (n, c) - assert c % 64 == 0, (n, c) - - -# --------------------------------------------------------------------------- -# T2 patchify / unpatchify -# --------------------------------------------------------------------------- - - -def test_patchify_unpatchify_roundtrip(): - model = tiny_model() - latents = torch.randn(2, 32, 16, 12) - tokens = model._patchify(latents) - assert tokens.shape == (2, (16 // 2) * (12 // 2), 32 * 2 * 2) - restored = model._unpatchify(tokens, 16, 12) - assert torch.equal(restored, latents) - - -def test_patchify_feature_order_is_channel_major(): - """Token features are ordered (channel, dy, dx), as the checkpoint expects.""" - model = tiny_model() - latents = torch.arange(2 * 2 * 2, dtype=torch.float32).reshape(1, 2, 2, 2) - tokens = model._patchify(latents) - # Single token; features = [c0 (0,0), c0 (0,1), c0 (1,0), c0 (1,1), c1 ...] - assert tokens.shape == (1, 1, 2 * 2 * 2) - expected = torch.tensor([[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]]) - assert torch.equal(tokens[0], expected) - - -# --------------------------------------------------------------------------- -# T3 RoPE equals the reference complex multiply -# --------------------------------------------------------------------------- - - -def test_apply_rotary_emb_matches_complex_reference(): - torch.manual_seed(0) - seq_len, heads, head_dim = 7, 2, 8 - x = torch.randn(1, seq_len, heads, head_dim) - half = head_dim // 2 - angles = torch.randn(1, seq_len, half) - cos, sin = torch.cos(angles), torch.sin(angles) - - got = apply_rotary_emb(x, cos, sin) - - # The reference does view_as_complex(x.reshape(..., -1, 2)) * polar(1, angle) - ref_complex = torch.complex(x.float()[..., 0::2], x.float()[..., 1::2]) - rotation = torch.polar(torch.ones_like(angles), angles).unsqueeze(2) - expected = torch.view_as_real(ref_complex * rotation).flatten(-2) - - assert torch.allclose(got.float(), expected, atol=1e-6) - - -def test_apply_rotary_emb_preserves_dtype_and_norm(): - x = torch.randn(1, 5, 2, 8, dtype=torch.float32) - angles = torch.randn(1, 5, 4) - out = apply_rotary_emb(x, torch.cos(angles), torch.sin(angles)) - assert out.dtype == x.dtype - assert torch.allclose(out.norm(dim=-1), x.norm(dim=-1), atol=1e-5), ( - "a rotation must preserve the per-head norm" - ) - - -# --------------------------------------------------------------------------- -# T4 RoPE planes: gist on 1..N, image always on max_num_gist_tokens + 1 -# --------------------------------------------------------------------------- - - -def _captured_rope_ids(model, **forward_kwargs): - """Run one forward and return the position ids handed to the RoPE table.""" - seen = [] - original = model.rope_embedder - - def recorder(ids): - seen.append(ids.clone()) - return original(ids) - - model.rope_embedder = recorder - try: - with torch.no_grad(): - model(**forward_kwargs) - finally: - model.rope_embedder = original - return seen - - -def test_rope_image_plane_is_fixed_at_max_gist_plus_one(): - """The image plane must be max_num_gist_tokens + 1 whatever the gist count.""" - model = tiny_model() # max_num_gist_tokens = 8 -> image plane 9 - batch, text_len, height, width = 1, 4, 16, 8 - hidden_states = torch.randn(batch, 32, height, width) - encoder_hidden_states = torch.randn(batch, 2, text_len, 48) - attention_mask = torch.ones(batch, text_len, dtype=torch.bool) - - ids = _captured_rope_ids( - model, - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=torch.full((batch,), 0.5), - encoder_attention_mask=attention_mask, - ) - # two calls: one for the patchified image+context block, one for the gist tokens - assert len(ids) == 2 - image_ids, gist_ids = ids - - image_plane = int(model.config.max_num_gist_tokens) + 1 - assert image_plane == 9 - assert image_ids[:, 0].unique().tolist() == [image_plane] - # rows and columns of the image patch grid - assert image_ids[:, 1].max().item() == height // 2 - 1 - assert image_ids[:, 2].max().item() == width // 2 - 1 - - # gist tokens walk along axis 0 starting at plane 1, rows/cols stay zero - assert gist_ids[:, 0].tolist() == list(range(1, gist_ids.shape[0] + 1)) - assert gist_ids[:, 1].abs().max().item() == 0 - assert gist_ids[:, 2].abs().max().item() == 0 - - -def test_rope_context_image_planes_follow_the_image_plane(): - model = tiny_model() - batch, text_len, height, width = 1, 4, 16, 8 - context_latents = [torch.randn(batch, 32, 8, 8), torch.randn(batch, 32, 16, 4)] - - ids = _captured_rope_ids( - model, - hidden_states=torch.randn(batch, 32, height, width), - encoder_hidden_states=torch.randn(batch, 2, text_len, 48), - timestep=torch.full((batch,), 0.5), - encoder_attention_mask=torch.ones(batch, text_len, dtype=torch.bool), - source_latents=context_latents, - ) - image_ids = ids[0] - image_plane = int(model.config.max_num_gist_tokens) + 1 - num_image_tokens = (height // 2) * (width // 2) - - assert image_ids[:num_image_tokens, 0].unique().tolist() == [image_plane] - # context image k sits on plane image_plane + k + 1 - first_ctx = image_ids[num_image_tokens : num_image_tokens + (8 // 2) * (8 // 2)] - assert first_ctx[:, 0].unique().tolist() == [image_plane + 1] - second_ctx = image_ids[num_image_tokens + (8 // 2) * (8 // 2) :] - assert second_ctx[:, 0].unique().tolist() == [image_plane + 2] - - -# --------------------------------------------------------------------------- -# T5 per-token modulation -# --------------------------------------------------------------------------- - - -def test_select_per_token_picks_noisy_and_clean_values(): - value_noisy = torch.tensor([[1.0, 2.0]]) - value_clean = torch.tensor([[10.0, 20.0]]) - noise_mask = torch.tensor([[1, 0, 1, 0]]) - got = select_per_token(value_noisy, value_clean, noise_mask, seq_len=4) - expected = torch.tensor([[[1.0, 2.0], [10.0, 20.0], [1.0, 2.0], [10.0, 20.0]]]) - assert got.shape == (1, 4, 2) - assert torch.equal(got, expected) - - -def test_noisy_tokens_use_real_timestep_and_clean_tokens_use_zero(): - """Noisy tokens take t_embedder(t * t_scale), clean tokens t_embedder(0).""" - model = tiny_model() - seen = [] - - def record_input(module, args): - seen.append(args[0].detach().clone()) - - handle = model.t_embedder.register_forward_pre_hook(record_input) - try: - with torch.no_grad(): - model( - hidden_states=torch.randn(1, 32, 16, 8), - encoder_hidden_states=torch.randn(1, 2, 4, 48), - timestep=torch.full((1,), 0.5), - encoder_attention_mask=torch.ones(1, 4, dtype=torch.bool), - ) - finally: - handle.remove() - - assert len(seen) == 2, "one embedding for the noisy tokens, one for the clean" - noisy_input, clean_input = seen - t_scale = model.config.t_scale - assert torch.allclose(noisy_input, torch.full((1,), 0.5) * t_scale) - assert torch.equal(clean_input, torch.zeros(1)) - - -def test_block_global_and_per_token_paths_agree_without_context(): - """Without context latents the refiner must modulate globally, not per token.""" - model = tiny_model() - hidden_states = torch.randn(1, 32, 16, 8) - encoder_hidden_states = torch.randn(1, 2, 4, 48) - timestep = torch.full((1,), 0.5) - mask = torch.ones(1, 4, dtype=torch.bool) - - with torch.no_grad(): - t2i = model( - hidden_states, - encoder_hidden_states, - timestep, - encoder_attention_mask=mask, - ) - assert torch.isfinite(t2i).all() - assert t2i.shape == hidden_states.shape - - -# --------------------------------------------------------------------------- -# T11 construction and forward (both branches) -# --------------------------------------------------------------------------- - - -def test_tiny_model_forward_text_to_image(): - model = tiny_model() - hidden_states = torch.randn(1, 32, 16, 16) - with torch.no_grad(): - out = model( - hidden_states, - torch.randn(1, 2, 5, 48), - torch.full((1,), 0.5), - encoder_attention_mask=torch.ones(1, 5, dtype=torch.bool), - ) - assert out.shape == hidden_states.shape - assert torch.isfinite(out).all() - - -def test_tiny_model_forward_with_source_latents_of_mixed_sizes(): - model = tiny_model() - hidden_states = torch.randn(1, 32, 16, 16) - context_latents = [torch.randn(1, 32, 8, 8), torch.randn(1, 32, 16, 4)] - with torch.no_grad(): - out = model( - hidden_states, - torch.randn(1, 2, 5, 48), - torch.full((1,), 0.5), - encoder_attention_mask=torch.ones(1, 5, dtype=torch.bool), - source_latents=context_latents, - ) - assert out.shape == hidden_states.shape - assert torch.isfinite(out).all() - - -def test_tiny_model_forward_without_text_mask(): - model = tiny_model() - hidden_states = torch.randn(2, 32, 8, 8) - with torch.no_grad(): - out = model(hidden_states, torch.randn(2, 2, 4, 48), torch.full((2,), 0.25)) - assert out.shape == hidden_states.shape - - -def test_batch_with_mixed_gist_counts_is_rejected(): - """A batch whose prompts need different gist counts must fail, not silently pad.""" - model = tiny_model() - batch = 2 - # text lengths 3 and 8 give different gist counts (4 vs 8) - attention_mask = torch.zeros(batch, 8, dtype=torch.bool) - attention_mask[0, :3] = True - attention_mask[1, :8] = True - with pytest.raises(ValueError, match="different numbers of gist tokens"): - with torch.no_grad(): - model( - torch.randn(batch, 32, 16, 8), - torch.randn(batch, 2, 8, 48), - torch.full((batch,), 0.5), - encoder_attention_mask=attention_mask, - ) - - -def test_no_context_images_are_registered_as_buffers(): - """The model owns no context-latent state; context arrives per call.""" - model = tiny_model() - names = {name for name, _ in model.named_buffers()} - assert not any("context" in name for name in names) - - -# --------------------------------------------------------------------------- -# T12 cross-attention dimensions -# --------------------------------------------------------------------------- - - -def test_cross_attention_reads_cap_feat_dim_and_writes_dim(): - model = tiny_model() - dim = model.config.dim - cap_feat_dim = model.config.cap_feat_dim - injection_layer_ids = list(model.config.injection_layer_ids) - assert injection_layer_ids == [1] - - block = model.layers[1] - assert block.cross_attention is not None - assert block.cross_attention.to_q.weight.shape == (dim, dim) - # the text bundle is cap_feat_dim wide, so k/v project from there - assert block.cross_attention.to_k.weight.shape == (dim, cap_feat_dim) - assert block.cross_attention.to_v.weight.shape == (dim, cap_feat_dim) - assert block.cross_attn_gate.shape == (1,) - - # blocks that are not injection blocks carry no cross-attention - assert model.layers[0].cross_attention is None - assert model.layers[2].cross_attention is None - - -def test_attention_is_not_fused(): - """q/k/v stay separate: cross-attention q and k/v have different input widths.""" - attn = BriaFibo2Attention(query_dim=8, heads=2, dim_head=4, kv_dim=12) - names = dict(attn.named_parameters()) - assert "to_q.weight" in names and names["to_q.weight"].shape == (8, 8) - assert "to_k.weight" in names and names["to_k.weight"].shape == (8, 12) - assert "to_v.weight" in names and names["to_v.weight"].shape == (8, 12) - # the reference keeps the output projection in a ModuleList -> `to_out.0` - assert "to_out.0.weight" in names - assert "norm_q.weight" in names and "norm_k.weight" in names - - -# --------------------------------------------------------------------------- -# T13 module placement -# --------------------------------------------------------------------------- - - -def test_adaln_modulation_is_replicated_and_full_width(): - """adaLN must emit the whole 4*dim vector: sharding it silently halves it.""" - model = tiny_model() - dim = model.config.dim - block = model.layers[0] - linear = block.adaLN_modulation[0] - assert isinstance(linear, torch.nn.Linear) - assert linear.weight.shape == (4 * dim, min(dim, ADALN_EMBED_DIM)) - assert linear.bias.shape == (4 * dim,) - - -def test_perceiver_latent_queries_shape(): - model = tiny_model() - queries = model.perceiver.latent_queries - assert isinstance(queries, torch.nn.Parameter) - assert queries.shape == ( - model.config.max_num_gist_tokens, - model.config.dim, - ) - - -def test_final_layer_outputs_one_patch_per_token(): - model = tiny_model() - patch_size = model.config.patch_size - out_features = patch_size * patch_size * model.config.out_channels - assert model.final_layer.linear.out_features == out_features - - -# --------------------------------------------------------------------------- -# Weight map: the checkpoint's own names and shapes must match one to one -# --------------------------------------------------------------------------- - - -def _load_key_shapes(path): - key_shapes = {} - with open(path, encoding="utf-8") as handle: - for line in handle: - line = line.rstrip("\n") - if not line.strip(): - continue - fields = line.split("\t") - if len(fields) < 3: - continue - name, shape = fields[0], fields[2] - dims = tuple(int(x) for x in shape.strip("[]").split(",") if x.strip()) - key_shapes[name] = dims - return key_shapes - - -@pytest.mark.skipif( - not os.path.isfile(os.path.join(META_CONFIG_DIR, "dit_keys_shapes.tsv")), - reason="checkpoint key/shape dump not available", -) -def test_state_dict_matches_checkpoint_keys_and_shapes(): - """Build the full 8B model on the meta device and compare against the dump. - - This is the guard against a silent weight reshuffle: every checkpoint key - must exist with the same shape, and the model must add nothing the - checkpoint does not ship. - """ - key_shapes = _load_key_shapes(os.path.join(META_CONFIG_DIR, "dit_keys_shapes.tsv")) - config = BriaFibo2DitConfig() - with torch.device("meta"): - model = BriaFibo2Transformer2DModel(config, hf_config={}) - mine = {name: tuple(t.shape) for name, t in model.state_dict().items()} - - missing = sorted(set(key_shapes) - set(mine)) - extra = sorted(set(mine) - set(key_shapes)) - mismatched = [ - (name, key_shapes[name], mine[name]) - for name in sorted(set(key_shapes) & set(mine)) - if key_shapes[name] != mine[name] - ] - - assert not missing, f"checkpoint keys missing from the model: {missing[:10]}" - assert not extra, f"model keys absent from the checkpoint: {extra[:10]}" - assert not mismatched, f"shape mismatches: {mismatched[:10]}" - - -def test_reference_parameter_count_matches_arch_config(): - """Spot-check the derived shapes the arch config is responsible for.""" - config = BriaFibo2ArchConfig() - assert config.dim == 3584 - assert config.num_layers == 44 - assert config.num_attention_heads == 28 - assert config.cap_feat_dim == 7680 - assert list(config.injection_layer_ids) == [6, 14, 22, 30, 38] - assert sum(config.axes_dims) == config.dim // config.num_attention_heads - # (dim * 8 // 3) // 64 * 64, the SwiGLU inner width the checkpoint ships - assert (config.dim * 8 // 3) // 64 * 64 == 9536 - - -def test_arch_config_rejects_inconsistent_axes(): - with pytest.raises(ValueError, match="axes_dims must sum"): - BriaFibo2ArchConfig(dim=64, num_attention_heads=2, axes_dims=(8, 12, 8)) diff --git a/python/sglang/multimodal_gen/test/unit/test_bria_fibo2_encoder.py b/python/sglang/multimodal_gen/test/unit/test_bria_fibo2_encoder.py deleted file mode 100644 index d9c8f31a4d19..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_bria_fibo2_encoder.py +++ /dev/null @@ -1,227 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""CPU-only tests for the Fibo 2.0 text-encoder side. - -The two failure modes guarded here both produce a model that runs and returns -wrong images without raising anything: using the reference layer indices -directly against the in-tree Qwen3-VL tuple (which is offset by one), and -building the encoder from the inherited Qwen3-VL-8B defaults instead of the 4B -values the checkpoint ships. -""" - -import json -import math -import os - -import pytest -import torch - -from sglang.multimodal_gen.configs.models.encoders.bria_fibo2 import ( - PROMPT_TEMPLATE_ENCODE, - SGLANG_HIDDEN_STATE_OFFSET, - TEXT_ENCODER_OUT_LAYERS, - BriaFibo2Qwen3VLArchConfig, - BriaFibo2Qwen3VLConfig, - build_caption_bundles, - sglang_layer_indices, - unique_hf_layers, -) - -MODEL_CFG_DIR = os.environ.get("FIBO2_MODEL_CFG_DIR", "/workspace/dsh/fibo2/model_cfg") - - -def _load_text_encoder_config() -> dict: - path = os.path.join(MODEL_CFG_DIR, "text_encoder_config.json") - with open(path, encoding="utf-8") as handle: - return json.load(handle) - - -# --------------------------------------------------------------------------- -# Reference constants -# --------------------------------------------------------------------------- - - -def test_out_layers_match_the_reference(): - assert TEXT_ENCODER_OUT_LAYERS == ( - (9, 20, 31), - (5, 16, 27), - (7, 18, 29), - (9, 20, 31), - (11, 22, 33), - (13, 24, 35), - ) - - -def test_fifteen_unique_layers_not_eighteen_slots(): - """Six bundles of three layers share one repeated bundle, so 15 not 18.""" - layers = unique_hf_layers() - assert len(layers) == 15 - assert layers == [5, 7, 9, 11, 13, 16, 18, 20, 22, 24, 27, 29, 31, 33, 35] - - -def test_sglang_indices_are_shifted_by_one(): - assert SGLANG_HIDDEN_STATE_OFFSET == 1 - assert sglang_layer_indices() == ( - (8, 19, 30), - (4, 15, 26), - (6, 17, 28), - (8, 19, 30), - (10, 21, 32), - (12, 23, 34), - ) - - -# --------------------------------------------------------------------------- -# R1: the layer offset. A one-hot state per layer identifies which layer a -# bundle actually read, so a missing shift changes the answer instead of quietly -# shifting the weights. -# --------------------------------------------------------------------------- - - -def _one_hot_states(hidden_size: int, num_layers: int, text_len: int = 2): - """states[i] has RMS 1 and a single non-zero channel at position i.""" - states = [] - for layer in range(num_layers): - state = torch.zeros(1, text_len, hidden_size) - state[..., layer] = math.sqrt(hidden_size) - states.append(state) - return tuple(states) - - -@pytest.mark.parametrize("bundle_index", range(len(TEXT_ENCODER_OUT_LAYERS))) -def test_bundles_read_the_shifted_layer(bundle_index): - hidden_size = 64 # wider than the layer count, so the marker is unambiguous - num_layers = 36 - states = _one_hot_states(hidden_size, num_layers) - - bundles = build_caption_bundles( - states, num_hidden_layers=num_layers, dtype=torch.float32 - ) - assert bundles.shape == (1, 6, 2, len(TEXT_ENCODER_OUT_LAYERS[0]) * hidden_size) - - bundle = bundles[0, bundle_index] - expected_layers = TEXT_ENCODER_OUT_LAYERS[bundle_index] - for block, hf_layer in enumerate(expected_layers): - block_values = bundle[0, block * hidden_size : (block + 1) * hidden_size] - marker = int(torch.argmax(block_values.abs()).item()) - # the in-tree tuple holds layer outputs, so the reference index is off by one - assert marker == hf_layer - SGLANG_HIDDEN_STATE_OFFSET, ( - f"bundle {bundle_index} block {block} read layer {marker}, " - f"expected {hf_layer - SGLANG_HIDDEN_STATE_OFFSET}" - ) - - -def test_bundle_values_are_unit_rms(): - hidden_size = 16 - num_layers = 36 - states = tuple( - torch.randn(1, 3, hidden_size) * (layer + 1) for layer in range(num_layers) - ) - bundles = build_caption_bundles( - states, num_hidden_layers=num_layers, dtype=torch.float32 - ) - for block in range(3): - block_values = bundles[..., block * hidden_size : (block + 1) * hidden_size] - rms = block_values.pow(2).mean(dim=-1).sqrt() - assert torch.allclose(rms, torch.ones_like(rms), atol=1e-5) - - -def test_bundles_cast_back_to_the_requested_dtype(): - states = _one_hot_states(64, 36) - bundles = build_caption_bundles(states, num_hidden_layers=36, dtype=torch.bfloat16) - assert bundles.dtype == torch.bfloat16 - - -def test_layer_count_mismatch_fails_loudly(): - """If the tower ever adds the embedding entry, this must raise, not misread.""" - hidden_size, num_layers = 16, 36 - # 37 entries: the `transformers` convention, i.e. one too many for us - states = tuple(torch.randn(1, 2, hidden_size) for _ in range(num_layers + 1)) - with pytest.raises(ValueError, match="one hidden state per decoder layer"): - build_caption_bundles(states, num_hidden_layers=num_layers) - - -def test_bundle_rejects_wrong_cap_feat_dim(): - """3 x text hidden size must equal the transformer's cap_feat_dim.""" - config = BriaFibo2Qwen3VLConfig( - arch_config=BriaFibo2Qwen3VLArchConfig( - hidden_size=2560, - num_attention_heads=32, - num_key_value_heads=8, - num_hidden_layers=36, - cap_feat_dim=7000, - ) - ) - with pytest.raises(ValueError, match="cap_feat_dim must be 3"): - config.post_diffusers_config_update() - - -# --------------------------------------------------------------------------- -# R5: flatten the real 4B shape out of the checkpoint config -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif( - not os.path.isfile(os.path.join(MODEL_CFG_DIR, "text_encoder_config.json")), - reason="model config dump not available", -) -def test_arch_is_flattened_from_the_real_text_config(): - config = BriaFibo2Qwen3VLConfig() - config.update_model_arch(_load_text_encoder_config()) - config.post_diffusers_config_update() - arch = config.arch_config - - # Qwen3-VL-4B-Instruct, not the inherited 8B defaults - assert arch.hidden_size == 2560 - assert arch.num_hidden_layers == 36 - assert arch.num_attention_heads == 32 - assert arch.num_key_value_heads == 8 - assert arch.head_dim == 128 - assert arch.intermediate_size == 9728 - assert arch.vocab_size == 151936 - assert arch.rms_norm_eps == pytest.approx(1e-6) - # rope_theta lives under `rope_parameters` in this checkpoint - assert arch.rope_theta == pytest.approx(5_000_000.0) - - vision = arch.vision_config - assert vision["depth"] == 24 - assert vision["hidden_size"] == 1024 - assert vision["out_hidden_size"] == 2560 - assert vision["spatial_merge_size"] == 2 - assert vision["patch_size"] == 16 - assert vision["deepstack_visual_indexes"] == [5, 11, 17] - - -@pytest.mark.skipif( - not os.path.isfile(os.path.join(MODEL_CFG_DIR, "text_encoder_config.json")), - reason="model config dump not available", -) -def test_bundle_width_matches_the_transformer_cross_attention_width(): - """3 x 2560 must be exactly the 7680 the DiT cross-attention reads.""" - from sglang.multimodal_gen.configs.models.dits.bria_fibo2 import ( - BriaFibo2ArchConfig, - ) - - config = BriaFibo2Qwen3VLConfig() - config.update_model_arch(_load_text_encoder_config()) - config.post_diffusers_config_update() - - dit_arch = BriaFibo2ArchConfig() - assert config.arch_config.cap_feat_dim == dit_arch.cap_feat_dim - assert 3 * config.arch_config.hidden_size == dit_arch.cap_feat_dim - - -def test_prompt_template_matches_the_reference(): - assert ( - PROMPT_TEMPLATE_ENCODE - == "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" - ) - - -def test_encoder_keeps_its_own_tokenizer_settings(): - """The Fibo encoder must not inherit JoyImage's prompt offsets.""" - arch = BriaFibo2Qwen3VLArchConfig() - assert arch.max_sequence_length == 4096 - assert arch.tokenizer_kwargs["max_length"] == 4096 - assert not hasattr(arch, "prompt_template_encode_start_idx") or ( - arch.prompt_template_encode_start_idx == {"image": 34, "video": 91} - ) diff --git a/python/sglang/multimodal_gen/test/unit/test_bria_fibo2_parity.py b/python/sglang/multimodal_gen/test/unit/test_bria_fibo2_parity.py deleted file mode 100644 index 6e9c83dd845d..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_bria_fibo2_parity.py +++ /dev/null @@ -1,229 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Numerical parity of the in-tree Fibo 2.0 transformer against diffusers. - -The strongest check we can run without weights or a GPU: build both models from -a tiny configuration, copy diffusers' parameters into ours by name, and compare -the outputs element-wise. Because the whole forward runs, a match proves far more -than shape assertions - patchify feature order, the three-axis RoPE planes and -rotation direction, per-token noisy/clean modulation, the Perceiver and gist -construction, the injection blocks' cross-attention wiring, the ``[gist | image | -context]`` token layout, the final layer reading only image tokens, and the -unpatchify flip are all exercised at once. - -The test skips when the installed diffusers predates the model (0.42.0). -""" - -import pytest -import torch - -diffusers = pytest.importorskip("diffusers") - - -def _ref_transformer_cls(): - try: - from diffusers.models.transformers.transformer_bria_fibo2 import ( - BriaFibo2Transformer2DModel as RefModel, - ) - except Exception as exc: # model not in this diffusers version - pytest.skip(f"diffusers has no BriaFibo2 transformer: {exc}") - return RefModel - - -# Small enough to run in a second on CPU, large enough to touch every branch: -# one injection block, a Perceiver layer, a two-axis-value text length. -REFERENCE_KWARGS = dict( - in_channels=32, - patch_size=2, - dim=64, - n_layers=3, - n_refiner_layers=1, - n_heads=2, - norm_eps=1e-5, - cap_feat_dim=48, - injection_layer_ids=(1,), - perceiver_num_layers=1, - min_num_gist_tokens=4, - max_num_gist_tokens=8, - gist_step=4, - gist_min_text_len=2, - gist_max_text_len=8, - rope_theta=256.0, - t_scale=1000.0, - # the sum must be the attention head dim, 64 // 2 - axes_dims=(8, 12, 12), - axes_lens=(32, 8, 8), -) - - -def _build_reference(): - torch.manual_seed(0) - return _ref_transformer_cls()(**REFERENCE_KWARGS).eval() - - -def _build_ours(): - from sglang.multimodal_gen.configs.models.dits.bria_fibo2 import ( - BriaFibo2DitConfig, - ) - from sglang.multimodal_gen.runtime.models.dits.bria_fibo2 import ( - BriaFibo2Transformer2DModel, - ) - - torch.manual_seed(0) - config = BriaFibo2DitConfig() - arch = config.arch_config - # Reference convention: noise levels come in as [0, 1] and the embedder - # scales them by `t_scale` itself. - arch.scheduler_timesteps = 1.0 - arch.in_channels = REFERENCE_KWARGS["in_channels"] - arch.patch_size = REFERENCE_KWARGS["patch_size"] - arch.dim = REFERENCE_KWARGS["dim"] - arch.num_layers = REFERENCE_KWARGS["n_layers"] - arch.n_refiner_layers = REFERENCE_KWARGS["n_refiner_layers"] - arch.num_attention_heads = REFERENCE_KWARGS["n_heads"] - arch.norm_eps = REFERENCE_KWARGS["norm_eps"] - arch.cap_feat_dim = REFERENCE_KWARGS["cap_feat_dim"] - arch.injection_layer_ids = REFERENCE_KWARGS["injection_layer_ids"] - arch.perceiver_num_layers = REFERENCE_KWARGS["perceiver_num_layers"] - arch.min_num_gist_tokens = REFERENCE_KWARGS["min_num_gist_tokens"] - arch.max_num_gist_tokens = REFERENCE_KWARGS["max_num_gist_tokens"] - arch.gist_step = REFERENCE_KWARGS["gist_step"] - arch.gist_min_text_len = REFERENCE_KWARGS["gist_min_text_len"] - arch.gist_max_text_len = REFERENCE_KWARGS["gist_max_text_len"] - arch.rope_theta = REFERENCE_KWARGS["rope_theta"] - arch.t_scale = REFERENCE_KWARGS["t_scale"] - arch.axes_dims = REFERENCE_KWARGS["axes_dims"] - arch.axes_lens = REFERENCE_KWARGS["axes_lens"] - arch.__post_init__() - return BriaFibo2Transformer2DModel(config, hf_config={}).eval() - - -def _copy_reference_weights(reference, ours) -> None: - """Every reference key must exist with the same shape, then be copied over.""" - ref_state = reference.state_dict() - our_state = ours.state_dict() - - missing = sorted(set(ref_state) - set(our_state)) - extra = sorted(set(our_state) - set(ref_state)) - shape_mismatch = [ - name - for name in set(ref_state) & set(our_state) - if ref_state[name].shape != our_state[name].shape - ] - assert not missing, f"reference keys absent from our model: {missing[:5]}" - assert not extra, f"our keys absent from the reference: {extra[:5]}" - assert not shape_mismatch, f"shape mismatches: {shape_mismatch[:5]}" - - with torch.no_grad(): - for name, value in ref_state.items(): - our_state[name].copy_(value) - - -def _assert_parity(tag: str, expected: torch.Tensor, actual: torch.Tensor) -> None: - assert expected.shape == actual.shape, (tag, expected.shape, actual.shape) - expected_flat = expected.flatten().double() - actual_flat = actual.flatten().double() - max_abs = float((expected_flat - actual_flat).abs().max()) - cosine = float( - (expected_flat * actual_flat).sum() - / (expected_flat.norm() * actual_flat.norm()) - ) - # fp32 on CPU: anything beyond rounding means a real semantic difference - assert max_abs < 1e-4, f"{tag}: max|delta| = {max_abs:.3e}" - assert cosine > 0.99999, f"{tag}: cosine = {cosine:.9f}" - - -@pytest.fixture(scope="module") -def models(): - reference = _build_reference() - ours = _build_ours() - _copy_reference_weights(reference, ours) - return reference, ours - - -def test_transformer_matches_diffusers_text_to_image(models): - reference, ours = models - batch, channels, height, width = 1, 32, 16, 16 - text_len = 6 - cap_feat_dim = REFERENCE_KWARGS["cap_feat_dim"] - num_bundles = len(REFERENCE_KWARGS["injection_layer_ids"]) + 1 - - torch.manual_seed(1) - hidden_states = torch.randn(batch, channels, height, width) - encoder_hidden_states = torch.randn(batch, num_bundles, text_len, cap_feat_dim) - attention_mask = torch.ones(batch, text_len, dtype=torch.bool) - timestep = torch.tensor([0.5]) - - with torch.no_grad(): - expected = reference( - hidden_states, - timestep, - encoder_hidden_states, - encoder_attention_mask=attention_mask, - return_dict=False, - )[0] - actual = ours( - hidden_states, - encoder_hidden_states, - timestep, - encoder_attention_mask=attention_mask, - ) - _assert_parity("text-to-image", expected, actual) - - -def test_transformer_matches_diffusers_with_context_latents(models): - """Editing conditions on clean latents of reference images of any size.""" - reference, ours = models - batch, channels, height, width = 1, 32, 16, 16 - text_len = 6 - cap_feat_dim = REFERENCE_KWARGS["cap_feat_dim"] - num_bundles = len(REFERENCE_KWARGS["injection_layer_ids"]) + 1 - - torch.manual_seed(2) - hidden_states = torch.randn(batch, channels, height, width) - encoder_hidden_states = torch.randn(batch, num_bundles, text_len, cap_feat_dim) - attention_mask = torch.ones(batch, text_len, dtype=torch.bool) - timestep = torch.tensor([0.5]) - # two reference images, deliberately different sizes - reference_latents = [ - torch.randn(batch, channels, height, width), - torch.randn(batch, channels, height // 2, width), - ] - - with torch.no_grad(): - expected = reference( - hidden_states, - timestep, - encoder_hidden_states, - encoder_attention_mask=attention_mask, - context_latents=reference_latents, - return_dict=False, - )[0] - actual = ours( - hidden_states, - encoder_hidden_states, - timestep, - encoder_attention_mask=attention_mask, - source_latents=reference_latents, - ) - _assert_parity("editing", expected, actual) - - -def test_transformer_matches_diffusers_without_text_mask(models): - """An unpadded prompt runs without a mask; gist count comes from the length.""" - reference, ours = models - batch, channels, height, width = 1, 32, 16, 16 - text_len = 5 - cap_feat_dim = REFERENCE_KWARGS["cap_feat_dim"] - num_bundles = len(REFERENCE_KWARGS["injection_layer_ids"]) + 1 - - torch.manual_seed(3) - hidden_states = torch.randn(batch, channels, height, width) - encoder_hidden_states = torch.randn(batch, num_bundles, text_len, cap_feat_dim) - timestep = torch.tensor([0.25]) - - with torch.no_grad(): - expected = reference( - hidden_states, timestep, encoder_hidden_states, return_dict=False - )[0] - actual = ours(hidden_states, encoder_hidden_states, timestep) - _assert_parity("no text mask", expected, actual)