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]