mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
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.
This commit is contained in:
CalamitousFelicitousness
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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',
|
||||
},
|
||||
}
|
||||
|
||||
@@ -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"]
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user