From 2f3d0e719db1829eae37d15e6a5a9bbbacd0c4cd Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 4 Jun 2026 01:00:01 +0100 Subject: [PATCH] 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