refactor(ideogram4): adopt diffusers-native pipeline

diffusers shipped first-party Ideogram 4 (transformer + pipeline) in
9b0818cf, so drop the in-tree port and keep only SD.Next integration glue.

Bump the diffusers pin to 9b0818cf and build diffusers' Ideogram4Pipeline
from a thin loader with per-transformer SDNQ. A small subclass keeps the
text encoder resident for the Qwen3-VL tap under balanced offload, and the
step callback denormalizes the preview latent from vae.bn before unpatchify.
Deletes the ported transformer, pipeline, scheduler, text encoder, and
latent-norm constants.
This commit is contained in:
CalamitousFelicitousness
2026-06-05 00:13:16 +01:00
parent 4866d2c8f3
commit bcebb6b81b
12 changed files with 69 additions and 1250 deletions
+1 -1
View File
@@ -533,7 +533,7 @@ def check_diffusers():
t_start = time.time()
if args.skip_all:
return
target_commit = "ed0711878cfed79aba2e4cf0712b2fcd8bead577" # diffusers commit hash == 0.39.0.dev0 == 06-02-2026
target_commit = "9b0818cf87413b4b9ca2501bf49406eed6d881af" # diffusers commit hash == 0.39.0.dev0 == 06-03-2026 (adds Ideogram 4)
# if args.use_rocm or args.use_zluda or args.use_directml:
# sha = '043ab2520f6a19fce78e6e060a68dbc947edb9f9' # lock diffusers versions for now
pkg = package_spec('diffusers')
+20
View File
@@ -173,6 +173,26 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
current_noise_pred = current_noise_pred.view(b, h_patches, w_patches, channels, 2, 2)
current_noise_pred = current_noise_pred.permute(0, 3, 1, 4, 2, 5).reshape(b, channels, h_patches * 2, w_patches * 2)
shared.state.current_noise_pred = current_noise_pred
elif 'Ideogram4' in pipe.__class__.__name__: # packed normalized [B, seq, 128] -> Flux.2 latent space for TAE FLUX.2
latents = kwargs['latents']
if latents.ndim == 3:
b, seq_len, packed_ch = latents.shape
vae_scale = getattr(pipe, 'vae_scale_factor', 8)
patch = getattr(pipe, 'patch_size', 2)
grid_h = getattr(p, 'height', 1024) // (vae_scale * patch)
grid_w = getattr(p, 'width', 1024) // (vae_scale * patch)
if grid_h * grid_w != seq_len: # fallback to square assumption
grid_h = grid_w = int(seq_len ** 0.5)
bn = pipe.vae.bn
mean = bn.running_mean.view(1, 1, -1).to(device=latents.device, dtype=torch.float32)
std = torch.sqrt(bn.running_var + pipe.vae.config.batch_norm_eps).view(1, 1, -1).to(device=latents.device, dtype=torch.float32)
z = latents.float() * std + mean
ae_ch = packed_ch // (patch * patch)
z = z.view(b, grid_h, grid_w, patch, patch, ae_ch).permute(0, 5, 1, 3, 2, 4).reshape(b, ae_ch, grid_h * patch, grid_w * patch)
shared.state.current_latent = z
else:
shared.state.current_latent = latents
shared.state.current_noise_pred = current_noise_pred
else:
shared.state.current_latent = kwargs['latents']
shared.state.current_noise_pred = current_noise_pred
-21
View File
@@ -1,21 +0,0 @@
"""Ideogram 4 architecture support for SD.Next (diffusers-native port).
Importing this package registers the ported classes under the ``diffusers``
namespace using the exact names Ideogram declared in the shipped
``model_index.json`` (``Ideogram4Pipeline`` / ``Ideogram4Transformer2DModel``),
so folder ``_class_name`` resolution and ``shared_items`` lookups find them.
"""
from __future__ import annotations
import diffusers
from pipelines.ideogram4.pipeline_ideogram4 import Ideogram4Pipeline
from pipelines.ideogram4.scheduler_ideogram4 import Ideogram4Scheduler
from pipelines.ideogram4.transformer_ideogram4 import Ideogram4Transformer2DModel
for cls in (Ideogram4Transformer2DModel, Ideogram4Pipeline, Ideogram4Scheduler):
if not hasattr(diffusers, cls.__name__):
setattr(diffusers, cls.__name__, cls)
__all__ = ["Ideogram4Pipeline", "Ideogram4Scheduler", "Ideogram4Transformer2DModel"]
-17
View File
@@ -1,17 +0,0 @@
"""Packed-sequence role indicators and Qwen3-VL tap layers for Ideogram 4.
Text and image latent tokens share one sequence; each token carries a role
indicator, and image grid positions are offset so they never collide with text
token positions in the shared MRoPE space.
"""
from __future__ import annotations
SEQUENCE_PADDING_INDICATOR = -1
OUTPUT_IMAGE_INDICATOR = 2
LLM_TOKEN_INDICATOR = 3
IMAGE_POSITION_OFFSET = 65536 # keeps image grid positions clear of text token indices
# Qwen3-VL hidden-state layers whose outputs are concatenated and fed to the DiT.
QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
-281
View File
@@ -1,281 +0,0 @@
"""Per-channel latent normalization constants for Ideogram 4.
The DiT works in a normalized latent space; decoding undoes that normalization
before the VAE decoder via ``z * scale + shift``. These 128-dim per-channel
constants are the reference values (they are not the VAE BatchNorm running stats,
which live in a different channel space).
"""
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
-264
View File
@@ -1,264 +0,0 @@
"""Ideogram 4 text-to-image pipeline (diffusers-native, SD.Next integration).
Owns its sampling loop: a packed text+image sequence is denoised with asymmetric
dual-branch classifier-free guidance over two transformers (conditional and a
separately-trained unconditional tower), then the image latents are denormalized,
unpatchified, and decoded by the VAE.
The ``__call__`` signature names every argument SD.Next forwards (it prunes kwargs
against the signature) and drives the diffusers ``callback_on_step_end`` contract
for progress, interrupt, and preview.
"""
from __future__ import annotations
import torch
from PIL import Image
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from pipelines.ideogram4.constants import (
IMAGE_POSITION_OFFSET,
LLM_TOKEN_INDICATOR,
OUTPUT_IMAGE_INDICATOR,
SEQUENCE_PADDING_INDICATOR,
)
from pipelines.ideogram4.latent_norm import get_latent_norm
from pipelines.ideogram4.scheduler_ideogram4 import (
get_schedule_for_resolution,
make_step_intervals,
)
from pipelines.ideogram4.text_encoder_ideogram4 import encode_text, tokenize
class Ideogram4Pipeline(DiffusionPipeline):
"""Ideogram 4 flow-matching text-to-image pipeline."""
def __init__(self, transformer, unconditional_transformer, text_encoder, tokenizer, vae, scheduler) -> None:
super().__init__()
self.register_modules(
transformer=transformer,
unconditional_transformer=unconditional_transformer,
text_encoder=text_encoder,
tokenizer=tokenizer,
vae=vae,
scheduler=scheduler,
)
self.patch_size = 2
self.ae_scale_factor = 8
self.max_text_tokens = 2048
self.latent_shift = None
self.latent_scale = None
self._num_timesteps = 0
@property
def num_timesteps(self) -> int:
return self._num_timesteps
def latent_norm(self, device, dtype) -> tuple[torch.Tensor, torch.Tensor]:
if self.latent_shift is None:
self.latent_shift, self.latent_scale = get_latent_norm()
return self.latent_shift.to(device=device, dtype=dtype), self.latent_scale.to(device=device, dtype=dtype)
def build_inputs(self, prompts: list[str], height: int, width: int, device) -> dict:
"""Build the packed (text tokens + image latent tokens) sequence for one batch."""
tokenized = [tokenize(self.tokenizer, p, self.max_text_tokens) for p in prompts]
batch_size = len(prompts)
patch = self.patch_size * self.ae_scale_factor
if height % patch != 0 or width % patch != 0:
raise ValueError(f"height/width must be divisible by patch_size*ae_scale_factor={patch}")
grid_h = height // patch
grid_w = width // patch
num_image_tokens = grid_h * grid_w
max_text_tokens = max(num_text for _, num_text in tokenized)
total_seq_len = max_text_tokens + num_image_tokens
# Image position ids (t=0, h, w), offset to stay disjoint from text positions.
h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1)
w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1)
t_idx = torch.zeros_like(h_idx)
image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET
token_ids = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
text_position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
segment_ids = torch.full((batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long)
indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
for b, (toks, num_text) in enumerate(tokenized):
pad_len = max_text_tokens - num_text
total_unpadded = num_text + num_image_tokens
offset = pad_len # layout: [pad] [text] [image]
token_ids[b, offset : offset + num_text] = toks
text_pos = torch.arange(num_text)
text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=1)
text_position_ids[b, offset : offset + num_text] = text_pos_3d
position_ids[b, offset : offset + num_text] = text_pos_3d
position_ids[b, offset + num_text :] = image_pos
indicator[b, offset : offset + num_text] = LLM_TOKEN_INDICATOR
indicator[b, offset + num_text :] = OUTPUT_IMAGE_INDICATOR
segment_ids[b, offset : offset + total_unpadded] = 1
return {
"token_ids": token_ids.to(device),
"text_position_ids": text_position_ids.to(device),
"position_ids": position_ids.to(device),
"segment_ids": segment_ids.to(device),
"indicator": indicator.to(device),
"num_image_tokens": num_image_tokens,
"grid_h": grid_h,
"grid_w": grid_w,
"max_text_tokens": max_text_tokens,
}
def init_noise(self, batch_size: int, num_image_tokens: int, latent_dim: int, generator, device) -> torch.Tensor:
if isinstance(generator, list):
samples = [
torch.randn((1, num_image_tokens, latent_dim), generator=generator[b % len(generator)], device=device, dtype=torch.float32)
for b in range(batch_size)
]
return torch.cat(samples, dim=0)
return torch.randn((batch_size, num_image_tokens, latent_dim), generator=generator, device=device, dtype=torch.float32)
def resolve_sampling(self, num_inference_steps: int, guidance_scale: float, device):
"""Map SD.Next's steps + CFG scale to a flat per-step guidance schedule."""
num_steps = int(num_inference_steps)
guidance_schedule = torch.full((num_steps,), float(guidance_scale), dtype=torch.float32, device=device)
return num_steps, guidance_schedule
@torch.no_grad()
def __call__(
self,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
num_inference_steps: int = 20,
guidance_scale: float = 7.0,
width: int = 1024,
height: int = 1024,
generator=None,
output_type: str = "pil",
return_dict: bool = True,
callback_on_step_end=None,
callback_on_step_end_tensor_inputs=None,
**kwargs,
):
device = self._execution_device
if prompt is None:
prompts = [""]
elif isinstance(prompt, str):
prompts = [prompt]
else:
prompts = list(prompt)
batch_size = len(prompts)
num_steps, gw_per_step = self.resolve_sampling(num_inference_steps, guidance_scale, device)
schedule = get_schedule_for_resolution((height, width), known_mean=self.scheduler.config.mu, std=self.scheduler.config.std)
step_intervals = make_step_intervals(num_steps).to(device)
inputs = self.build_inputs(prompts, height=height, width=width, device=device)
num_image_tokens = inputs["num_image_tokens"]
grid_h, grid_w = inputs["grid_h"], inputs["grid_w"]
max_text_tokens = inputs["max_text_tokens"]
latent_dim = self.transformer.config.in_channels
# The tapped forward bypasses the offload hook, so move the encoder on-device, then free it after.
from modules import devices, shared
self.text_encoder.to(device)
llm_features = encode_text(self.text_encoder, inputs["token_ids"], inputs["text_position_ids"], inputs["indicator"])
self.text_encoder.to(devices.cpu)
# At guidance 1.0 the unconditional velocity has zero weight, so skip the second
# tower; only build the negative branch when some step needs it.
gw_values = gw_per_step.tolist()
use_cfg = any(gw != 1.0 for gw in gw_values)
neg_position_ids = neg_segment_ids = neg_indicator = neg_llm_features = None
if use_cfg:
# Negative branch is image-only (asymmetric CFG) with zeroed conditioning.
neg_position_ids = inputs["position_ids"][:, max_text_tokens:]
neg_segment_ids = inputs["segment_ids"][:, max_text_tokens:]
neg_indicator = inputs["indicator"][:, max_text_tokens:]
neg_llm_features = torch.zeros(batch_size, num_image_tokens, llm_features.shape[-1], dtype=llm_features.dtype, device=device)
z = self.init_noise(batch_size, num_image_tokens, latent_dim, generator, device)
text_z_padding = torch.zeros(batch_size, max_text_tokens, latent_dim, dtype=torch.float32, device=device)
self._num_timesteps = num_steps
with self.progress_bar(total=num_steps) as progress_bar:
for i in range(num_steps - 1, -1, -1):
t_val = float(schedule(step_intervals[i + 1].unsqueeze(0)).item())
s_val = float(schedule(step_intervals[i].unsqueeze(0)).item())
t = torch.full((batch_size,), t_val, dtype=torch.float32, device=device)
pos_z = torch.cat([text_z_padding, z], dim=1)
pos_out = self.transformer(
llm_features=llm_features,
x=pos_z,
t=t,
position_ids=inputs["position_ids"],
segment_ids=inputs["segment_ids"],
indicator=inputs["indicator"],
)
pos_v = pos_out[:, max_text_tokens:]
gw_i = gw_values[i]
if gw_i == 1.0:
v = pos_v # unconditional weight is zero; skip the second tower
else:
neg_v = self.unconditional_transformer(
llm_features=neg_llm_features,
x=z,
t=t,
position_ids=neg_position_ids,
segment_ids=neg_segment_ids,
indicator=neg_indicator,
)
v = gw_i * pos_v + (1.0 - gw_i) * neg_v
z = z + v * (s_val - t_val)
if callback_on_step_end is not None:
cb = callback_on_step_end(self, num_steps - 1 - i, t_val, {"latents": z})
if isinstance(cb, dict):
z = cb.get("latents", z)
# Live preview: denorm+unpack into Flux.2 latent space for TAE FLUX.2.
shared.state.current_latent = self.denorm_unpack(z, grid_h, grid_w)
progress_bar.update()
if output_type == "latent":
images = z
else:
images = self.decode_latents(z, grid_h=grid_h, grid_w=grid_w, output_type=output_type)
if not return_dict:
return (images,)
return ImagePipelineOutput(images=images)
def denorm_unpack(self, z: torch.Tensor, grid_h: int, grid_w: int) -> torch.Tensor:
"""Denormalize the packed latent and unpatchify to (B, ae_channels, H, W) in VAE space."""
batch_size = z.shape[0]
patch = self.patch_size
shift, scale = self.latent_norm(z.device, torch.float32)
z = z.float() * scale + shift
ae_channels = z.shape[-1] // (patch * patch)
z = z.view(batch_size, grid_h, grid_w, patch, patch, ae_channels)
z = z.permute(0, 5, 1, 3, 2, 4).contiguous()
return z.view(batch_size, ae_channels, grid_h * patch, grid_w * patch)
def decode_latents(self, z: torch.Tensor, grid_h: int, grid_w: int, output_type: str = "pil"):
"""Denormalize, unpatchify, and VAE-decode the image latents."""
z = self.denorm_unpack(z, grid_h, grid_w)
vae_dtype = next((p.dtype for p in self.vae.parameters() if torch.is_floating_point(p)), torch.float32)
decoded = self.vae.decode(z.to(vae_dtype), return_dict=False)[0]
decoded = decoded.float().clamp(-1.0, 1.0)
decoded = ((decoded + 1.0) * 127.5).round().to(torch.uint8)
decoded = decoded.permute(0, 2, 3, 1).cpu().numpy()
if output_type == "np":
return decoded
return [Image.fromarray(arr) for arr in decoded]
@@ -1,68 +0,0 @@
"""Logit-normal flow-matching schedule for Ideogram 4.
Sampling integrates the flow-matching ODE with Euler steps over a logit-normal
timestep schedule whose mean shifts with resolution (more steps at high noise for
larger images). Guidance is a per-step weight applied as an asymmetric blend of
the conditional and unconditional velocity. Ported from the reference
(github.com/ideogram-oss/ideogram4).
"""
from __future__ import annotations
import math
from dataclasses import dataclass
import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
@dataclass(frozen=True)
class LogitNormalSchedule:
mean: float
std: float = 1.0
logsnr_min: float = -15.0
logsnr_max: float = 18.0
def __call__(self, t: torch.Tensor) -> torch.Tensor:
t = t.to(torch.float64)
z = torch.special.ndtri(t)
y = self.mean + self.std * z
t_ = torch.special.expit(y)
t_ = 1 - t_
t_min = 1.0 / (1 + math.exp(0.5 * self.logsnr_max))
t_max = 1.0 / (1 + math.exp(0.5 * self.logsnr_min))
return t_.clamp(t_min, t_max).to(torch.float32)
def get_schedule_for_resolution(
image_resolution: tuple[int, int],
known_resolution: tuple[int, int] = (512, 512),
known_mean: float = 1.0,
std: float = 1.0,
) -> LogitNormalSchedule:
"""Resolution-aware schedule: the mean shifts by half the log pixel-count ratio."""
num_pixels = image_resolution[0] * image_resolution[1]
known_pixels = known_resolution[0] * known_resolution[1]
mean = known_mean + 0.5 * math.log(num_pixels / known_pixels)
return LogitNormalSchedule(mean=mean, std=std)
def make_step_intervals(num_steps: int) -> torch.Tensor:
"""Linear step schedule mapped through the logit-normal schedule at sample time."""
return torch.linspace(0.0, 1.0, num_steps + 1, dtype=torch.float32)
class Ideogram4Scheduler(SchedulerMixin, ConfigMixin):
"""Thin diffusers scheduler holding the logit-normal defaults.
The denoise loop lives in the pipeline (asymmetric dual-branch CFG over a
flat per-step guidance), so this only carries the ``mu``/``std`` that
parameterize the resolution-aware schedule and serves as the registered
``scheduler`` component.
"""
@register_to_config
def __init__(self, mu: float = 0.0, std: float = 1.75) -> None:
pass
@@ -1,95 +0,0 @@
"""Qwen3-VL text conditioning for Ideogram 4.
The prompt (plain text or a structured JSON caption) is wrapped in the Qwen3
chat template and run through the Qwen3-VL language model. Hidden states are
captured from 13 intermediate layers (pre final-norm) and concatenated along the
feature dim, giving the DiT multi-scale semantic features.
v1 feeds the prompt verbatim (no magic-prompt expansion, no caption verifier).
These weights require a structured JSON caption: a plain-text prompt lands
out-of-distribution and the model renders a baked-in "Image blocked by safety
filter" placeholder. The caption schema is a ``compositional_deconstruction``
object (with ``background`` + ``elements``); see the upstream prompting guide.
"""
from __future__ import annotations
import torch
from transformers.masking_utils import create_causal_mask
from pipelines.ideogram4.constants import LLM_TOKEN_INDICATOR, QWEN3_VL_ACTIVATION_LAYERS
def tokenize(tokenizer, prompt: str, max_text_tokens: int) -> tuple[torch.Tensor, int]:
"""Chat-template tokenize a single prompt (passed verbatim)."""
messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
encoded = tokenizer(text, return_tensors="pt", add_special_tokens=False)
token_ids = encoded["input_ids"][0]
num_text_tokens = int(token_ids.shape[0])
if num_text_tokens > max_text_tokens:
raise ValueError(f"prompt has {num_text_tokens} tokens, exceeds max_text_tokens={max_text_tokens}")
return token_ids, num_text_tokens
def qwen3_vl_layer_features(text_encoder, token_ids: torch.Tensor, attention_mask: torch.Tensor, pos_2d: torch.Tensor) -> list[torch.Tensor]:
"""Run the Qwen3-VL language model and return hidden states at the tap layers.
The layer loop is driven manually so the tapped states are captured before the
model's final norm, matching the reference.
"""
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
return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
def encode_text(text_encoder, token_ids: torch.Tensor, text_position_ids: torch.Tensor, indicator: torch.Tensor) -> torch.Tensor:
"""Stack the tap-layer hidden states into (B, L, hidden_size * num_taps) float32.
Non-LLM positions (left padding / image slots) are zeroed so the DiT only sees
real text features at LLM_TOKEN_INDICATOR positions.
"""
batch_size, seq_len = token_ids.shape
attention_mask = (indicator == LLM_TOKEN_INDICATOR).to(torch.long)
pos_2d = text_position_ids[..., 0].contiguous()
with torch.no_grad():
selected = qwen3_vl_layer_features(text_encoder, token_ids, attention_mask, pos_2d)
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) # (B, L, H * num_taps)
text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
stacked = stacked * text_mask
return stacked.to(torch.float32)
@@ -1,330 +0,0 @@
"""Ideogram 4 flow-matching DiT transformer.
A single-stream Diffusion Transformer: Qwen3-VL text features and noisy image
latent tokens are concatenated into one sequence and processed by 34 shared
blocks (QK-RMSNorm attention, SwiGLU MLP, tanh-gated AdaLN), with 3D multimodal
RoPE giving text and image tokens a unified positional space. The model predicts
a flow-matching velocity on the image tokens.
Ported from the reference implementation (github.com/ideogram-oss/ideogram4) as a
diffusers ModelMixin so it loads from the shipped diffusers-layout checkpoint and
can be SDNQ-quantized at load. Parameter names match the reference state dict
1:1, so no key remapping is needed.
"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from pipelines.ideogram4.constants import LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR
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
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 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 (t, h, w).
assert position_ids.ndim == 3 and position_ids.shape[-1] == 3
batch_size, _, _ = position_ids.shape
pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32) # (3, B, L)
inv_freq = self.inv_freq.to(dtype=torch.float32)[None, None, :, None].expand(3, batch_size, -1, 1)
freqs = inv_freq @ pos.unsqueeze(2) # (3, B, F, L)
freqs = freqs.transpose(2, 3) # (3, B, L, F)
# 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, segment_ids: 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)
# Block-diagonal mask from segment ids: (B, 1, L, L), True = attend.
attn_mask = (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1)
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, adaln_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(adaln_dim, 4 * hidden_size, bias=True)
def forward(self, x: torch.Tensor, segment_ids: 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, segment_ids=segment_ids, 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
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 holds a scalar per token; keep its (float) dtype as the compute dtype so
# SDNQ-quantized Linears (whose .weight is int) don't drive an int cast.
compute_dtype = x.dtype if torch.is_floating_point(x) else torch.float32
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(getattr(self.mlp_in, "compute_dtype", None) or compute_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, adaln_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(adaln_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(ModelMixin, ConfigMixin):
"""Ideogram 4 flow-matching transformer (single-stream DiT)."""
_no_split_modules = ["Ideogram4TransformerBlock"]
_supports_gradient_checkpointing = False
@register_to_config
def __init__(
self,
num_attention_heads: int = 18,
attention_head_dim: int = 256,
num_layers: int = 34,
intermediate_size: int = 12288,
adaln_dim: int = 512,
in_channels: int = 128,
llm_features_dim: int = 53248,
mrope_section: tuple[int, ...] = (24, 20, 20),
rope_theta: int = 5_000_000,
norm_eps: float = 1e-5,
) -> None:
super().__init__()
emb_dim = num_attention_heads * attention_head_dim
head_dim = attention_head_dim
self.num_heads = num_attention_heads
self.emb_dim = emb_dim
self.input_proj = nn.Linear(in_channels, emb_dim, bias=True)
self.llm_cond_norm = Ideogram4RMSNorm(llm_features_dim, eps=1e-6)
self.llm_cond_proj = nn.Linear(llm_features_dim, emb_dim, bias=True)
self.t_embedding = Ideogram4EmbedScalar(emb_dim, input_range=(0.0, 1.0))
self.adaln_proj = nn.Linear(emb_dim, adaln_dim, bias=True)
self.embed_image_indicator = nn.Embedding(2, emb_dim)
self.rotary_emb = Ideogram4MRoPE(head_dim=head_dim, base=rope_theta, mrope_section=tuple(mrope_section))
self.layers = nn.ModuleList(
[
Ideogram4TransformerBlock(
hidden_size=emb_dim,
intermediate_size=intermediate_size,
num_heads=num_attention_heads,
norm_eps=norm_eps,
adaln_dim=adaln_dim,
)
for _ in range(num_layers)
]
)
self.final_layer = Ideogram4FinalLayer(hidden_size=emb_dim, out_channels=in_channels, adaln_dim=adaln_dim)
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 / OUTPUT_IMAGE_INDICATOR).
Returns:
(B, L, in_channels) velocity in float32; only OUTPUT_IMAGE_INDICATOR
positions are meaningful.
"""
_, _, in_channels = x.shape
assert in_channels == self.config["in_channels"]
# Compute dtype from a norm weight (never SDNQ-quantized), honoring an
# explicit Fp8Linear compute_dtype if one is present.
param_dtype = getattr(self.input_proj, "compute_dtype", None)
if param_dtype is None:
w = self.input_proj.weight
param_dtype = w.dtype if torch.is_floating_point(w) else self.llm_cond_norm.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 the adaln 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)
for layer in self.layers:
h = layer(h, segment_ids=segment_ids, cos=cos, sin=sin, adaln_input=adaln_input)
out = self.final_layer(h, c=adaln_input)
return out.to(torch.float32)
+29 -13
View File
@@ -4,9 +4,23 @@ from transformers.models.qwen3_vl import Qwen3VLModel
from modules import shared, devices, sd_models
from modules.logger import log
from pipelines import generic
from pipelines.ideogram4.pipeline_ideogram4 import Ideogram4Pipeline
from pipelines.ideogram4.scheduler_ideogram4 import Ideogram4Scheduler
from pipelines.ideogram4.transformer_ideogram4 import Ideogram4Transformer2DModel
class Ideogram4Pipeline(diffusers.Ideogram4Pipeline):
"""SD.Next integration subclass for the diffusers-native Ideogram 4 pipeline.
``encode_prompt`` drives the Qwen3-VL tap by calling ``language_model`` submodules
directly, which bypasses the balanced-offload pre-forward hook. Move the encoder
on-device for the tap and release it afterward so it does not pin VRAM.
"""
def encode_prompt(self, *args, **kwargs):
self.text_encoder.to(self._execution_device)
try:
return super().encode_prompt(*args, **kwargs)
finally:
if shared.opts.diffusers_offload_mode != 'none':
self.text_encoder.to(devices.cpu)
def load_ideogram4(checkpoint_info, diffusers_load_config=None):
@@ -19,26 +33,28 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None):
if repo_id is None or repo_id.lower() == 'none':
return None
# Each transformer loads independently from its subfolder.
transformer = generic.load_transformer(repo_id, cls_name=Ideogram4Transformer2DModel, subfolder="transformer", load_config=diffusers_load_config)
unconditional_transformer = generic.load_transformer(repo_id, cls_name=Ideogram4Transformer2DModel, subfolder="unconditional_transformer", load_config=diffusers_load_config)
# Each transformer loads independently from its subfolder, so each gets its own SDNQ config.
cls = diffusers.Ideogram4Transformer2DModel
transformer = generic.load_transformer(repo_id, cls_name=cls, subfolder="transformer", load_config=diffusers_load_config)
unconditional_transformer = generic.load_transformer(repo_id, cls_name=cls, subfolder="unconditional_transformer", load_config=diffusers_load_config)
# shared_te_map redirects to the shared Qwen3-VL repo (deduped with VQA + prompt-enhance);
# the bundled text_encoder is the fallback when sharing is off.
text_encoder = generic.load_text_encoder(repo_id, cls_name=Qwen3VLModel, load_config=diffusers_load_config)
tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder="tokenizer", cache_dir=shared.opts.diffusers_dir)
vae = diffusers.AutoencoderKLFlux2.from_pretrained(repo_id, subfolder="vae", cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype)
scheduler = Ideogram4Scheduler()
scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler", cache_dir=shared.opts.diffusers_dir)
pipe = Ideogram4Pipeline(
transformer=transformer,
unconditional_transformer=unconditional_transformer,
scheduler=scheduler,
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
vae=vae,
scheduler=scheduler,
transformer=transformer,
unconditional_transformer=unconditional_transformer,
)
# The pipeline decodes the packed latent itself, so keep SD.Next from re-decoding it.
pipe.task_args = {'output_type': 'pil'}
# The pipeline decodes internally; the CFG scale slider drives guidance_scale, which is
# mutually exclusive with the pipeline's default per-step guidance_schedule.
pipe.task_args = {'output_type': 'pil', 'guidance_schedule': None} # pylint: disable=attribute-defined-outside-init
del transformer, unconditional_transformer, text_encoder, vae
devices.torch_gc(force=True, reason='load')
-138
View File
@@ -1,138 +0,0 @@
#!/usr/bin/env python
"""Numerical parity gate for the ported Ideogram 4 transformer.
Validates that ``pipelines.ideogram4.transformer_ideogram4.Ideogram4Transformer2DModel``
reproduces the upstream reference (github.com/ideogram-oss/ideogram4) exactly:
identical parameter names/shapes (checked via ``load_state_dict``) and identical
forward outputs on fixed inputs. A small config is used so it runs on CPU in
seconds rather than instantiating the real 9B model.
The upstream reference module is fetched at run time from a pinned commit into a
temp dir and imported as the parity oracle; the test SKIPS cleanly when it cannot
be fetched (offline). Run from the repo root:
python test/test-ideogram4-parity.py
"""
from __future__ import annotations
import importlib
import os
import sys
import tempfile
import urllib.request
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if REPO_ROOT not in sys.path:
sys.path.insert(0, REPO_ROOT)
import torch
REF_COMMIT = "19fc3af67fd7a98b7accf0844e50eda50af9bdc9"
REF_BASE = f"https://raw.githubusercontent.com/ideogram-oss/ideogram4/{REF_COMMIT}/src/ideogram4"
REF_FILES = ("constants.py", "modeling_ideogram4.py")
# Small config (both reference and port use it) so parity runs on CPU in seconds.
SMALL = {
"num_attention_heads": 4,
"attention_head_dim": 16,
"num_layers": 2,
"intermediate_size": 128,
"adaln_dim": 32,
"in_channels": 16,
"llm_features_dim": 48,
"mrope_section": (2, 2, 2),
"rope_theta": 5_000_000,
"norm_eps": 1e-5,
}
def fetch_reference():
"""Download the pinned upstream modeling module into a temp package and import it."""
tmp = tempfile.mkdtemp(prefix="ideogram4_ref_")
pkg = os.path.join(tmp, "ideogram4_ref")
os.makedirs(pkg, exist_ok=True)
with open(os.path.join(pkg, "__init__.py"), "w", encoding="utf8"):
pass
for fn in REF_FILES:
with urllib.request.urlopen(f"{REF_BASE}/{fn}", timeout=30) as resp:
data = resp.read().decode("utf8")
# upstream imports `from ideogram4.constants import ...`; point it at the temp package
data = data.replace("from ideogram4.constants", "from ideogram4_ref.constants")
with open(os.path.join(pkg, fn), "w", encoding="utf8") as f:
f.write(data)
sys.path.insert(0, tmp)
return importlib.import_module("ideogram4_ref.modeling_ideogram4")
def check_latent_norm() -> None:
"""Guard the latent-norm constants (offline): they are the reference values, not VAE BatchNorm stats."""
from pipelines.ideogram4.latent_norm import get_latent_norm
shift, scale = get_latent_norm()
assert shift.shape == (128,) and scale.shape == (128,)
ref_shift = torch.tensor([0.01984364, 0.10149707, 0.29689495, 0.27188619])
ref_scale = torch.tensor([1.63933691, 1.70204478, 1.73642566, 1.90004803])
assert torch.allclose(shift[:4], ref_shift, atol=1e-6), f"latent shift regressed: {shift[:4].tolist()}"
assert torch.allclose(scale[:4], ref_scale, atol=1e-6), f"latent scale regressed: {scale[:4].tolist()}"
print("PASS: latent-norm constants match reference")
def main() -> int:
check_latent_norm()
try:
ref = fetch_reference()
except Exception as e:
print(f"SKIP: could not fetch upstream reference ({e})")
return 0
from pipelines.ideogram4.transformer_ideogram4 import Ideogram4Transformer2DModel
torch.manual_seed(0)
ref_cfg = ref.Ideogram4Config(
emb_dim=SMALL["num_attention_heads"] * SMALL["attention_head_dim"],
num_layers=SMALL["num_layers"],
num_heads=SMALL["num_attention_heads"],
intermediate_size=SMALL["intermediate_size"],
adanln_dim=SMALL["adaln_dim"],
in_channels=SMALL["in_channels"],
llm_features_dim=SMALL["llm_features_dim"],
rope_theta=SMALL["rope_theta"],
mrope_section=SMALL["mrope_section"],
norm_eps=SMALL["norm_eps"],
)
ref_model = ref.Ideogram4Transformer(ref_cfg).eval()
mine = Ideogram4Transformer2DModel(**SMALL).eval()
# 1. structural parity: names + shapes must match 1:1 (the shipped checkpoint
# was saved from the reference, so a clean load proves load compatibility).
missing, unexpected = mine.load_state_dict(ref_model.state_dict(), strict=False)
assert not missing, f"missing keys in port: {missing}"
assert not unexpected, f"unexpected keys in port: {unexpected}"
# 2. numerical parity on fixed inputs.
n_text, n_img = 3, 4
seq_len = n_text + n_img
gen = torch.Generator().manual_seed(123)
x = torch.randn(1, seq_len, SMALL["in_channels"], generator=gen)
t = torch.rand(1, generator=gen)
llm = torch.randn(1, seq_len, SMALL["llm_features_dim"], generator=gen)
position_ids = torch.randint(0, 64, (1, seq_len, 3), generator=gen)
segment_ids = torch.ones(1, seq_len, dtype=torch.long)
indicator = torch.tensor([[ref.LLM_TOKEN_INDICATOR] * n_text + [ref.OUTPUT_IMAGE_INDICATOR] * n_img], dtype=torch.long)
kwargs = {"llm_features": llm, "x": x, "t": t, "position_ids": position_ids, "segment_ids": segment_ids, "indicator": indicator}
with torch.no_grad():
out_ref = ref_model(**kwargs)
out_mine = mine(**kwargs)
image_tokens = indicator[0] == ref.OUTPUT_IMAGE_INDICATOR
diff = (out_ref[:, image_tokens] - out_mine[:, image_tokens]).abs().max().item()
print(f"max abs diff (image tokens): {diff:.3e}")
assert diff < 1e-4, f"parity FAILED: max diff {diff}"
print("PASS: ported transformer matches upstream reference")
return 0
if __name__ == "__main__":
sys.exit(main())
+19 -22
View File
@@ -1,15 +1,16 @@
#!/usr/bin/env python
"""Standalone end-to-end smoke for the Ideogram 4 port.
"""Standalone end-to-end smoke for diffusers-native Ideogram 4.
Loads the converted bf16 diffusers folder and quantizes it with SDNQ at load (the
same path SD.Next uses), loads the shared Qwen3-VL text encoder, builds
Ideogram4Pipeline, and generates an image. Exercises the real pipeline: both
transformers (real weights), the Qwen3-VL 13-layer tap, the dual-branch loop, the
logit-normal schedule, latent norm, and VAE decode. SDNQ int4 fits the two towers
plus the encoder on a 24GB GPU.
Loads the split-projection bf16 diffusers folder and quantizes both transformers
with SDNQ at load (the same path SD.Next uses), loads the shared Qwen3-VL text
encoder, builds the diffusers ``Ideogram4Pipeline``, and generates an image. This
exercises the real pipeline: both transformers under SDNQ, the Qwen3-VL 13-layer
tap, the dual-branch asymmetric CFG loop, the logit-normal schedule, the vae.bn
latent denorm, and VAE decode. SDNQ int4 fits the two towers plus the encoder on a
24GB GPU.
Usage:
python test/test-ideogram4-smoke.py --model /path/to/Ideogram-4-bf16 --output out.png
python test/test-ideogram4-smoke.py --model /path/to/Ideogram-4-bf16-split --output out.png
"""
import argparse
@@ -18,7 +19,7 @@ import sys
import time
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, help="converted bf16 diffusers folder")
parser.add_argument("--model", required=True, help="split-projection bf16 diffusers folder")
parser.add_argument("--output", default="ideogram4_smoke.png")
parser.add_argument("--prompt", default="a ginger cat wearing a tiny wizard hat reading a glowing spellbook, detailed digital illustration")
parser.add_argument("--height", type=int, default=1024)
@@ -53,10 +54,6 @@ from transformers.models.qwen3_vl import Qwen3VLModel
from modules import devices
from modules.sdnq import SDNQConfig
from pipelines.ideogram4.pipeline_ideogram4 import Ideogram4Pipeline
from pipelines.ideogram4.scheduler_ideogram4 import Ideogram4Scheduler
from pipelines.ideogram4.transformer_ideogram4 import Ideogram4Transformer2DModel
TE_REPO = "Qwen/Qwen3-VL-8B-Instruct"
@@ -65,30 +62,30 @@ def main() -> int:
cfg = SDNQConfig(weights_dtype=args.weights_dtype)
print(f"loading transformer (sdnq {args.weights_dtype}) ...", flush=True)
transformer = Ideogram4Transformer2DModel.from_pretrained(args.model, subfolder="transformer", quantization_config=cfg, torch_dtype=torch.bfloat16).to(device)
transformer = diffusers.Ideogram4Transformer2DModel.from_pretrained(args.model, subfolder="transformer", quantization_config=cfg, torch_dtype=torch.bfloat16).to(device)
print("loading unconditional_transformer ...", flush=True)
uncond = Ideogram4Transformer2DModel.from_pretrained(args.model, subfolder="unconditional_transformer", quantization_config=cfg, torch_dtype=torch.bfloat16).to(device)
uncond = diffusers.Ideogram4Transformer2DModel.from_pretrained(args.model, subfolder="unconditional_transformer", quantization_config=cfg, torch_dtype=torch.bfloat16).to(device)
print("loading text encoder Qwen3-VL ...", flush=True)
te_kwargs = {"cache_dir": args.hf_cache} if args.hf_cache else {}
text_encoder = Qwen3VLModel.from_pretrained(TE_REPO, quantization_config=SDNQConfig(weights_dtype=args.weights_dtype), torch_dtype=torch.bfloat16, **te_kwargs).to(device)
tokenizer = AutoTokenizer.from_pretrained(args.model, subfolder="tokenizer")
print("loading vae ...", flush=True)
vae = diffusers.AutoencoderKLFlux2.from_pretrained(args.model, subfolder="vae", torch_dtype=torch.bfloat16).to(device)
scheduler = Ideogram4Scheduler()
scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained(args.model, subfolder="scheduler")
pipe = Ideogram4Pipeline(
transformer=transformer,
unconditional_transformer=uncond,
pipe = diffusers.Ideogram4Pipeline(
scheduler=scheduler,
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
vae=vae,
scheduler=scheduler,
transformer=transformer,
unconditional_transformer=uncond,
)
generator = torch.Generator(device=device).manual_seed(args.seed)
print(f"generating {args.width}x{args.height} steps={args.steps} ...", flush=True)
start = time.time()
out = pipe(prompt=args.prompt, num_inference_steps=args.steps, guidance_scale=7.0, width=args.width, height=args.height, generator=generator)
out = pipe(prompt=args.prompt, num_inference_steps=args.steps, guidance_scale=7.0, guidance_schedule=None, width=args.width, height=args.height, generator=generator)
elapsed = time.time() - start
image = out.images[0]