github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added experimental support for Ideogram 4

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-04 09:03:00 -06:00
parent b5f21ae695
commit 63b3181262
12 changed files with 2502 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
extensions_built_in/diffusion_models/ideogram4/src/__init__.py Normal file
View File

272
extensions_built_in/diffusion_models/ideogram4/src/latent_norm.py Normal file
View File

@@ -0,0 +1,272 @@
from __future__ import annotations
import torch
LATENT_SHIFT: tuple[float, ...] = (
0.01984364,
0.10149707,
0.29689495,
0.27188619,
-0.21445648,
-0.15979549,
0.05021099,
-0.15083604,
-0.15360136,
-0.20131799,
0.01922352,
0.0622626,
0.10140969,
-0.06739428,
0.3758261,
-0.233712,
0.35164491,
-0.02590912,
-0.0271935,
-0.10833897,
-0.1476848,
-0.01130957,
-0.2298372,
0.23526423,
-0.10893522,
0.11957631,
0.04047799,
0.3134589,
-0.17225064,
-0.18646109,
-0.34691978,
-0.03571246,
0.02583857,
0.10190072,
0.28402294,
0.26952152,
-0.21634675,
-0.17938656,
0.04358909,
-0.15007621,
-0.1548502,
-0.18971131,
0.02710861,
0.05609494,
0.10697846,
-0.06854968,
0.38167698,
-0.24269937,
0.35705471,
-0.03063305,
-0.02946109,
-0.11244286,
-0.14336038,
-0.01362137,
-0.21863696,
0.23228983,
-0.11739769,
0.11693044,
0.02563311,
0.31356594,
-0.17420591,
-0.19006285,
-0.34905377,
-0.04025005,
0.01924137,
0.07652984,
0.2995608,
0.2628057,
-0.22011674,
-0.12715361,
0.04879879,
-0.14075719,
-0.15935895,
-0.2123584,
0.01974813,
0.05523547,
0.10011992,
-0.06428964,
0.37781868,
-0.21491644,
0.34254215,
-0.03153528,
-0.0310082,
-0.10761415,
-0.14730405,
-0.02475182,
-0.2285588,
0.2515081,
-0.10445128,
0.12446,
0.07062869,
0.30880162,
-0.18016875,
-0.18869164,
-0.34533499,
-0.0129177,
0.02578168,
0.07993659,
0.28642181,
0.26038408,
-0.22459419,
-0.14820155,
0.04059549,
-0.14043529,
-0.16111187,
-0.2020305,
0.02602069,
0.04852717,
0.10432153,
-0.06309942,
0.38402443,
-0.22397003,
0.34814481,
-0.03774432,
-0.03381438,
-0.11245691,
-0.14128767,
-0.02853208,
-0.21752016,
0.24872463,
-0.11399775,
0.1222687,
0.05620835,
0.309178,
-0.18065738,
-0.19401479,
-0.34495114,
-0.01760592,
)
LATENT_SCALE: tuple[float, ...] = (
1.63933691,
1.70204478,
1.73642566,
1.90004803,
1.6675316,
1.69059584,
1.56853198,
1.62314944,
1.89106626,
1.58086668,
1.60822129,
1.60962993,
1.63322129,
1.56074359,
1.73419528,
1.7919265,
1.64040632,
1.66802808,
1.60390303,
1.75480492,
1.63187587,
1.64334594,
1.61722884,
1.60146046,
1.63459219,
1.55291476,
1.68771497,
1.68415657,
1.78966054,
1.66631641,
1.65626686,
1.65976433,
1.63487607,
1.69513249,
1.72933756,
1.91310663,
1.67035057,
1.72286863,
1.56719251,
1.61934825,
1.88628859,
1.56911539,
1.59455129,
1.60829869,
1.62470611,
1.56052853,
1.73677003,
1.77563606,
1.63732541,
1.66370527,
1.59508952,
1.75153949,
1.63029275,
1.64517667,
1.61659342,
1.59722044,
1.64103121,
1.5408531,
1.68610394,
1.67772755,
1.78998563,
1.66621713,
1.65458955,
1.66041308,
1.64710857,
1.68163503,
1.74000294,
1.92784786,
1.67411194,
1.67395548,
1.57406532,
1.62199356,
1.87618195,
1.5584375,
1.57438785,
1.61711053,
1.63094305,
1.55644029,
1.73124302,
1.80666627,
1.6463621,
1.65932006,
1.60816188,
1.75682671,
1.64695873,
1.63121722,
1.61380832,
1.60478651,
1.63396035,
1.53505068,
1.65534289,
1.67132281,
1.80317197,
1.6767314,
1.65700938,
1.68426259,
1.65339716,
1.67540638,
1.73298504,
1.94067348,
1.67893609,
1.70635117,
1.5730906,
1.61928553,
1.87148809,
1.56244866,
1.56697152,
1.61584394,
1.62759496,
1.55480378,
1.73484107,
1.79055143,
1.64688773,
1.66121492,
1.60135887,
1.75254572,
1.64798332,
1.62989921,
1.61381592,
1.60792883,
1.63939668,
1.53075757,
1.65371318,
1.66801185,
1.80029087,
1.67591476,
1.65655173,
1.68533454,
)
def get_latent_norm() -> tuple[torch.Tensor, torch.Tensor]:
shift = torch.tensor(LATENT_SHIFT, dtype=torch.float32)
scale = torch.tensor(LATENT_SCALE, dtype=torch.float32)
assert shift.shape == (128,) and scale.shape == (128,)
return shift, scale

310
extensions_built_in/diffusion_models/ideogram4/src/pipeline.py Normal file
View File

@@ -0,0 +1,310 @@
"""Packing / sampling helpers for Ideogram 4.
This module holds the glue that turns image latents + Qwen3-VL text features into
the single packed sequence the transformer consumes, plus a minimal flow-matching
sampling pipeline used to render preview images during training.
"""
from __future__ import annotations
from typing import List, Optional
import torch
from PIL import Image
from diffusers.utils.torch_utils import randn_tensor
from transformers.masking_utils import create_causal_mask
from .transformer import (
IMAGE_POSITION_OFFSET,
LLM_TOKEN_INDICATOR,
OUTPUT_IMAGE_INDICATOR,
QWEN3_VL_ACTIVATION_LAYERS,
SEQUENCE_PADDING_INDICATOR,
Ideogram4Transformer2DModel,
)
# ---------------------------------------------------------------------------
# Latent (un)patchification.
#
# The VAE produces (B, ae_ch=32, H/8, W/8) latents. The transformer works on
# tokens of dim ae_ch * patch**2 = 128. We store the patchified latent in a 4-D
# (B, 128, gh, gw) layout so the rest of ai-toolkit (noise, add_noise, loss) can
# treat it like an ordinary image latent. The channel ordering here matches the
# reference Ideogram 4 decode exactly: 128 = (patch_h, patch_w, ae_ch) with ae_ch
# the fastest-varying axis.
# ---------------------------------------------------------------------------
def patchify_latents(z: torch.Tensor, patch_size: int = 2) -> torch.Tensor:
"""(B, ae_ch, H8, W8) -> (B, ae_ch * patch**2, gh, gw)."""
b, ae_ch, h8, w8 = z.shape
ph = pw = patch_size
gh, gw = h8 // ph, w8 // pw
z = z.view(b, ae_ch, gh, ph, gw, pw)
# -> (B, ph, pw, ae_ch, gh, gw) then merge (ph, pw, ae_ch) -> channels
z = z.permute(0, 3, 5, 1, 2, 4).reshape(b, ph * pw * ae_ch, gh, gw)
return z
def unpatchify_latents(z: torch.Tensor, patch_size: int = 2) -> torch.Tensor:
"""(B, ae_ch * patch**2, gh, gw) -> (B, ae_ch, H8, W8)."""
b, c, gh, gw = z.shape
ph = pw = patch_size
ae_ch = c // (ph * pw)
z = z.view(b, ph, pw, ae_ch, gh, gw)
# -> (B, ae_ch, gh, ph, gw, pw) then merge spatial
z = z.permute(0, 3, 4, 1, 5, 2).reshape(b, ae_ch, gh * ph, gw * pw)
return z
# ---------------------------------------------------------------------------
# Qwen3-VL hidden-state extraction.
# ---------------------------------------------------------------------------
@torch.no_grad()
def get_qwen3_vl_features(
text_encoder,
token_ids: torch.Tensor,
attention_mask: torch.Tensor,
pos_2d: torch.Tensor,
) -> torch.Tensor:
"""Run Qwen3-VL and concat the hidden states from the activation layers.
Returns a (B, L, hidden_size * num_layers) tensor (in the encoder's dtype),
zeroed at non-text (padding) positions.
"""
language_model = text_encoder.language_model
inputs_embeds = language_model.embed_tokens(token_ids)
position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1)
text_position_ids = position_ids_4d[0]
mrope_position_ids = position_ids_4d[1:]
causal_mask = create_causal_mask(
config=language_model.config,
inputs_embeds=inputs_embeds,
attention_mask=attention_mask,
past_key_values=None,
position_ids=text_position_ids,
)
position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)
tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
captured: dict[int, torch.Tensor] = {}
hidden_states = inputs_embeds
for layer_idx, decoder_layer in enumerate(language_model.layers):
hidden_states = decoder_layer(
hidden_states,
attention_mask=causal_mask,
position_ids=text_position_ids,
past_key_values=None,
position_embeddings=position_embeddings,
)
if layer_idx in tap_set:
captured[layer_idx] = hidden_states
selected = [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
batch_size, seq_len = token_ids.shape
stacked = torch.stack(selected, dim=0) # (num_taps, B, L, H)
stacked = torch.permute(stacked, (1, 2, 3, 0)) # (B, L, H, num_taps)
stacked = stacked.reshape(batch_size, seq_len, -1)
text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
stacked = stacked * text_mask
return stacked
# ---------------------------------------------------------------------------
# Packing + velocity prediction.
# ---------------------------------------------------------------------------
def predict_velocity(
transformer: Ideogram4Transformer2DModel,
latents: torch.Tensor, # (B, 128, gh, gw)
t: torch.Tensor, # (B,) toolkit flow time in [0, 1] (1 = pure noise)
llm_features: torch.Tensor, # (B, Lt, llm_dim)
text_mask: torch.Tensor, # (B, Lt) 1 for real text tokens
) -> torch.Tensor:
"""Run the transformer on the packed [text | image] sequence.
``t`` is in the ai-toolkit flow-matching convention: ``t=1`` is pure noise,
``t=0`` is clean, and the returned velocity is ``noise - clean`` (matching the
toolkit scheduler / loss target).
Ideogram's transformer uses the opposite convention internally (``t=1`` is
clean) and predicts ``clean - noise``, so we feed it ``1 - t`` and negate its
output. Returns the velocity reshaped to the (B, 128, gh, gw) latent layout.
"""
device = latents.device
b, c, gh, gw = latents.shape
num_image_tokens = gh * gw
num_text_tokens = llm_features.shape[1]
seq_len = num_text_tokens + num_image_tokens
# image latents -> tokens (row-major: h outer, w inner)
image_tokens = latents.permute(0, 2, 3, 1).reshape(b, num_image_tokens, c)
# The mask may arrive as a float (PromptEmbeds.to casts it to the embed
# dtype); work in long so cumsum positions stay exact for long prompts.
text_mask_bool = text_mask.to(device) > 0
text_mask_long = text_mask_bool.long()
# noise tokens: text region is zeroed (masked out anyway)
x = torch.cat(
[
torch.zeros(b, num_text_tokens, c, device=device, dtype=image_tokens.dtype),
image_tokens,
],
dim=1,
)
# llm features: image region is zero
llm_full = torch.cat(
[
llm_features,
torch.zeros(
b,
num_image_tokens,
llm_features.shape[-1],
device=device,
dtype=llm_features.dtype,
),
],
dim=1,
)
# indicator: real text -> 3, image -> 2, text pad -> 0
indicator = torch.zeros(b, seq_len, dtype=torch.long, device=device)
indicator[:, :num_text_tokens] = text_mask_long * LLM_TOKEN_INDICATOR
indicator[:, num_text_tokens:] = OUTPUT_IMAGE_INDICATOR
# segment ids: real text + image -> 1, text pad -> -1 (its own padding segment)
segment_ids = torch.ones(b, seq_len, dtype=torch.long, device=device)
segment_ids[:, :num_text_tokens] = torch.where(
text_mask_bool,
torch.ones_like(text_mask_long),
torch.full_like(text_mask_long, SEQUENCE_PADDING_INDICATOR),
)
# position ids (t, h, w)
# text positions: 0..num_real-1 at the real slots (relative; pad -> 0)
text_pos = (text_mask_long.cumsum(dim=-1) - 1).clamp(min=0) # (B, Lt)
text_pos_3d = text_pos.unsqueeze(-1).expand(-1, -1, 3)
h_idx = torch.arange(gh, device=device).view(-1, 1).expand(gh, gw).reshape(-1)
w_idx = torch.arange(gw, device=device).view(1, -1).expand(gh, gw).reshape(-1)
t_idx = torch.zeros_like(h_idx)
image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET
image_pos_3d = image_pos.unsqueeze(0).expand(b, -1, -1)
position_ids = torch.cat([text_pos_3d, image_pos_3d], dim=1)
# Flip into the model's time convention (t=1 -> clean).
model_t = 1.0 - t
out = transformer(
llm_features=llm_full,
x=x,
t=model_t,
position_ids=position_ids,
segment_ids=segment_ids,
indicator=indicator,
)
image_velocity = out[:, num_text_tokens:] # (B, Li, 128)
image_velocity = image_velocity.reshape(b, gh, gw, c).permute(0, 3, 1, 2)
# Model predicts clean - noise; negate to return toolkit velocity (noise - clean).
return -image_velocity
# ---------------------------------------------------------------------------
# Minimal sampling pipeline (for training previews).
# ---------------------------------------------------------------------------
class Ideogram4Pipeline:
"""Lightweight flow-matching sampler used by ai-toolkit's preview generation."""
def __init__(self, model):
# ``model`` is the Ideogram4Model so we can reuse its encode/decode and
# latent helpers without duplicating state.
self.model = model
@property
def device(self):
return self.model.device_torch
def to(self, *args, **kwargs):
return self
@torch.no_grad()
def __call__(
self,
conditional_embeds,
unconditional_embeds,
height: int = 1024,
width: int = 1024,
num_inference_steps: int = 30,
guidance_scale: float = 7.0,
latents: Optional[torch.Tensor] = None,
generator: Optional[torch.Generator] = None,
**kwargs,
) -> List[Image.Image]:
model = self.model
device = model.device_torch
dtype = model.torch_dtype
transformer = model.transformer
patch = model.patch_size
# Use a fresh scheduler so we never mutate the training scheduler's state.
scheduler = model.get_train_scheduler()
scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = scheduler.timesteps
ae_scale = model.vae_scale_factor # 8
gh = height // (ae_scale * patch)
gw = width // (ae_scale * patch)
latent_channels = transformer.config.in_channels
do_cfg = unconditional_embeds is not None and guidance_scale != 1.0
if latents is None:
shape = (1, latent_channels, gh, gw)
latents = randn_tensor(
shape, generator=generator, device=device, dtype=torch.float32
)
latents = latents.to(device, dtype=torch.float32)
latents = latents * scheduler.init_noise_sigma
cond_feats = conditional_embeds.text_embeds.to(device, dtype=dtype)
cond_mask = conditional_embeds.attention_mask.to(device)
if do_cfg:
uncond_feats = unconditional_embeds.text_embeds.to(device, dtype=dtype)
uncond_mask = unconditional_embeds.attention_mask.to(device)
for t in timesteps:
t01 = (t / 1000.0).to(device).expand(latents.shape[0])
v_cond = predict_velocity(
transformer, latents.to(dtype), t01, cond_feats, cond_mask
)
if do_cfg:
v_uncond = predict_velocity(
transformer, latents.to(dtype), t01, uncond_feats, uncond_mask
)
v = v_uncond + guidance_scale * (v_cond - v_uncond)
else:
v = v_cond
latents = scheduler.step(
v.to(torch.float32), t, latents, return_dict=False
)[0]
images = model.decode_latents(latents, device=device, dtype=dtype)
images = images.float().clamp(-1.0, 1.0)
images = ((images + 1.0) * 127.5).round().to(torch.uint8)
images = images.permute(0, 2, 3, 1).cpu().numpy()
return [Image.fromarray(arr) for arr in images]

406
extensions_built_in/diffusion_models/ideogram4/src/transformer.py Normal file
View File

@@ -0,0 +1,406 @@
"""Ideogram4 transformer backbone.
Ported from the reference ``modeling_ideogram4.py`` for ai-toolkit. The
transformer consumes Qwen3-VL hidden states and flow-matching noise tokens
(packed into a single sequence) and produces velocity predictions on the image
latent tokens.
"""
from __future__ import annotations
import math
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
# Per-token role indicators used inside the packed sequence.
SEQUENCE_PADDING_INDICATOR = -1
OUTPUT_IMAGE_INDICATOR = 2
LLM_TOKEN_INDICATOR = 3
# Image grid coordinates start at this offset so they never collide with text
# token indices (text positions start at 0 and never exceed max_text_tokens).
IMAGE_POSITION_OFFSET = 65536
# Layers of Qwen3-VL whose hidden states are concatenated and fed to the transformer.
QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
@dataclass
class Ideogram4Config:
emb_dim: int = 4608
num_layers: int = 34
num_heads: int = 18
intermediate_size: int = 12288
adanln_dim: int = 512
# Latent dimension after patchification: ae_channels (32) * patch_size**2 (4) = 128.
in_channels: int = 128
# Hidden size of Qwen3-VL-8B-Instruct multiplied by the number of extracted layers.
# Qwen3-VL hidden size = 4096
llm_features_dim: int = 4096 * len(QWEN3_VL_ACTIVATION_LAYERS)
rope_theta: int = 5_000_000
mrope_section: tuple[int, ...] = (24, 20, 20)
norm_eps: float = 1e-5
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
half = x.shape[-1] // 2
x1 = x[..., :half]
x2 = x[..., half:]
return torch.cat((-x2, x1), dim=-1)
def _apply_rotary_pos_emb(
q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
# q, k: (B, num_heads, L, head_dim); cos/sin: (B, L, head_dim).
cos = cos.unsqueeze(1)
sin = sin.unsqueeze(1)
q_embed = (q * cos) + (_rotate_half(q) * sin)
k_embed = (k * cos) + (_rotate_half(k) * sin)
return q_embed, k_embed
class Ideogram4MRoPE(nn.Module):
inv_freq: torch.Tensor
def __init__(
self,
head_dim: int,
base: int,
mrope_section: tuple[int, ...],
) -> None:
super().__init__()
inv_freq = 1.0 / (
base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
)
self.register_buffer("inv_freq", inv_freq, persistent=False)
self.mrope_section = tuple(mrope_section)
self.head_dim = head_dim
@torch.no_grad()
def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
# position_ids: (B, L, 3) of int.
assert position_ids.ndim == 3 and position_ids.shape[-1] == 3
batch_size, seq_len, _ = position_ids.shape
# (3, B, inv_freq_size, L)
pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32)
inv_freq = self.inv_freq.to(dtype=torch.float32)[None, None, :, None].expand(
3, batch_size, -1, 1
)
freqs = inv_freq @ pos.unsqueeze(2)
freqs = freqs.transpose(2, 3) # (3, B, L, inv_freq_size)
# interleaved mrope: pull H freqs into idx 1 mod 3, W freqs into idx 2 mod 3.
freqs_t = freqs[0].clone()
for axis, offset in ((1, 1), (2, 2)):
length = self.mrope_section[axis] * 3
idx = torch.arange(offset, length, 3, device=freqs_t.device)
freqs_t[..., idx] = freqs[axis][..., idx]
emb = torch.cat((freqs_t, freqs_t), dim=-1)
return emb.cos(), emb.sin()
class Ideogram4RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.rms_norm(x, self.weight.shape, self.weight, self.eps)
class Ideogram4Attention(nn.Module):
def __init__(self, hidden_size: int, num_heads: int, eps: float = 1e-5) -> None:
super().__init__()
assert hidden_size % num_heads == 0
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps)
self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps)
self.o = nn.Linear(hidden_size, hidden_size, bias=False)
def forward(
self,
x: torch.Tensor,
attn_mask: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
) -> torch.Tensor:
batch_size, seq_len, _ = x.shape
qkv = self.qkv(x)
qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
q, k, v = qkv.unbind(dim=2)
q = self.norm_q(q)
k = self.norm_k(k)
# SDPA expects (B, num_heads, L, head_dim).
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
q, k = _apply_rotary_pos_emb(q, k, cos, sin)
out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
out = out.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)
return self.o(out)
class Ideogram4MLP(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 Ideogram4TransformerBlock(nn.Module):
def __init__(
self,
hidden_size: int,
intermediate_size: int,
num_heads: int,
norm_eps: float,
adanln_dim: int,
) -> None:
super().__init__()
self.attention = Ideogram4Attention(hidden_size, num_heads, eps=1e-5)
self.feed_forward = Ideogram4MLP(hidden_size, intermediate_size)
self.attention_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
self.attention_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
self.ffn_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
self.adaln_modulation = nn.Linear(adanln_dim, 4 * hidden_size, bias=True)
def forward(
self,
x: torch.Tensor,
attn_mask: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
adaln_input: torch.Tensor,
) -> torch.Tensor:
mod = self.adaln_modulation(adaln_input)
scale_msa, gate_msa, scale_mlp, gate_mlp = mod.chunk(4, dim=-1)
gate_msa = torch.tanh(gate_msa)
gate_mlp = torch.tanh(gate_mlp)
scale_msa = 1.0 + scale_msa
scale_mlp = 1.0 + scale_mlp
attn_out = self.attention(
self.attention_norm1(x) * scale_msa,
attn_mask=attn_mask,
cos=cos,
sin=sin,
)
x = x + gate_msa * self.attention_norm2(attn_out)
x = x + gate_mlp * self.ffn_norm2(
self.feed_forward(self.ffn_norm1(x) * scale_mlp)
)
return x
def _sinusoidal_embedding(
t: torch.Tensor, dim: int, scale: float = 1e4
) -> torch.Tensor:
t = t.to(torch.float32)
half = dim // 2
freq = math.log(scale) / (half - 1)
freq = torch.exp(torch.arange(half, dtype=torch.float32, device=t.device) * -freq)
emb = t.unsqueeze(-1) * freq
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
if dim % 2 == 1:
emb = F.pad(emb, (0, 1))
return emb
class Ideogram4EmbedScalar(nn.Module):
def __init__(self, dim: int, input_range: tuple[float, float]) -> None:
super().__init__()
self.dim = dim
self.range_min, self.range_max = input_range
assert self.range_max > self.range_min
self.mlp_in = nn.Linear(dim, dim, bias=True)
self.mlp_out = nn.Linear(dim, dim, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x is shape (..., 1) or (...,) holding a scalar per token.
x = x.to(torch.float32)
scaled = 1e4 * (x - self.range_min) / (self.range_max - self.range_min)
emb = _sinusoidal_embedding(scaled, self.dim)
emb = emb.to(self.mlp_in.weight.dtype)
emb = F.silu(self.mlp_in(emb))
return self.mlp_out(emb)
class Ideogram4FinalLayer(nn.Module):
def __init__(self, hidden_size: int, out_channels: int, adanln_dim: int) -> None:
super().__init__()
self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
self.linear = nn.Linear(hidden_size, out_channels, bias=True)
self.adaln_modulation = nn.Linear(adanln_dim, hidden_size, bias=True)
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
scale = 1.0 + self.adaln_modulation(F.silu(c))
return self.linear(self.norm_final(x) * scale)
class Ideogram4Transformer2DModel(nn.Module):
"""Ideogram 4 flow-matching transformer."""
def __init__(self, config: Ideogram4Config) -> None:
super().__init__()
self.config = config
self.gradient_checkpointing = False
head_dim = config.emb_dim // config.num_heads
self.input_proj = nn.Linear(config.in_channels, config.emb_dim, bias=True)
self.llm_cond_norm = Ideogram4RMSNorm(config.llm_features_dim, eps=1e-6)
self.llm_cond_proj = nn.Linear(
config.llm_features_dim, config.emb_dim, bias=True
)
self.t_embedding = Ideogram4EmbedScalar(config.emb_dim, input_range=(0.0, 1.0))
self.adaln_proj = nn.Linear(config.emb_dim, config.adanln_dim, bias=True)
self.embed_image_indicator = nn.Embedding(2, config.emb_dim)
self.rotary_emb = Ideogram4MRoPE(
head_dim=head_dim,
base=config.rope_theta,
mrope_section=config.mrope_section,
)
self.layers = nn.ModuleList(
[
Ideogram4TransformerBlock(
hidden_size=config.emb_dim,
intermediate_size=config.intermediate_size,
num_heads=config.num_heads,
norm_eps=config.norm_eps,
adanln_dim=config.adanln_dim,
)
for _ in range(config.num_layers)
]
)
self.final_layer = Ideogram4FinalLayer(
hidden_size=config.emb_dim,
out_channels=config.in_channels,
adanln_dim=config.adanln_dim,
)
@property
def device(self) -> torch.device:
return next(self.parameters()).device
@property
def dtype(self) -> torch.dtype:
return next(self.parameters()).dtype
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
def disable_gradient_checkpointing(self):
self.gradient_checkpointing = False
def forward(
self,
*,
llm_features: torch.Tensor,
x: torch.Tensor,
t: torch.Tensor,
position_ids: torch.Tensor,
segment_ids: torch.Tensor,
indicator: torch.Tensor,
) -> torch.Tensor:
"""Velocity prediction.
Args:
llm_features: (B, L, llm_features_dim) Qwen3-VL conditioning features.
x: (B, L, in_channels) noise tokens.
t: (B,) or (B, L) flow-matching time in [0, 1].
position_ids: (B, L, 3) (t, h, w) positions for MRoPE.
segment_ids: (B, L) sample id within a packed batch.
indicator: (B, L) per-token role: LLM_TOKEN_INDICATOR or OUTPUT_IMAGE_INDICATOR.
Returns:
(B, L, in_channels) velocity prediction in float32. Only the positions
with ``indicator == OUTPUT_IMAGE_INDICATOR`` are meaningful.
"""
batch_size, seq_len, in_channels = x.shape
assert in_channels == self.config.in_channels
param_dtype = self.input_proj.weight.dtype
x = x.to(param_dtype)
t = t.to(param_dtype)
llm_features = llm_features.to(param_dtype)
indicator = indicator.to(torch.long)
llm_token_mask = (indicator == LLM_TOKEN_INDICATOR).to(x.dtype).unsqueeze(-1)
output_image_mask = (
(indicator == OUTPUT_IMAGE_INDICATOR).to(x.dtype).unsqueeze(-1)
)
llm_features = llm_features * llm_token_mask
x = x * output_image_mask
x = self.input_proj(x) * output_image_mask
# Keep shape (B, 1, ...) when t is per-sample so downstream adaln_modulation
# projections don't pay for L identical copies.
t_cond = self.t_embedding(t)
if t.dim() == 1:
t_cond = t_cond.unsqueeze(1)
adaln_input = F.silu(self.adaln_proj(t_cond))
llm_features = self.llm_cond_norm(llm_features)
llm_features = self.llm_cond_proj(llm_features) * llm_token_mask
h = x + llm_features
image_indicator_embedding = self.embed_image_indicator(
(indicator == OUTPUT_IMAGE_INDICATOR).to(torch.long)
)
h = h + image_indicator_embedding
cos, sin = self.rotary_emb(position_ids)
cos = cos.to(h.dtype)
sin = sin.to(h.dtype)
# Block-diagonal mask from segment ids: (B, 1, L, L), True = attend.
attn_mask = (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1)
for layer in self.layers:
if self.gradient_checkpointing and torch.is_grad_enabled():
h = checkpoint(
layer, h, attn_mask, cos, sin, adaln_input, use_reentrant=False
)
else:
h = layer(h, attn_mask, cos, sin, adaln_input)
out = self.final_layer(h, c=adaln_input)
return out.to(torch.float32)

431
extensions_built_in/diffusion_models/ideogram4/src/vae.py Normal file
View File

@@ -0,0 +1,431 @@
"""Flux2 KL autoencoder."""
from __future__ import annotations
import math
import re
from dataclasses import dataclass, field
import torch
from einops import rearrange
from torch import Tensor, nn
@dataclass
class AutoEncoderParams:
resolution: int = 256
in_channels: int = 3
ch: int = 128
out_ch: int = 3
ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
num_res_blocks: int = 2
z_channels: int = 32
def swish(x: Tensor) -> Tensor:
return x * torch.sigmoid(x)
class AttnBlock(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.in_channels = in_channels
self.norm = nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
def attention(self, h_: Tensor) -> Tensor:
h_ = self.norm(h_)
q = self.q(h_)
k = self.k(h_)
v = self.v(h_)
b, c, h, w = q.shape
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
def forward(self, x: Tensor) -> Tensor:
return x + self.proj_out(self.attention(x))
class ResnetBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.norm1 = nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
self.conv1 = nn.Conv2d(
in_channels, out_channels, kernel_size=3, stride=1, padding=1
)
self.norm2 = nn.GroupNorm(
num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
)
self.conv2 = nn.Conv2d(
out_channels, out_channels, kernel_size=3, stride=1, padding=1
)
if self.in_channels != self.out_channels:
self.nin_shortcut = nn.Conv2d(
in_channels, out_channels, kernel_size=1, stride=1, padding=0
)
def forward(self, x):
h = x
h = self.norm1(h)
h = swish(h)
h = self.conv1(h)
h = self.norm2(h)
h = swish(h)
h = self.conv2(h)
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return x + h
class Downsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
# no asymmetric padding in torch conv, must do it ourselves
self.conv = nn.Conv2d(
in_channels, in_channels, kernel_size=3, stride=2, padding=0
)
def forward(self, x: Tensor):
pad = (0, 1, 0, 1)
x = nn.functional.pad(x, pad, mode="constant", value=0)
x = self.conv(x)
return x
class Upsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.conv = nn.Conv2d(
in_channels, in_channels, kernel_size=3, stride=1, padding=1
)
def forward(self, x: Tensor):
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
x = self.conv(x)
return x
class Encoder(nn.Module):
def __init__(
self,
resolution: int,
in_channels: int,
ch: int,
ch_mult: list[int],
num_res_blocks: int,
z_channels: int,
):
super().__init__()
self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
# downsampling
self.conv_in = nn.Conv2d(
in_channels, self.ch, kernel_size=3, stride=1, padding=1
)
curr_res = resolution
in_ch_mult = (1,) + tuple(ch_mult)
self.in_ch_mult = in_ch_mult
self.down = nn.ModuleList()
block_in = self.ch
for i_level in range(self.num_resolutions):
block = nn.ModuleList()
attn = nn.ModuleList()
block_in = ch * in_ch_mult[i_level]
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
down = nn.Module()
down.block = block
down.attn = attn
if i_level != self.num_resolutions - 1:
down.downsample = Downsample(block_in)
curr_res = curr_res // 2
self.down.append(down)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# end
self.norm_out = nn.GroupNorm(
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
)
self.conv_out = nn.Conv2d(
block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
)
def forward(self, x: Tensor) -> Tensor:
# downsampling
hs = [self.conv_in(x)]
for i_level in range(self.num_resolutions):
for i_block in range(self.num_res_blocks):
h = self.down[i_level].block[i_block](hs[-1]) # type: ignore[index, operator]
if len(self.down[i_level].attn) > 0: # type: ignore[arg-type]
h = self.down[i_level].attn[i_block](h) # type: ignore[index, operator]
hs.append(h)
if i_level != self.num_resolutions - 1:
hs.append(self.down[i_level].downsample(hs[-1])) # type: ignore[operator]
# middle
h = hs[-1]
h = self.mid.block_1(h) # type: ignore[operator]
h = self.mid.attn_1(h) # type: ignore[operator]
h = self.mid.block_2(h) # type: ignore[operator]
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
h = self.quant_conv(h)
return h
class Decoder(nn.Module):
def __init__(
self,
ch: int,
out_ch: int,
ch_mult: list[int],
num_res_blocks: int,
in_channels: int,
resolution: int,
z_channels: int,
):
super().__init__()
self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
self.ffactor = 2 ** (self.num_resolutions - 1)
# compute in_ch_mult, block_in and curr_res at lowest res
block_in = ch * ch_mult[self.num_resolutions - 1]
curr_res = resolution // 2 ** (self.num_resolutions - 1)
self.z_shape = (1, z_channels, curr_res, curr_res)
# z to block_in
self.conv_in = nn.Conv2d(
z_channels, block_in, kernel_size=3, stride=1, padding=1
)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# upsampling
self.up = nn.ModuleList()
for i_level in reversed(range(self.num_resolutions)):
block = nn.ModuleList()
attn = nn.ModuleList()
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks + 1):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
up = nn.Module()
up.block = block
up.attn = attn
if i_level != 0:
up.upsample = Upsample(block_in)
curr_res = curr_res * 2
self.up.insert(0, up) # prepend to get consistent order
# end
self.norm_out = nn.GroupNorm(
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
)
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
def forward(self, z: Tensor) -> Tensor:
z = self.post_quant_conv(z)
# get dtype for proper tracing
upscale_dtype = next(self.up.parameters()).dtype
# z to block_in
h = self.conv_in(z)
# middle
h = self.mid.block_1(h) # type: ignore[operator]
h = self.mid.attn_1(h) # type: ignore[operator]
h = self.mid.block_2(h) # type: ignore[operator]
# cast to proper dtype
h = h.to(upscale_dtype)
# upsampling
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks + 1):
h = self.up[i_level].block[i_block](h) # type: ignore[index, operator]
if len(self.up[i_level].attn) > 0: # type: ignore[arg-type]
h = self.up[i_level].attn[i_block](h) # type: ignore[index, operator]
if i_level != 0:
h = self.up[i_level].upsample(h) # type: ignore[operator]
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
return h
class AutoEncoder(nn.Module):
def __init__(self, params: AutoEncoderParams):
super().__init__()
self.params = params
self.encoder = Encoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=params.ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
self.decoder = Decoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=params.ch,
out_ch=params.out_ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
self.bn_eps = 1e-4
self.bn_momentum = 0.1
self.ps = [2, 2]
self.bn = torch.nn.BatchNorm2d(
math.prod(self.ps) * params.z_channels,
eps=self.bn_eps,
momentum=self.bn_momentum,
affine=False,
track_running_stats=True,
)
@property
def device(self) -> torch.device:
return next(self.parameters()).device
@property
def dtype(self) -> torch.dtype:
return next(self.parameters()).dtype
_NUM_RESOLUTIONS = 4
def convert_diffusers_state_dict(src: dict[str, Tensor]) -> dict[str, Tensor]:
out: dict[str, Tensor] = {}
attn_substrings = (".mid.attn_1.",)
for src_key, tensor in src.items():
dst_key = _rewrite_diffusers_key(src_key)
if dst_key is None:
raise KeyError(f"Unrecognized diffusers VAE state-dict key: {src_key}")
if (
any(s in dst_key for s in attn_substrings)
and dst_key.endswith(".weight")
and tensor.ndim == 2
):
tensor = tensor.unsqueeze(-1).unsqueeze(-1)
out[dst_key] = tensor
return out
def _rewrite_diffusers_key(key: str) -> str | None:
if key.startswith("bn."):
return key
if key.startswith("quant_conv."):
return key.replace("quant_conv.", "encoder.quant_conv.", 1)
if key.startswith("post_quant_conv."):
return key.replace("post_quant_conv.", "decoder.post_quant_conv.", 1)
if key == "encoder.conv_norm_out.weight":
return "encoder.norm_out.weight"
if key == "encoder.conv_norm_out.bias":
return "encoder.norm_out.bias"
if key == "decoder.conv_norm_out.weight":
return "decoder.norm_out.weight"
if key == "decoder.conv_norm_out.bias":
return "decoder.norm_out.bias"
m = re.match(r"^(encoder|decoder)\.mid_block\.resnets\.(\d+)\.(.+)$", key)
if m:
side, idx, rest = m.group(1), int(m.group(2)), m.group(3)
rest = rest.replace("conv_shortcut", "nin_shortcut")
return f"{side}.mid.block_{idx + 1}.{rest}"
m = re.match(r"^(encoder|decoder)\.mid_block\.attentions\.0\.(.+)$", key)
if m:
side, rest = m.group(1), m.group(2)
rest = (
rest.replace("group_norm.", "norm.")
.replace("to_q.", "q.")
.replace("to_k.", "k.")
.replace("to_v.", "v.")
.replace("to_out.0.", "proj_out.")
)
return f"{side}.mid.attn_1.{rest}"
m = re.match(r"^encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key)
if m:
level, res_idx, rest = m.group(1), m.group(2), m.group(3)
rest = rest.replace("conv_shortcut", "nin_shortcut")
return f"encoder.down.{level}.block.{res_idx}.{rest}"
m = re.match(r"^encoder\.down_blocks\.(\d+)\.downsamplers\.0\.conv\.(.+)$", key)
if m:
return f"encoder.down.{m.group(1)}.downsample.conv.{m.group(2)}"
m = re.match(r"^decoder\.up_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key)
if m:
diffusers_idx = int(m.group(1))
res_idx = m.group(2)
rest = m.group(3).replace("conv_shortcut", "nin_shortcut")
return (
f"decoder.up.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.block.{res_idx}.{rest}"
)
m = re.match(r"^decoder\.up_blocks\.(\d+)\.upsamplers\.0\.conv\.(.+)$", key)
if m:
diffusers_idx = int(m.group(1))
return f"decoder.up.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.upsample.conv.{m.group(2)}"
if key.startswith(
(
"encoder.conv_in.",
"encoder.conv_out.",
"decoder.conv_in.",
"decoder.conv_out.",
)
):
return key
return None