From 2f3d0e719db1829eae37d15e6a5a9bbbacd0c4cd Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 4 Jun 2026 01:00:01 +0100 Subject: [PATCH 1/7] feat(ideogram4): add Ideogram 4 model support Diffusers-native port of the 9.3B flow-matching DiT: dual-transformer asymmetric CFG, a 13-layer Qwen3-VL tap encoder deduped with VQA and prompt-enhance, the Flux.2 VAE, and a logit-normal schedule. Loads a published bf16 repo with SDNQ at load. --- CHANGELOG.md | 4 + data/reference.json | 8 + modules/modeldata.py | 2 + modules/sd_detect.py | 2 + modules/sd_models.py | 4 + modules/sd_samplers_common.py | 2 +- modules/shared_items.py | 1 + pipelines/generic_shared.py | 6 + pipelines/ideogram4/__init__.py | 21 ++ pipelines/ideogram4/constants.py | 17 + pipelines/ideogram4/latent_norm.py | 281 +++++++++++++++ pipelines/ideogram4/pipeline_ideogram4.py | 255 ++++++++++++++ pipelines/ideogram4/scheduler_ideogram4.py | 68 ++++ pipelines/ideogram4/text_encoder_ideogram4.py | 95 +++++ pipelines/ideogram4/transformer_ideogram4.py | 330 ++++++++++++++++++ pipelines/model_ideogram4.py | 45 +++ 16 files changed, 1140 insertions(+), 1 deletion(-) create mode 100644 pipelines/ideogram4/__init__.py create mode 100644 pipelines/ideogram4/constants.py create mode 100644 pipelines/ideogram4/latent_norm.py create mode 100644 pipelines/ideogram4/pipeline_ideogram4.py create mode 100644 pipelines/ideogram4/scheduler_ideogram4.py create mode 100644 pipelines/ideogram4/text_encoder_ideogram4.py create mode 100644 pipelines/ideogram4/transformer_ideogram4.py create mode 100644 pipelines/model_ideogram4.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a61a55e9..43175fdc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,10 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas oh, that 12B encoder is MoE with 3.6B activated plus its prequantized using `mxfp4` *note* Lens comes with its own prompt-refiner, enable in settings -> model options (disabled by default) *note* original Lens implements only text-2-image, SD.Next adds image-2-image and inpaint workflows as well + - [Ideogram 4](https://huggingface.co/ideogram-ai/ideogram-4) open-weight 9.3B flow-matching single-stream DiT + Qwen3-VL-8B text encoder (shared and deduped) and Flux2 VAE, with dual-transformer asymmetric CFG + converted to a bf16-Diffusers repo with SDNQ-at-load + *note* requires structured JSON-caption prompts, a plain-text prompt returns the model's built-in safety placeholder - **Features** - **SDNQ** new quantization algorithm: *Hadamard Rotations* much higher quality than base SDNQ, but runs slightly slower diff --git a/data/reference.json b/data/reference.json index b28af9c3d..ebef80479 100644 --- a/data/reference.json +++ b/data/reference.json @@ -184,6 +184,14 @@ "size": 20.3, "date": "2025 November" }, + "Ideogram 4": { + "path": "CalamitousFelicitousness/Ideogram-4-bf16-Diffusers", + "desc": "Ideogram 4 is Ideogram's first open-weight text-to-image model: a 9.3B flow-matching single-stream DiT that uses a Qwen3-VL vision-language model as its text encoder, with strong in-image text rendering. Requires structured JSON-caption prompts; a plain-text prompt returns a built-in safety placeholder. Non-commercial license.", + "skip": true, + "extras": "sampler: Default, cfg_scale: 7.0, steps: 20, width: 1024, height: 1024", + "size": 50.0, + "date": "2026 June" + }, "Baidu ERNIE-Image": { "path": "baidu/ERNIE-Image", "preview": "baidu--ERNIE-Image.jpg", diff --git a/modules/modeldata.py b/modules/modeldata.py index 8ca37098f..75162eecd 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -55,6 +55,8 @@ def get_model_type(pipe): model_type = 'f1' elif "ZImage" in name or "Z-Image" in name: model_type = 'zimage' + elif "Ideogram4" in name: + model_type = 'ideogram4' elif "LuminaDiMOO" in name: model_type = 'luminadimoo' elif "Lumina2" in name: diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 84d793876..f87d7ae25 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -151,6 +151,8 @@ def guess_by_name(fn, current_guess): new_guess = 'NucleusImage' elif 'z-image' in fn.lower() or 'z_image' in fn.lower(): new_guess = 'ZImage' + elif 'ideogram' in fn.lower(): + new_guess = 'Ideogram4' elif 'longcat-image' in fn.lower(): new_guess = 'LongCat' elif 'ovis-image' in fn.lower(): diff --git a/modules/sd_models.py b/modules/sd_models.py index 935e07361..5c5444104 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -543,6 +543,10 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_z_image import load_z_image sd_model = load_z_image(checkpoint_info, diffusers_load_config) allow_post_quant = False + elif model_type in ['Ideogram4']: + from pipelines.model_ideogram4 import load_ideogram4 + sd_model = load_ideogram4(checkpoint_info, diffusers_load_config) + allow_post_quant = False elif model_type in ['LongCat']: from pipelines.model_longcat import load_longcat sd_model = load_longcat(checkpoint_info, diffusers_load_config) diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 151e5a174..7bbae7eda 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -11,7 +11,7 @@ from modules.image import convert SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options']) approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE": 3 } -flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'zimage', 'lumina2', 'cogview4', 'h1', 'cosmos', 'anima', 'chroma', 'omnigen', 'omnigen2', 'longcat'] +flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'zimage', 'lumina2', 'cogview4', 'h1', 'cosmos', 'anima', 'chroma', 'omnigen', 'omnigen2', 'longcat', 'ideogram4'] warned = False queue_lock = threading.Lock() diff --git a/modules/shared_items.py b/modules/shared_items.py index 4fd642458..3e98da1f2 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -40,6 +40,7 @@ pipelines = { 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None), 'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None), 'HunyuanImage': getattr(diffusers, 'HunyuanImagePipeline', None), + 'Ideogram4': getattr(diffusers, 'Ideogram4Pipeline', None), 'JoyEdit': getattr(diffusers, 'JoyImageEditPipeline', None), 'Kandinsky21': getattr(diffusers, 'KandinskyCombinedPipeline', None), 'Kandinsky22': getattr(diffusers, 'KandinskyV22CombinedPipeline', None), diff --git a/pipelines/generic_shared.py b/pipelines/generic_shared.py index 588c51602..069427049 100644 --- a/pipelines/generic_shared.py +++ b/pipelines/generic_shared.py @@ -1,5 +1,6 @@ import os import transformers +from transformers.models.qwen3_vl import Qwen3VLModel shared_te_map = { @@ -95,4 +96,9 @@ shared_te_map = { 'target_repo': 'vladmandic/Anima-1.0-Base', 'target_subfolder': 'text_encoder', }, + + 'Qwen3-VL 8B Base': { + 'cls': Qwen3VLModel, + 'target_repo': 'Qwen/Qwen3-VL-8B-Instruct', + }, } diff --git a/pipelines/ideogram4/__init__.py b/pipelines/ideogram4/__init__.py new file mode 100644 index 000000000..35573bfa1 --- /dev/null +++ b/pipelines/ideogram4/__init__.py @@ -0,0 +1,21 @@ +"""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"] diff --git a/pipelines/ideogram4/constants.py b/pipelines/ideogram4/constants.py new file mode 100644 index 000000000..d0db514ab --- /dev/null +++ b/pipelines/ideogram4/constants.py @@ -0,0 +1,17 @@ +"""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) diff --git a/pipelines/ideogram4/latent_norm.py b/pipelines/ideogram4/latent_norm.py new file mode 100644 index 000000000..7d24f3588 --- /dev/null +++ b/pipelines/ideogram4/latent_norm.py @@ -0,0 +1,281 @@ +"""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 diff --git a/pipelines/ideogram4/pipeline_ideogram4.py b/pipelines/ideogram4/pipeline_ideogram4.py new file mode 100644 index 000000000..d5930ed1a --- /dev/null +++ b/pipelines/ideogram4/pipeline_ideogram4.py @@ -0,0 +1,255 @@ +"""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 + + llm_features = encode_text(self.text_encoder, inputs["token_ids"], inputs["text_position_ids"], inputs["indicator"]) + + # 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 + 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) + + 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 decode_latents(self, z: torch.Tensor, grid_h: int, grid_w: int, output_type: str = "pil"): + """Denormalize, unpatchify, and VAE-decode the image latents.""" + 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() + z = z.view(batch_size, ae_channels, grid_h * patch, grid_w * patch) + + 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] diff --git a/pipelines/ideogram4/scheduler_ideogram4.py b/pipelines/ideogram4/scheduler_ideogram4.py new file mode 100644 index 000000000..f735c81cc --- /dev/null +++ b/pipelines/ideogram4/scheduler_ideogram4.py @@ -0,0 +1,68 @@ +"""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 diff --git a/pipelines/ideogram4/text_encoder_ideogram4.py b/pipelines/ideogram4/text_encoder_ideogram4.py new file mode 100644 index 000000000..66421c07f --- /dev/null +++ b/pipelines/ideogram4/text_encoder_ideogram4.py @@ -0,0 +1,95 @@ +"""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) diff --git a/pipelines/ideogram4/transformer_ideogram4.py b/pipelines/ideogram4/transformer_ideogram4.py new file mode 100644 index 000000000..49203ceb7 --- /dev/null +++ b/pipelines/ideogram4/transformer_ideogram4.py @@ -0,0 +1,330 @@ +"""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) diff --git a/pipelines/model_ideogram4.py b/pipelines/model_ideogram4.py new file mode 100644 index 000000000..87351a597 --- /dev/null +++ b/pipelines/model_ideogram4.py @@ -0,0 +1,45 @@ +import diffusers +from transformers import AutoTokenizer +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 + + +def load_ideogram4(checkpoint_info, diffusers_load_config=None): + if diffusers_load_config is None: + diffusers_load_config = {} + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + log.debug(f'Load model: type=Ideogram4 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}') + + 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) + # 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() + + pipe = Ideogram4Pipeline( + transformer=transformer, + unconditional_transformer=unconditional_transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ) + # The pipeline decodes the packed latent itself, so keep SD.Next from re-decoding it. + pipe.task_args = {'output_type': 'pil'} + + del transformer, unconditional_transformer, text_encoder, vae + devices.torch_gc(force=True, reason='load') + return pipe From b5fdb1995e177c58596104ed46473dd5117824ce Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 4 Jun 2026 01:01:19 +0100 Subject: [PATCH 2/7] test(ideogram4): add parity and smoke tests Offline transformer parity gate against the reference plus an SDNQ-int4 end-to-end smoke. --- test/test-ideogram4-parity.py | 138 ++++++++++++++++++++++++++++++++++ test/test-ideogram4-smoke.py | 103 +++++++++++++++++++++++++ 2 files changed, 241 insertions(+) create mode 100644 test/test-ideogram4-parity.py create mode 100644 test/test-ideogram4-smoke.py diff --git a/test/test-ideogram4-parity.py b/test/test-ideogram4-parity.py new file mode 100644 index 000000000..601b046f3 --- /dev/null +++ b/test/test-ideogram4-parity.py @@ -0,0 +1,138 @@ +#!/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()) diff --git a/test/test-ideogram4-smoke.py b/test/test-ideogram4-smoke.py new file mode 100644 index 000000000..467a158c6 --- /dev/null +++ b/test/test-ideogram4-smoke.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python +"""Standalone end-to-end smoke for the Ideogram 4 port. + +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. + +Usage: + python test/test-ideogram4-smoke.py --model /path/to/Ideogram-4-bf16 --output out.png +""" + +import argparse +import os +import sys +import time + +parser = argparse.ArgumentParser() +parser.add_argument("--model", required=True, help="converted 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) +parser.add_argument("--width", type=int, default=1024) +parser.add_argument("--steps", type=int, default=20) +parser.add_argument("--seed", type=int, default=0) +parser.add_argument("--weights-dtype", default="uint4", help="SDNQ weights dtype (uint4, int8, ...)") +parser.add_argument("--hf-cache", default=None, help="HF cache_dir for the shared Qwen3-VL encoder (default: HF default cache)") +args = parser.parse_args() + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) +os.chdir(REPO_ROOT) +os.environ["SD_INSTALL_QUIET"] = "1" + +# Our own args are already parsed; clear argv (and leave it cleared) so sdnext's +# shared.py / devices, which re-parse argv on import, don't see this script's flags. +sys.argv = [sys.argv[0]] + +import modules.cmd_args +import installer + +modules.cmd_args.parse_args() +installer.add_args(modules.cmd_args.parser) +modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + +import torch +import diffusers +from transformers import AutoTokenizer +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" + + +def main() -> int: + device = devices.device + 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) + print("loading unconditional_transformer ...", flush=True) + uncond = 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() + + pipe = Ideogram4Pipeline( + transformer=transformer, + unconditional_transformer=uncond, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ) + + 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) + elapsed = time.time() - start + + image = out.images[0] + image.save(args.output) + if torch.cuda.is_available(): + print(f"peak VRAM: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB", flush=True) + print(f"PASS: generated {args.output} in {elapsed:.1f}s, size={image.size}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 4866d2c8f3b4ccb527fdfee40fa3e9fb85cc0972 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 4 Jun 2026 02:36:42 +0100 Subject: [PATCH 3/7] fix(ideogram4): offload-safe encode, step progress, live preview - move the text encoder on-device for the tapped forward (bypasses the offload hook) - wrap the denoise loop in the diffusers progress bar - live preview through the shared TAE FLUX.2 decoder (Flux.2 VAE) --- modules/vae/sd_vae_taesd.py | 2 +- pipelines/ideogram4/pipeline_ideogram4.py | 81 +++++++++++++---------- 2 files changed, 46 insertions(+), 37 deletions(-) diff --git a/modules/vae/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py index cbae08815..b661205e7 100644 --- a/modules/vae/sd_vae_taesd.py +++ b/modules/vae/sd_vae_taesd.py @@ -71,7 +71,7 @@ def get_model(model_cls, variant=None): elif model_cls in {'f1', 'h1', 'zimage', 'lumina2', 'chroma', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4', 'ultraflux'}: model_cls = 'f1' variant = 'TAE FLUX.1' - elif model_cls in {'f2', 'ernieimage', 'lens'}: + elif model_cls in {'f2', 'ernieimage', 'lens', 'ideogram4'}: model_cls = 'f2' variant = 'TAE FLUX.2' elif model_cls in {'sd3'}: diff --git a/pipelines/ideogram4/pipeline_ideogram4.py b/pipelines/ideogram4/pipeline_ideogram4.py index d5930ed1a..556778908 100644 --- a/pipelines/ideogram4/pipeline_ideogram4.py +++ b/pipelines/ideogram4/pipeline_ideogram4.py @@ -167,7 +167,11 @@ class Ideogram4Pipeline(DiffusionPipeline): 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. @@ -186,41 +190,45 @@ class Ideogram4Pipeline(DiffusionPipeline): text_z_padding = torch.zeros(batch_size, max_text_tokens, latent_dim, dtype=torch.float32, device=device) self._num_timesteps = num_steps - 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) + 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, + 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=neg_position_ids, - segment_ids=neg_segment_ids, - indicator=neg_indicator, + position_ids=inputs["position_ids"], + segment_ids=inputs["segment_ids"], + indicator=inputs["indicator"], ) - v = gw_i * pos_v + (1.0 - gw_i) * neg_v - z = z + v * (s_val - t_val) + pos_v = pos_out[:, max_text_tokens:] - 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) + 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 @@ -231,19 +239,20 @@ class Ideogram4Pipeline(DiffusionPipeline): return (images,) return ImagePipelineOutput(images=images) - 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.""" + 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() - z = z.view(batch_size, ae_channels, grid_h * patch, grid_w * patch) + 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] From bcebb6b81b149e7343880721ea89120168e8ee91 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 5 Jun 2026 00:13:16 +0100 Subject: [PATCH 4/7] 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. --- installer.py | 2 +- modules/processing_callbacks.py | 20 ++ pipelines/ideogram4/__init__.py | 21 -- pipelines/ideogram4/constants.py | 17 - pipelines/ideogram4/latent_norm.py | 281 --------------- pipelines/ideogram4/pipeline_ideogram4.py | 264 -------------- pipelines/ideogram4/scheduler_ideogram4.py | 68 ---- pipelines/ideogram4/text_encoder_ideogram4.py | 95 ----- pipelines/ideogram4/transformer_ideogram4.py | 330 ------------------ pipelines/model_ideogram4.py | 42 ++- test/test-ideogram4-parity.py | 138 -------- test/test-ideogram4-smoke.py | 41 +-- 12 files changed, 69 insertions(+), 1250 deletions(-) delete mode 100644 pipelines/ideogram4/__init__.py delete mode 100644 pipelines/ideogram4/constants.py delete mode 100644 pipelines/ideogram4/latent_norm.py delete mode 100644 pipelines/ideogram4/pipeline_ideogram4.py delete mode 100644 pipelines/ideogram4/scheduler_ideogram4.py delete mode 100644 pipelines/ideogram4/text_encoder_ideogram4.py delete mode 100644 pipelines/ideogram4/transformer_ideogram4.py delete mode 100644 test/test-ideogram4-parity.py diff --git a/installer.py b/installer.py index 51ea09b65..cfcdcfce4 100644 --- a/installer.py +++ b/installer.py @@ -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') diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index c14576675..0df281c47 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -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 diff --git a/pipelines/ideogram4/__init__.py b/pipelines/ideogram4/__init__.py deleted file mode 100644 index 35573bfa1..000000000 --- a/pipelines/ideogram4/__init__.py +++ /dev/null @@ -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"] diff --git a/pipelines/ideogram4/constants.py b/pipelines/ideogram4/constants.py deleted file mode 100644 index d0db514ab..000000000 --- a/pipelines/ideogram4/constants.py +++ /dev/null @@ -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) diff --git a/pipelines/ideogram4/latent_norm.py b/pipelines/ideogram4/latent_norm.py deleted file mode 100644 index 7d24f3588..000000000 --- a/pipelines/ideogram4/latent_norm.py +++ /dev/null @@ -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 diff --git a/pipelines/ideogram4/pipeline_ideogram4.py b/pipelines/ideogram4/pipeline_ideogram4.py deleted file mode 100644 index 556778908..000000000 --- a/pipelines/ideogram4/pipeline_ideogram4.py +++ /dev/null @@ -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] diff --git a/pipelines/ideogram4/scheduler_ideogram4.py b/pipelines/ideogram4/scheduler_ideogram4.py deleted file mode 100644 index f735c81cc..000000000 --- a/pipelines/ideogram4/scheduler_ideogram4.py +++ /dev/null @@ -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 diff --git a/pipelines/ideogram4/text_encoder_ideogram4.py b/pipelines/ideogram4/text_encoder_ideogram4.py deleted file mode 100644 index 66421c07f..000000000 --- a/pipelines/ideogram4/text_encoder_ideogram4.py +++ /dev/null @@ -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) diff --git a/pipelines/ideogram4/transformer_ideogram4.py b/pipelines/ideogram4/transformer_ideogram4.py deleted file mode 100644 index 49203ceb7..000000000 --- a/pipelines/ideogram4/transformer_ideogram4.py +++ /dev/null @@ -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) diff --git a/pipelines/model_ideogram4.py b/pipelines/model_ideogram4.py index 87351a597..9e3f7f2a7 100644 --- a/pipelines/model_ideogram4.py +++ b/pipelines/model_ideogram4.py @@ -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') diff --git a/test/test-ideogram4-parity.py b/test/test-ideogram4-parity.py deleted file mode 100644 index 601b046f3..000000000 --- a/test/test-ideogram4-parity.py +++ /dev/null @@ -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()) diff --git a/test/test-ideogram4-smoke.py b/test/test-ideogram4-smoke.py index 467a158c6..971671dff 100644 --- a/test/test-ideogram4-smoke.py +++ b/test/test-ideogram4-smoke.py @@ -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] From 5c1a52ee4c27f317e067a8041db42bf53b9dc534 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 5 Jun 2026 01:37:03 +0100 Subject: [PATCH 5/7] perf(ideogram4): keep both transformers resident when they fit Each denoise step runs both transformers, so balanced offload ping-pongs them across PCIe every step. When both fit gpu_memory * max watermark, set offload_never so the per-step pre-sweep skips them and they stay resident; otherwise the normal offload path is kept, so smaller GPUs and bf16 fall back to offloading unchanged. --- pipelines/model_ideogram4.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/pipelines/model_ideogram4.py b/pipelines/model_ideogram4.py index 9e3f7f2a7..d6ebe0ccc 100644 --- a/pipelines/model_ideogram4.py +++ b/pipelines/model_ideogram4.py @@ -23,6 +23,28 @@ class Ideogram4Pipeline(diffusers.Ideogram4Pipeline): self.text_encoder.to(devices.cpu) +def pin_transformers_if_fit(transformer, unconditional_transformer) -> bool: + """Keep both transformers resident under balanced offload when they fit the budget. + + Every denoise step runs both transformers, so balanced offload ping-pongs them across + PCIe each step. ``offload_never`` makes ``offload_allowed`` skip the per-step pre-sweep so + they stay resident, but only when both fit ``gpu_memory * max watermark`` (which leaves the + watermark headroom for activations); otherwise the normal offload path is kept. + """ + if shared.opts.diffusers_offload_mode != 'balanced' or shared.gpu_memory <= 0: + return False + if transformer is None or unconditional_transformer is None: + return False + size_gb = sum(p.numel() * p.element_size() for m in (transformer, unconditional_transformer) for p in m.parameters()) / (1024 ** 3) + budget_gb = shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory + fits = size_gb <= budget_gb + if fits: + transformer.offload_never = True + unconditional_transformer.offload_never = True + log.info(f'Load model: type=Ideogram4 offload=balanced transformers={size_gb:.1f} budget={budget_gb:.1f} action={"pin-resident" if fits else "offload"}') + return fits + + def load_ideogram4(checkpoint_info, diffusers_load_config=None): if diffusers_load_config is None: diffusers_load_config = {} @@ -37,6 +59,7 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None): 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) + pin_transformers_if_fit(transformer, unconditional_transformer) # 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) From 92e1cb9927a45932d8ef31dcf728bbc38bf89f57 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 5 Jun 2026 01:37:55 +0100 Subject: [PATCH 6/7] fix(ideogram4): pass JSON captions through unmangled The dynamic-prompt brace processor in apply_styles_to_prompts strips the {} and [] out of a JSON caption, leaving non-JSON that trips the model's weight-baked safety placeholder. Let a model opt out of style and wildcard processing via keep_prompts and set it for Ideogram4, then normalize the prompt in encode_prompt: valid JSON to the compact training form, plain text wrapped into a minimal caption so basic prompts still generate. --- modules/processing.py | 2 +- pipelines/model_ideogram4.py | 32 +++++++++++++++++++++++++++----- 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/modules/processing.py b/modules/processing.py index 56d6da3f1..d73ea8973 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -278,7 +278,7 @@ def process_init(p: StableDiffusionProcessing): else: p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))] if reset_prompts: - if not hasattr(p, 'keep_prompts'): + if not hasattr(p, 'keep_prompts') and not getattr(shared.sd_model, 'keep_prompts', False): p.all_prompts, p.all_negative_prompts = shared.prompt_styles.apply_styles_to_prompts(p.all_prompts, p.all_negative_prompts, p.styles, p.all_seeds, p=p) p.prompts = p.all_prompts[(p.iteration * p.batch_size):((p.iteration+1) * p.batch_size)] p.negative_prompts = p.all_negative_prompts[(p.iteration * p.batch_size):((p.iteration+1) * p.batch_size)] diff --git a/pipelines/model_ideogram4.py b/pipelines/model_ideogram4.py index d6ebe0ccc..b7cae290f 100644 --- a/pipelines/model_ideogram4.py +++ b/pipelines/model_ideogram4.py @@ -1,3 +1,4 @@ +import json import diffusers from transformers import AutoTokenizer from transformers.models.qwen3_vl import Qwen3VLModel @@ -6,18 +7,37 @@ from modules.logger import log from pipelines import generic +def prompt_to_json(prompt): + """Normalize a JSON caption to the compact form Ideogram 4 trained on, or wrap plain text. + + Ideogram 4 expects a structured JSON caption serialized compactly. A valid JSON prompt is + re-serialized to that form; a plain-text prompt is wrapped in a minimal caption so it stays + in distribution instead of tripping the weight-baked "blocked by safety filter" placeholder. + """ + if isinstance(prompt, list): + return [prompt_to_json(p) for p in prompt] + if not isinstance(prompt, str) or len(prompt) == 0: + return prompt + try: + return json.dumps(json.loads(prompt), ensure_ascii=False, separators=(',', ':')) + except ValueError: + caption = {'high_level_description': prompt, 'compositional_deconstruction': {'background': prompt, 'elements': []}} + return json.dumps(caption, ensure_ascii=False, separators=(',', ':')) + + 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. + ``encode_prompt`` normalizes the prompt into the structured JSON the model expects, then + drives the Qwen3-VL tap. The tap calls ``language_model`` submodules directly, bypassing the + balanced-offload pre-forward hook, so the encoder is moved on-device for it and released after. """ - def encode_prompt(self, *args, **kwargs): + def encode_prompt(self, prompt, *args, **kwargs): + prompt = prompt_to_json(prompt) self.text_encoder.to(self._execution_device) try: - return super().encode_prompt(*args, **kwargs) + return super().encode_prompt(prompt, *args, **kwargs) finally: if shared.opts.diffusers_offload_mode != 'none': self.text_encoder.to(devices.cpu) @@ -78,6 +98,8 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None): # 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 + # JSON captions must pass through verbatim; skip styles/wildcards that would strip the braces. + pipe.keep_prompts = True # pylint: disable=attribute-defined-outside-init del transformer, unconditional_transformer, text_encoder, vae devices.torch_gc(force=True, reason='load') From 63fd441a14ae2c54e94ce915f1883919cb9d82fd Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 5 Jun 2026 02:07:39 +0100 Subject: [PATCH 7/7] refactor(ideogram4): register pipeline and load via from_pretrained Register the subclass with generic.set_pipeline and build it with from_pretrained, passing the SDNQ transformers and shared text encoder while the vae, scheduler, and tokenizer load from the repo. --- modules/shared_items.py | 2 +- pipelines/model_ideogram4.py | 25 ++++++++++++------------- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/modules/shared_items.py b/modules/shared_items.py index 3e98da1f2..613908ea7 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -40,7 +40,6 @@ pipelines = { 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None), 'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None), 'HunyuanImage': getattr(diffusers, 'HunyuanImagePipeline', None), - 'Ideogram4': getattr(diffusers, 'Ideogram4Pipeline', None), 'JoyEdit': getattr(diffusers, 'JoyImageEditPipeline', None), 'Kandinsky21': getattr(diffusers, 'KandinskyCombinedPipeline', None), 'Kandinsky22': getattr(diffusers, 'KandinskyV22CombinedPipeline', None), @@ -70,6 +69,7 @@ pipelines = { 'FLEX': None, 'HiDreamO1': None, 'HunyuanImage3': None, + 'Ideogram4': None, 'Lens': None, 'LuminaDiMOO': None, 'Meissonic': None, diff --git a/pipelines/model_ideogram4.py b/pipelines/model_ideogram4.py index b7cae290f..d77cdb6a0 100644 --- a/pipelines/model_ideogram4.py +++ b/pipelines/model_ideogram4.py @@ -1,8 +1,7 @@ import json import diffusers -from transformers import AutoTokenizer from transformers.models.qwen3_vl import Qwen3VLModel -from modules import shared, devices, sd_models +from modules import shared, devices, sd_models, model_quant from modules.logger import log from pipelines import generic @@ -70,8 +69,10 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None): diffusers_load_config = {} repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - log.debug(f'Load model: type=Ideogram4 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}') + load_args, _ = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + log.debug(f'Load model: type=Ideogram4 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + generic.set_pipeline('Ideogram4', Ideogram4Pipeline) if repo_id is None or repo_id.lower() == 'none': return None @@ -81,19 +82,17 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None): unconditional_transformer = generic.load_transformer(repo_id, cls_name=cls, subfolder="unconditional_transformer", load_config=diffusers_load_config) pin_transformers_if_fit(transformer, unconditional_transformer) # 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. + # the bundled text_encoder is the fallback when sharing is off. The vae, tokenizer, and + # scheduler load from the repo via from_pretrained. 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 = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler", cache_dir=shared.opts.diffusers_dir) - pipe = Ideogram4Pipeline( - scheduler=scheduler, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, + pipe = Ideogram4Pipeline.from_pretrained( + repo_id, + cache_dir=shared.opts.diffusers_dir, transformer=transformer, unconditional_transformer=unconditional_transformer, + text_encoder=text_encoder, + **load_args, ) # The pipeline decodes internally; the CFG scale slider drives guidance_scale, which is # mutually exclusive with the pipeline's default per-step guidance_schedule. @@ -101,6 +100,6 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None): # JSON captions must pass through verbatim; skip styles/wildcards that would strip the braces. pipe.keep_prompts = True # pylint: disable=attribute-defined-outside-init - del transformer, unconditional_transformer, text_encoder, vae + del transformer, unconditional_transformer, text_encoder devices.torch_gc(force=True, reason='load') return pipe