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

2
extensions_built_in/diffusion_models/__init__.py
View File

@@ -13,6 +13,7 @@ from .ernie_image import ErnieImageModel
from .nucleus_image import NucleusImageModel
from .hidream.hidream_o1_model import HidreamO1Model
from .z_image.z_image_l2p_model import ZImageL2PModel
from .ideogram4 import Ideogram4Model
AI_TOOLKIT_MODELS = [
# put a list of models here
@@ -40,4 +41,5 @@ AI_TOOLKIT_MODELS = [
NucleusImageModel,
HidreamO1Model,
ZImageL2PModel,
Ideogram4Model,
]

1
extensions_built_in/diffusion_models/ideogram4/__init__.py Normal file
View File

@@ -0,0 +1 @@
from .ideogram4 import Ideogram4Model

534
extensions_built_in/diffusion_models/ideogram4/ideogram4.py Normal file
View File

@@ -0,0 +1,534 @@
import os
from typing import List, Optional
import torch
import yaml
from safetensors.torch import load_file, save_file
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.models.base_model import BaseModel
from toolkit.basic import flush
from toolkit.print import print_acc
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.accelerator import unwrap_model
from toolkit.metadata import get_meta_for_safetensors
from toolkit.memory_management import MemoryManager
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from optimum.quanto import freeze, QTensor
import huggingface_hub
from huggingface_hub.errors import EntryNotFoundError
from transformers import AutoModel, AutoTokenizer
from .src.transformer import Ideogram4Config, Ideogram4Transformer2DModel
from .src.vae import AutoEncoder, AutoEncoderParams, convert_diffusers_state_dict
from .src.latent_norm import get_latent_norm
from .src.pipeline import (
Ideogram4Pipeline,
get_qwen3_vl_features,
patchify_latents,
predict_velocity,
unpatchify_latents,
)
scheduler_config = {
"base_image_seq_len": 256,
"base_shift": 0.5,
"invert_sigmas": False,
"max_image_seq_len": 4096,
"max_shift": 1.15,
"num_train_timesteps": 1000,
"shift": 1.0,
"shift_terminal": None,
"stochastic_sampling": False,
"time_shift_type": "exponential",
"use_beta_sigmas": False,
"use_dynamic_shifting": False,
"use_exponential_sigmas": False,
"use_karras_sigmas": False,
}
# Weight-only FP8 (e4m3) Linear weights carry a per-output-channel float32 scale
# saved alongside as ``<name>.weight_scale``. Folding it back gives bf16 weights.
FP8_SCALE_SUFFIX = ".weight_scale"
# The text encoder is frozen, stock Qwen3-VL-8B-Instruct.
QWEN3_VL_PATH = "Qwen/Qwen3-VL-8B-Instruct"
HF_TOKEN = os.getenv("HF_TOKEN", None)
def _dequantize_fp8_state_dict(
state_dict: dict,
dtype: torch.dtype,
device: torch.device,
low_vram: bool,
) -> dict:
"""Fold weight-only FP8 scales back into the weights, casting to ``dtype``.
Linear weights stored as float8 with a sibling ``.weight_scale`` are
reconstructed as ``weight_fp8.to(float32) * scale[:, None]``. Everything else
is simply cast to ``dtype`` (non-floating tensors are left untouched). If the
checkpoint isn't quantized this is just a dtype cast.
The fold/cast runs on ``device`` (GPU is much faster than CPU). With
``low_vram=True`` each tensor is moved to ``device``, processed, then moved
back to CPU so the whole bf16 model never sits on the GPU at once; otherwise
the dequantized tensors are left on ``device`` ready to load.
"""
work_device = torch.device(device)
def _finish(t: torch.Tensor) -> torch.Tensor:
return t.to("cpu") if low_vram else t
num_fp8 = sum(1 for k in state_dict if k.endswith(FP8_SCALE_SUFFIX))
if num_fp8 > 0:
print_acc(f" dequantizing {num_fp8} fp8 weights -> {dtype} on {work_device}")
else:
print_acc(f" casting weights -> {dtype} on {work_device}")
out = {}
for key, tensor in state_dict.items():
if key.endswith(FP8_SCALE_SUFFIX):
continue
scale_key = key + "_scale"
if key.endswith(".weight") and scale_key in state_dict:
w = tensor.to(work_device, torch.float32)
scale = state_dict[scale_key].to(work_device, torch.float32)
out[key] = _finish((w * scale.unsqueeze(1)).to(dtype))
elif tensor.is_floating_point():
out[key] = _finish(tensor.to(work_device, dtype))
else:
out[key] = tensor
return out
def _load_component_state_dict(base: str, subfolder: str, basename: str) -> dict:
"""Load a component's weights whether local or on the hub, sharded or single."""
index_name = f"{basename}.safetensors.index.json"
single_name = f"{basename}.safetensors"
# Local directory layout: <base>/<subfolder>/<file>
local_dir = os.path.join(base, subfolder)
if os.path.isdir(local_dir):
index_path = os.path.join(local_dir, index_name)
if os.path.exists(index_path):
return _load_sharded(local_dir, index_path, is_local=True)
return load_file(os.path.join(local_dir, single_name))
# Hub repo layout: <subfolder>/<file>
prefix = f"{subfolder}/" if subfolder else ""
try:
index_path = huggingface_hub.hf_hub_download(
repo_id=base, filename=f"{prefix}{index_name}", token=HF_TOKEN
)
return _load_sharded(base, index_path, is_local=False, prefix=prefix)
except EntryNotFoundError:
single_path = huggingface_hub.hf_hub_download(
repo_id=base, filename=f"{prefix}{single_name}", token=HF_TOKEN
)
return load_file(single_path)
def _load_sharded(base, index_path, is_local, prefix="") -> dict:
import json
with open(index_path) as f:
index = json.load(f)
shard_files = sorted(set(index["weight_map"].values()))
state_dict = {}
num_shards = len(shard_files)
for i, shard in enumerate(shard_files):
if is_local:
shard_path = os.path.join(base, shard)
else:
print_acc(f" downloading shard {i + 1}/{num_shards}: {shard}")
shard_path = huggingface_hub.hf_hub_download(
repo_id=base, filename=f"{prefix}{shard}", token=HF_TOKEN
)
print_acc(f" loading shard {i + 1}/{num_shards}: {shard}")
state_dict.update(load_file(shard_path))
return state_dict
class Ideogram4Model(BaseModel):
arch = "ideogram4"
def __init__(
self,
device,
model_config: ModelConfig,
dtype="bf16",
custom_pipeline=None,
noise_scheduler=None,
**kwargs,
):
super().__init__(
device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
)
self.is_flow_matching = True
self.is_transformer = True
self.target_lora_modules = ["Ideogram4Transformer2DModel"]
self.patch_size = 2
self.vae_scale_factor = 8
# Length the Qwen3-VL caption is padded/truncated to. Fixed so cached and
# batched embeddings always share a sequence length.
self.max_text_length = int(
self.model_config.model_kwargs.get("max_text_length", 512)
)
self._latent_shift = None
self._latent_scale = None
@staticmethod
def get_train_scheduler():
return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
def get_bucket_divisibility(self):
# 8 for the VAE downsample, 2 for the patch size.
return self.vae_scale_factor * self.patch_size
# ------------------------------------------------------------------
# Loading
# ------------------------------------------------------------------
def _load_text_encoder(self, base: str):
dtype = self.torch_dtype
# The text encoder is frozen, stock Qwen3-VL-8B-Instruct. The ideogram repo
# only ships an fp8 copy of it, so load the public bf16 model directly --
# faster and higher precision than dequantizing the fp8 weights.
te_path = self.model_config.model_kwargs.get("text_encoder_path", QWEN3_VL_PATH)
self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN)
text_encoder = AutoModel.from_pretrained(
te_path, torch_dtype=dtype, token=HF_TOKEN
)
flush()
text_encoder.eval()
text_encoder.requires_grad_(False)
return tokenizer, text_encoder
def _load_transformer(self, base: str):
dtype = self.torch_dtype
self.print_and_status_update("Loading transformer")
transformer_config = Ideogram4Config()
with torch.device("meta"):
transformer = Ideogram4Transformer2DModel(transformer_config)
self.print_and_status_update(" - fetching transformer weights")
state_dict = _load_component_state_dict(
base, "transformer", "diffusion_pytorch_model"
)
self.print_and_status_update(" - dequantizing transformer weights")
state_dict = _dequantize_fp8_state_dict(
state_dict, dtype, self.device_torch, self.model_config.low_vram
)
self.print_and_status_update(" - loading transformer state dict")
transformer.load_state_dict(state_dict, assign=True)
del state_dict
flush()
# inv_freq is a non-persistent buffer absent from the checkpoint; rebuild
# it now that the module is off the meta device.
head_dim = transformer_config.emb_dim // transformer_config.num_heads
inv_freq = 1.0 / (
transformer_config.rope_theta
** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
)
transformer.rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)
return transformer
def _load_vae(self, base: str):
dtype = self.torch_dtype
self.print_and_status_update("Loading VAE")
vae_sd = _load_component_state_dict(base, "vae", "diffusion_pytorch_model")
vae_sd = convert_diffusers_state_dict(vae_sd)
vae = AutoEncoder(AutoEncoderParams())
vae.load_state_dict(vae_sd)
del vae_sd
vae.to(self.vae_device_torch, dtype=dtype)
vae.eval()
vae.requires_grad_(False)
return vae
def load_model(self):
dtype = self.torch_dtype
self.print_and_status_update("Loading Ideogram4 model")
base = self.model_config.name_or_path
transformer = self._load_transformer(base)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
else:
transformer.to(self.device_torch, dtype=dtype)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[transformer.input_proj, transformer.llm_cond_proj],
)
elif self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
else:
# quantize_model leaves the model on CPU; make sure it lands on device.
transformer.to(self.device_torch)
flush()
tokenizer, text_encoder = self._load_text_encoder(base)
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
text_encoder.to(self.device_torch)
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
elif self.model_config.low_vram:
self.print_and_status_update("Moving text encoder to CPU")
text_encoder.to("cpu")
else:
self.print_and_status_update("Moving text encoder to device")
text_encoder.to(self.device_torch)
flush()
vae = self._load_vae(base)
self.noise_scheduler = Ideogram4Model.get_train_scheduler()
shift, scale = get_latent_norm()
self._latent_shift = shift.view(1, -1, 1, 1)
self._latent_scale = scale.view(1, -1, 1, 1)
self.vae = vae
self.text_encoder = text_encoder
self.tokenizer = tokenizer
self.model = transformer
self.pipeline = Ideogram4Pipeline(self)
self.print_and_status_update("Model Loaded")
# ------------------------------------------------------------------
# Generation
# ------------------------------------------------------------------
def get_generation_pipeline(self):
return Ideogram4Pipeline(self)
def generate_single_image(
self,
pipeline: Ideogram4Pipeline,
gen_config: GenerateImageConfig,
conditional_embeds: PromptEmbeds,
unconditional_embeds: PromptEmbeds,
generator: torch.Generator,
extra: dict,
):
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
sc = self.get_bucket_divisibility()
gen_config.width = int(gen_config.width // sc * sc)
gen_config.height = int(gen_config.height // sc * sc)
img = pipeline(
conditional_embeds=conditional_embeds,
unconditional_embeds=unconditional_embeds,
height=gen_config.height,
width=gen_config.width,
num_inference_steps=gen_config.num_inference_steps,
guidance_scale=gen_config.guidance_scale,
latents=gen_config.latents,
generator=generator,
)[0]
return img
# ------------------------------------------------------------------
# Training hooks
# ------------------------------------------------------------------
def get_noise_prediction(
self,
latent_model_input: torch.Tensor, # (B, 128, gh, gw)
timestep: torch.Tensor, # 0 to 1000 scale
text_embeddings: PromptEmbeds,
**kwargs,
):
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
t01 = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0
if t01.dim() == 0:
t01 = t01.unsqueeze(0)
if t01.shape[0] != latent_model_input.shape[0]:
t01 = t01.expand(latent_model_input.shape[0])
llm_features = text_embeddings.text_embeds.to(self.device_torch)
text_mask = text_embeddings.attention_mask.to(self.device_torch)
pred = predict_velocity(
self.transformer,
latent_model_input.to(self.device_torch),
t01,
llm_features,
text_mask,
)
return pred
def get_prompt_embeds(self, prompt) -> PromptEmbeds:
if isinstance(prompt, str):
prompt = [prompt]
if self.text_encoder.device == torch.device("cpu"):
self.text_encoder.to(self.device_torch)
device = self.text_encoder.device
pad_id = self.tokenizer.pad_token_id
if pad_id is None:
pad_id = self.tokenizer.eos_token_id or 0
token_id_list = []
for p in prompt:
messages = [{"role": "user", "content": [{"type": "text", "text": p}]}]
text = self.tokenizer.apply_chat_template(
messages, add_generation_prompt=True, tokenize=False
)
ids = self.tokenizer(
text,
add_special_tokens=False,
truncation=True,
max_length=self.max_text_length,
)["input_ids"]
token_id_list.append(ids)
seq_len = self.max_text_length
batch_size = len(token_id_list)
token_ids = torch.full((batch_size, seq_len), pad_id, dtype=torch.long)
attention_mask = torch.zeros((batch_size, seq_len), dtype=torch.long)
for b, ids in enumerate(token_id_list):
n = min(len(ids), seq_len)
# left pad
token_ids[b, seq_len - n :] = torch.tensor(ids[:n], dtype=torch.long)
attention_mask[b, seq_len - n :] = 1
token_ids = token_ids.to(device)
attention_mask = attention_mask.to(device)
pos_2d = (attention_mask.cumsum(dim=-1) - 1).clamp(min=0).to(torch.long)
features = get_qwen3_vl_features(
self.text_encoder, token_ids, attention_mask, pos_2d
)
pe = PromptEmbeds(features.to(self.torch_dtype))
pe.attention_mask = attention_mask
return pe
def get_model_has_grad(self):
return False
def get_te_has_grad(self):
return False
# ------------------------------------------------------------------
# VAE
# ------------------------------------------------------------------
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
if self.vae.device == torch.device("cpu"):
self.vae.to(self.vae_device_torch)
if isinstance(image_list, list):
images = torch.stack(image_list, dim=0)
else:
images = image_list
images = images.to(device, dtype=dtype)
ae_channels = self.vae.params.z_channels
moments = self.vae.encoder(images)
mean = moments[:, :ae_channels]
patched = patchify_latents(mean, self.patch_size)
shift = self._latent_shift.to(patched.device, patched.dtype)
scale = self._latent_scale.to(patched.device, patched.dtype)
latents = (patched - shift) / scale
return latents.to(device, dtype=dtype)
def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
if self.vae.device == torch.device("cpu"):
self.vae.to(self.vae_device_torch)
latents = latents.to(device, dtype=dtype)
shift = self._latent_shift.to(device, dtype)
scale = self._latent_scale.to(device, dtype)
patched = latents * scale + shift
z = unpatchify_latents(patched, self.patch_size)
images = self.vae.decoder(z)
return images
# ------------------------------------------------------------------
# Saving / misc
# ------------------------------------------------------------------
def get_loss_target(self, *args, **kwargs):
noise = kwargs.get("noise")
batch = kwargs.get("batch")
return (noise - batch.latents).detach()
def save_model(self, output_path, meta, save_dtype):
if not output_path.endswith(".safetensors"):
output_path = output_path + ".safetensors"
transformer: Ideogram4Transformer2DModel = unwrap_model(self.model)
state_dict = transformer.state_dict()
save_dict = {}
for k, v in state_dict.items():
if isinstance(v, QTensor):
v = v.dequantize()
save_dict[k] = v.clone().to("cpu", dtype=save_dtype)
meta = get_meta_for_safetensors(meta, name="ideogram4")
save_file(save_dict, output_path, metadata=meta)
def get_base_model_version(self):
return "ideogram4"
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["layers"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
return new_sd
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd

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