362 lines
14 KiB
Python
362 lines
14 KiB
Python
"""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 pad_text_features(
|
|
features_list: List[torch.Tensor],
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Right-pad a list of per-sample (Lt_i, D) features into a batch.
|
|
|
|
Captions are stored at their natural length (one tensor per batch item) and
|
|
only padded to the batch max here, right before the model call. Returns
|
|
``(features (B, Lt, D), attention_mask (B, Lt))``; the mask is 1 for real
|
|
tokens and 0 for padding (which the transformer masks out anyway).
|
|
"""
|
|
lengths = [f.shape[0] for f in features_list]
|
|
max_len = max(lengths)
|
|
dim = features_list[0].shape[-1]
|
|
batch_size = len(features_list)
|
|
|
|
features = torch.zeros(batch_size, max_len, dim, device=device, dtype=dtype)
|
|
mask = torch.zeros(batch_size, max_len, dtype=torch.long, device=device)
|
|
for i, f in enumerate(features_list):
|
|
n = f.shape[0]
|
|
features[i, :n] = f.to(device, dtype)
|
|
mask[i, :n] = 1
|
|
return features, mask
|
|
|
|
|
|
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
|
|
|
|
# Ideogram uses asymmetric CFG: the unconditional branch is image-only
|
|
# (no text tokens) with zeroed text features -- it does NOT run a negative
|
|
# prompt through the text encoder. So we ignore unconditional_embeds and
|
|
# build an empty (0-length) text sequence for the uncond pass below.
|
|
do_cfg = 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, cond_mask = pad_text_features(
|
|
conditional_embeds.text_embeds, device, dtype
|
|
)
|
|
if do_cfg:
|
|
# Image-only unconditional: zero-length text sequence. predict_velocity
|
|
# then produces an image-token-only forward pass with zeroed llm
|
|
# features, matching the reference's asymmetric CFG.
|
|
batch_size = latents.shape[0]
|
|
text_dim = cond_feats.shape[-1]
|
|
uncond_feats = torch.zeros(
|
|
batch_size, 0, text_dim, device=device, dtype=dtype
|
|
)
|
|
uncond_mask = torch.zeros(batch_size, 0, dtype=torch.long, device=device)
|
|
|
|
# The unconditional LoRA (if present) must be active *only* on the
|
|
# unconditional pass. We force it off before each conditional pass since the
|
|
# outer sampling context (``with network:``) may switch it on globally.
|
|
uncond_lora = getattr(model, "unconditional_lora", None)
|
|
|
|
for t in timesteps:
|
|
t01 = (t / 1000.0).to(device).expand(latents.shape[0])
|
|
if uncond_lora is not None:
|
|
uncond_lora.is_active = False
|
|
v_cond = predict_velocity(
|
|
transformer, latents.to(dtype), t01, cond_feats, cond_mask
|
|
)
|
|
if do_cfg:
|
|
if uncond_lora is not None:
|
|
uncond_lora.is_active = True
|
|
try:
|
|
v_uncond = predict_velocity(
|
|
transformer, latents.to(dtype), t01, uncond_feats, uncond_mask
|
|
)
|
|
finally:
|
|
if uncond_lora is not None:
|
|
uncond_lora.is_active = False
|
|
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]
|