mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
refactor(ideogram4): adopt diffusers-native pipeline
diffusers shipped first-party Ideogram 4 (transformer + pipeline) in 9b0818cf, so drop the in-tree port and keep only SD.Next integration glue. Bump the diffusers pin to 9b0818cf and build diffusers' Ideogram4Pipeline from a thin loader with per-transformer SDNQ. A small subclass keeps the text encoder resident for the Qwen3-VL tap under balanced offload, and the step callback denormalizes the preview latent from vae.bn before unpatchify. Deletes the ported transformer, pipeline, scheduler, text encoder, and latent-norm constants.
This commit is contained in:
+1
-1
@@ -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')
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -1,264 +0,0 @@
|
||||
"""Ideogram 4 text-to-image pipeline (diffusers-native, SD.Next integration).
|
||||
|
||||
Owns its sampling loop: a packed text+image sequence is denoised with asymmetric
|
||||
dual-branch classifier-free guidance over two transformers (conditional and a
|
||||
separately-trained unconditional tower), then the image latents are denormalized,
|
||||
unpatchified, and decoded by the VAE.
|
||||
|
||||
The ``__call__`` signature names every argument SD.Next forwards (it prunes kwargs
|
||||
against the signature) and drives the diffusers ``callback_on_step_end`` contract
|
||||
for progress, interrupt, and preview.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
||||
|
||||
from pipelines.ideogram4.constants import (
|
||||
IMAGE_POSITION_OFFSET,
|
||||
LLM_TOKEN_INDICATOR,
|
||||
OUTPUT_IMAGE_INDICATOR,
|
||||
SEQUENCE_PADDING_INDICATOR,
|
||||
)
|
||||
from pipelines.ideogram4.latent_norm import get_latent_norm
|
||||
from pipelines.ideogram4.scheduler_ideogram4 import (
|
||||
get_schedule_for_resolution,
|
||||
make_step_intervals,
|
||||
)
|
||||
from pipelines.ideogram4.text_encoder_ideogram4 import encode_text, tokenize
|
||||
|
||||
|
||||
class Ideogram4Pipeline(DiffusionPipeline):
|
||||
"""Ideogram 4 flow-matching text-to-image pipeline."""
|
||||
|
||||
def __init__(self, transformer, unconditional_transformer, text_encoder, tokenizer, vae, scheduler) -> None:
|
||||
super().__init__()
|
||||
self.register_modules(
|
||||
transformer=transformer,
|
||||
unconditional_transformer=unconditional_transformer,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
vae=vae,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
self.patch_size = 2
|
||||
self.ae_scale_factor = 8
|
||||
self.max_text_tokens = 2048
|
||||
self.latent_shift = None
|
||||
self.latent_scale = None
|
||||
self._num_timesteps = 0
|
||||
|
||||
@property
|
||||
def num_timesteps(self) -> int:
|
||||
return self._num_timesteps
|
||||
|
||||
def latent_norm(self, device, dtype) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if self.latent_shift is None:
|
||||
self.latent_shift, self.latent_scale = get_latent_norm()
|
||||
return self.latent_shift.to(device=device, dtype=dtype), self.latent_scale.to(device=device, dtype=dtype)
|
||||
|
||||
def build_inputs(self, prompts: list[str], height: int, width: int, device) -> dict:
|
||||
"""Build the packed (text tokens + image latent tokens) sequence for one batch."""
|
||||
tokenized = [tokenize(self.tokenizer, p, self.max_text_tokens) for p in prompts]
|
||||
batch_size = len(prompts)
|
||||
|
||||
patch = self.patch_size * self.ae_scale_factor
|
||||
if height % patch != 0 or width % patch != 0:
|
||||
raise ValueError(f"height/width must be divisible by patch_size*ae_scale_factor={patch}")
|
||||
grid_h = height // patch
|
||||
grid_w = width // patch
|
||||
num_image_tokens = grid_h * grid_w
|
||||
|
||||
max_text_tokens = max(num_text for _, num_text in tokenized)
|
||||
total_seq_len = max_text_tokens + num_image_tokens
|
||||
|
||||
# Image position ids (t=0, h, w), offset to stay disjoint from text positions.
|
||||
h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1)
|
||||
w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1)
|
||||
t_idx = torch.zeros_like(h_idx)
|
||||
image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET
|
||||
|
||||
token_ids = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
|
||||
text_position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
|
||||
position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
|
||||
segment_ids = torch.full((batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long)
|
||||
indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
|
||||
|
||||
for b, (toks, num_text) in enumerate(tokenized):
|
||||
pad_len = max_text_tokens - num_text
|
||||
total_unpadded = num_text + num_image_tokens
|
||||
offset = pad_len # layout: [pad] [text] [image]
|
||||
|
||||
token_ids[b, offset : offset + num_text] = toks
|
||||
|
||||
text_pos = torch.arange(num_text)
|
||||
text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=1)
|
||||
text_position_ids[b, offset : offset + num_text] = text_pos_3d
|
||||
position_ids[b, offset : offset + num_text] = text_pos_3d
|
||||
position_ids[b, offset + num_text :] = image_pos
|
||||
|
||||
indicator[b, offset : offset + num_text] = LLM_TOKEN_INDICATOR
|
||||
indicator[b, offset + num_text :] = OUTPUT_IMAGE_INDICATOR
|
||||
segment_ids[b, offset : offset + total_unpadded] = 1
|
||||
|
||||
return {
|
||||
"token_ids": token_ids.to(device),
|
||||
"text_position_ids": text_position_ids.to(device),
|
||||
"position_ids": position_ids.to(device),
|
||||
"segment_ids": segment_ids.to(device),
|
||||
"indicator": indicator.to(device),
|
||||
"num_image_tokens": num_image_tokens,
|
||||
"grid_h": grid_h,
|
||||
"grid_w": grid_w,
|
||||
"max_text_tokens": max_text_tokens,
|
||||
}
|
||||
|
||||
def init_noise(self, batch_size: int, num_image_tokens: int, latent_dim: int, generator, device) -> torch.Tensor:
|
||||
if isinstance(generator, list):
|
||||
samples = [
|
||||
torch.randn((1, num_image_tokens, latent_dim), generator=generator[b % len(generator)], device=device, dtype=torch.float32)
|
||||
for b in range(batch_size)
|
||||
]
|
||||
return torch.cat(samples, dim=0)
|
||||
return torch.randn((batch_size, num_image_tokens, latent_dim), generator=generator, device=device, dtype=torch.float32)
|
||||
|
||||
def resolve_sampling(self, num_inference_steps: int, guidance_scale: float, device):
|
||||
"""Map SD.Next's steps + CFG scale to a flat per-step guidance schedule."""
|
||||
num_steps = int(num_inference_steps)
|
||||
guidance_schedule = torch.full((num_steps,), float(guidance_scale), dtype=torch.float32, device=device)
|
||||
return num_steps, guidance_schedule
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt: str | list[str] | None = None,
|
||||
negative_prompt: str | list[str] | None = None,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float = 7.0,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
generator=None,
|
||||
output_type: str = "pil",
|
||||
return_dict: bool = True,
|
||||
callback_on_step_end=None,
|
||||
callback_on_step_end_tensor_inputs=None,
|
||||
**kwargs,
|
||||
):
|
||||
device = self._execution_device
|
||||
|
||||
if prompt is None:
|
||||
prompts = [""]
|
||||
elif isinstance(prompt, str):
|
||||
prompts = [prompt]
|
||||
else:
|
||||
prompts = list(prompt)
|
||||
batch_size = len(prompts)
|
||||
|
||||
num_steps, gw_per_step = self.resolve_sampling(num_inference_steps, guidance_scale, device)
|
||||
schedule = get_schedule_for_resolution((height, width), known_mean=self.scheduler.config.mu, std=self.scheduler.config.std)
|
||||
step_intervals = make_step_intervals(num_steps).to(device)
|
||||
|
||||
inputs = self.build_inputs(prompts, height=height, width=width, device=device)
|
||||
num_image_tokens = inputs["num_image_tokens"]
|
||||
grid_h, grid_w = inputs["grid_h"], inputs["grid_w"]
|
||||
max_text_tokens = inputs["max_text_tokens"]
|
||||
latent_dim = self.transformer.config.in_channels
|
||||
|
||||
# The tapped forward bypasses the offload hook, so move the encoder on-device, then free it after.
|
||||
from modules import devices, shared
|
||||
self.text_encoder.to(device)
|
||||
llm_features = encode_text(self.text_encoder, inputs["token_ids"], inputs["text_position_ids"], inputs["indicator"])
|
||||
self.text_encoder.to(devices.cpu)
|
||||
|
||||
# At guidance 1.0 the unconditional velocity has zero weight, so skip the second
|
||||
# tower; only build the negative branch when some step needs it.
|
||||
gw_values = gw_per_step.tolist()
|
||||
use_cfg = any(gw != 1.0 for gw in gw_values)
|
||||
|
||||
neg_position_ids = neg_segment_ids = neg_indicator = neg_llm_features = None
|
||||
if use_cfg:
|
||||
# Negative branch is image-only (asymmetric CFG) with zeroed conditioning.
|
||||
neg_position_ids = inputs["position_ids"][:, max_text_tokens:]
|
||||
neg_segment_ids = inputs["segment_ids"][:, max_text_tokens:]
|
||||
neg_indicator = inputs["indicator"][:, max_text_tokens:]
|
||||
neg_llm_features = torch.zeros(batch_size, num_image_tokens, llm_features.shape[-1], dtype=llm_features.dtype, device=device)
|
||||
|
||||
z = self.init_noise(batch_size, num_image_tokens, latent_dim, generator, device)
|
||||
text_z_padding = torch.zeros(batch_size, max_text_tokens, latent_dim, dtype=torch.float32, device=device)
|
||||
|
||||
self._num_timesteps = num_steps
|
||||
with self.progress_bar(total=num_steps) as progress_bar:
|
||||
for i in range(num_steps - 1, -1, -1):
|
||||
t_val = float(schedule(step_intervals[i + 1].unsqueeze(0)).item())
|
||||
s_val = float(schedule(step_intervals[i].unsqueeze(0)).item())
|
||||
t = torch.full((batch_size,), t_val, dtype=torch.float32, device=device)
|
||||
|
||||
pos_z = torch.cat([text_z_padding, z], dim=1)
|
||||
pos_out = self.transformer(
|
||||
llm_features=llm_features,
|
||||
x=pos_z,
|
||||
t=t,
|
||||
position_ids=inputs["position_ids"],
|
||||
segment_ids=inputs["segment_ids"],
|
||||
indicator=inputs["indicator"],
|
||||
)
|
||||
pos_v = pos_out[:, max_text_tokens:]
|
||||
|
||||
gw_i = gw_values[i]
|
||||
if gw_i == 1.0:
|
||||
v = pos_v # unconditional weight is zero; skip the second tower
|
||||
else:
|
||||
neg_v = self.unconditional_transformer(
|
||||
llm_features=neg_llm_features,
|
||||
x=z,
|
||||
t=t,
|
||||
position_ids=neg_position_ids,
|
||||
segment_ids=neg_segment_ids,
|
||||
indicator=neg_indicator,
|
||||
)
|
||||
v = gw_i * pos_v + (1.0 - gw_i) * neg_v
|
||||
z = z + v * (s_val - t_val)
|
||||
|
||||
if callback_on_step_end is not None:
|
||||
cb = callback_on_step_end(self, num_steps - 1 - i, t_val, {"latents": z})
|
||||
if isinstance(cb, dict):
|
||||
z = cb.get("latents", z)
|
||||
# Live preview: denorm+unpack into Flux.2 latent space for TAE FLUX.2.
|
||||
shared.state.current_latent = self.denorm_unpack(z, grid_h, grid_w)
|
||||
progress_bar.update()
|
||||
|
||||
if output_type == "latent":
|
||||
images = z
|
||||
else:
|
||||
images = self.decode_latents(z, grid_h=grid_h, grid_w=grid_w, output_type=output_type)
|
||||
|
||||
if not return_dict:
|
||||
return (images,)
|
||||
return ImagePipelineOutput(images=images)
|
||||
|
||||
def denorm_unpack(self, z: torch.Tensor, grid_h: int, grid_w: int) -> torch.Tensor:
|
||||
"""Denormalize the packed latent and unpatchify to (B, ae_channels, H, W) in VAE space."""
|
||||
batch_size = z.shape[0]
|
||||
patch = self.patch_size
|
||||
shift, scale = self.latent_norm(z.device, torch.float32)
|
||||
z = z.float() * scale + shift
|
||||
ae_channels = z.shape[-1] // (patch * patch)
|
||||
z = z.view(batch_size, grid_h, grid_w, patch, patch, ae_channels)
|
||||
z = z.permute(0, 5, 1, 3, 2, 4).contiguous()
|
||||
return z.view(batch_size, ae_channels, grid_h * patch, grid_w * patch)
|
||||
|
||||
def decode_latents(self, z: torch.Tensor, grid_h: int, grid_w: int, output_type: str = "pil"):
|
||||
"""Denormalize, unpatchify, and VAE-decode the image latents."""
|
||||
z = self.denorm_unpack(z, grid_h, grid_w)
|
||||
vae_dtype = next((p.dtype for p in self.vae.parameters() if torch.is_floating_point(p)), torch.float32)
|
||||
decoded = self.vae.decode(z.to(vae_dtype), return_dict=False)[0]
|
||||
|
||||
decoded = decoded.float().clamp(-1.0, 1.0)
|
||||
decoded = ((decoded + 1.0) * 127.5).round().to(torch.uint8)
|
||||
decoded = decoded.permute(0, 2, 3, 1).cpu().numpy()
|
||||
if output_type == "np":
|
||||
return decoded
|
||||
return [Image.fromarray(arr) for arr in decoded]
|
||||
@@ -1,68 +0,0 @@
|
||||
"""Logit-normal flow-matching schedule for Ideogram 4.
|
||||
|
||||
Sampling integrates the flow-matching ODE with Euler steps over a logit-normal
|
||||
timestep schedule whose mean shifts with resolution (more steps at high noise for
|
||||
larger images). Guidance is a per-step weight applied as an asymmetric blend of
|
||||
the conditional and unconditional velocity. Ported from the reference
|
||||
(github.com/ideogram-oss/ideogram4).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LogitNormalSchedule:
|
||||
mean: float
|
||||
std: float = 1.0
|
||||
logsnr_min: float = -15.0
|
||||
logsnr_max: float = 18.0
|
||||
|
||||
def __call__(self, t: torch.Tensor) -> torch.Tensor:
|
||||
t = t.to(torch.float64)
|
||||
z = torch.special.ndtri(t)
|
||||
y = self.mean + self.std * z
|
||||
t_ = torch.special.expit(y)
|
||||
t_ = 1 - t_
|
||||
t_min = 1.0 / (1 + math.exp(0.5 * self.logsnr_max))
|
||||
t_max = 1.0 / (1 + math.exp(0.5 * self.logsnr_min))
|
||||
return t_.clamp(t_min, t_max).to(torch.float32)
|
||||
|
||||
|
||||
def get_schedule_for_resolution(
|
||||
image_resolution: tuple[int, int],
|
||||
known_resolution: tuple[int, int] = (512, 512),
|
||||
known_mean: float = 1.0,
|
||||
std: float = 1.0,
|
||||
) -> LogitNormalSchedule:
|
||||
"""Resolution-aware schedule: the mean shifts by half the log pixel-count ratio."""
|
||||
num_pixels = image_resolution[0] * image_resolution[1]
|
||||
known_pixels = known_resolution[0] * known_resolution[1]
|
||||
mean = known_mean + 0.5 * math.log(num_pixels / known_pixels)
|
||||
return LogitNormalSchedule(mean=mean, std=std)
|
||||
|
||||
|
||||
def make_step_intervals(num_steps: int) -> torch.Tensor:
|
||||
"""Linear step schedule mapped through the logit-normal schedule at sample time."""
|
||||
return torch.linspace(0.0, 1.0, num_steps + 1, dtype=torch.float32)
|
||||
|
||||
|
||||
class Ideogram4Scheduler(SchedulerMixin, ConfigMixin):
|
||||
"""Thin diffusers scheduler holding the logit-normal defaults.
|
||||
|
||||
The denoise loop lives in the pipeline (asymmetric dual-branch CFG over a
|
||||
flat per-step guidance), so this only carries the ``mu``/``std`` that
|
||||
parameterize the resolution-aware schedule and serves as the registered
|
||||
``scheduler`` component.
|
||||
"""
|
||||
|
||||
@register_to_config
|
||||
def __init__(self, mu: float = 0.0, std: float = 1.75) -> None:
|
||||
pass
|
||||
@@ -1,95 +0,0 @@
|
||||
"""Qwen3-VL text conditioning for Ideogram 4.
|
||||
|
||||
The prompt (plain text or a structured JSON caption) is wrapped in the Qwen3
|
||||
chat template and run through the Qwen3-VL language model. Hidden states are
|
||||
captured from 13 intermediate layers (pre final-norm) and concatenated along the
|
||||
feature dim, giving the DiT multi-scale semantic features.
|
||||
|
||||
v1 feeds the prompt verbatim (no magic-prompt expansion, no caption verifier).
|
||||
These weights require a structured JSON caption: a plain-text prompt lands
|
||||
out-of-distribution and the model renders a baked-in "Image blocked by safety
|
||||
filter" placeholder. The caption schema is a ``compositional_deconstruction``
|
||||
object (with ``background`` + ``elements``); see the upstream prompting guide.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
|
||||
from pipelines.ideogram4.constants import LLM_TOKEN_INDICATOR, QWEN3_VL_ACTIVATION_LAYERS
|
||||
|
||||
|
||||
def tokenize(tokenizer, prompt: str, max_text_tokens: int) -> tuple[torch.Tensor, int]:
|
||||
"""Chat-template tokenize a single prompt (passed verbatim)."""
|
||||
messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
|
||||
text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
|
||||
encoded = tokenizer(text, return_tensors="pt", add_special_tokens=False)
|
||||
token_ids = encoded["input_ids"][0]
|
||||
num_text_tokens = int(token_ids.shape[0])
|
||||
if num_text_tokens > max_text_tokens:
|
||||
raise ValueError(f"prompt has {num_text_tokens} tokens, exceeds max_text_tokens={max_text_tokens}")
|
||||
return token_ids, num_text_tokens
|
||||
|
||||
|
||||
def qwen3_vl_layer_features(text_encoder, token_ids: torch.Tensor, attention_mask: torch.Tensor, pos_2d: torch.Tensor) -> list[torch.Tensor]:
|
||||
"""Run the Qwen3-VL language model and return hidden states at the tap layers.
|
||||
|
||||
The layer loop is driven manually so the tapped states are captured before the
|
||||
model's final norm, matching the reference.
|
||||
"""
|
||||
language_model = text_encoder.language_model
|
||||
|
||||
inputs_embeds = language_model.embed_tokens(token_ids)
|
||||
|
||||
position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1)
|
||||
text_position_ids = position_ids_4d[0]
|
||||
mrope_position_ids = position_ids_4d[1:]
|
||||
|
||||
causal_mask = create_causal_mask(
|
||||
config=language_model.config,
|
||||
inputs_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
past_key_values=None,
|
||||
position_ids=text_position_ids,
|
||||
)
|
||||
position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)
|
||||
|
||||
tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
|
||||
captured: dict[int, torch.Tensor] = {}
|
||||
hidden_states = inputs_embeds
|
||||
for layer_idx, decoder_layer in enumerate(language_model.layers):
|
||||
hidden_states = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=causal_mask,
|
||||
position_ids=text_position_ids,
|
||||
past_key_values=None,
|
||||
position_embeddings=position_embeddings,
|
||||
)
|
||||
if layer_idx in tap_set:
|
||||
captured[layer_idx] = hidden_states
|
||||
|
||||
return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
|
||||
|
||||
|
||||
def encode_text(text_encoder, token_ids: torch.Tensor, text_position_ids: torch.Tensor, indicator: torch.Tensor) -> torch.Tensor:
|
||||
"""Stack the tap-layer hidden states into (B, L, hidden_size * num_taps) float32.
|
||||
|
||||
Non-LLM positions (left padding / image slots) are zeroed so the DiT only sees
|
||||
real text features at LLM_TOKEN_INDICATOR positions.
|
||||
"""
|
||||
batch_size, seq_len = token_ids.shape
|
||||
|
||||
attention_mask = (indicator == LLM_TOKEN_INDICATOR).to(torch.long)
|
||||
pos_2d = text_position_ids[..., 0].contiguous()
|
||||
|
||||
with torch.no_grad():
|
||||
selected = qwen3_vl_layer_features(text_encoder, token_ids, attention_mask, pos_2d)
|
||||
|
||||
stacked = torch.stack(selected, dim=0) # (num_taps, B, L, H)
|
||||
stacked = torch.permute(stacked, (1, 2, 3, 0)) # (B, L, H, num_taps)
|
||||
stacked = stacked.reshape(batch_size, seq_len, -1) # (B, L, H * num_taps)
|
||||
|
||||
text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
|
||||
stacked = stacked * text_mask
|
||||
return stacked.to(torch.float32)
|
||||
@@ -1,330 +0,0 @@
|
||||
"""Ideogram 4 flow-matching DiT transformer.
|
||||
|
||||
A single-stream Diffusion Transformer: Qwen3-VL text features and noisy image
|
||||
latent tokens are concatenated into one sequence and processed by 34 shared
|
||||
blocks (QK-RMSNorm attention, SwiGLU MLP, tanh-gated AdaLN), with 3D multimodal
|
||||
RoPE giving text and image tokens a unified positional space. The model predicts
|
||||
a flow-matching velocity on the image tokens.
|
||||
|
||||
Ported from the reference implementation (github.com/ideogram-oss/ideogram4) as a
|
||||
diffusers ModelMixin so it loads from the shipped diffusers-layout checkpoint and
|
||||
can be SDNQ-quantized at load. Parameter names match the reference state dict
|
||||
1:1, so no key remapping is needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
|
||||
from pipelines.ideogram4.constants import LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR
|
||||
|
||||
|
||||
def rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
half = x.shape[-1] // 2
|
||||
x1 = x[..., :half]
|
||||
x2 = x[..., half:]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# q, k: (B, num_heads, L, head_dim); cos/sin: (B, L, head_dim).
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
q_embed = (q * cos) + (rotate_half(q) * sin)
|
||||
k_embed = (k * cos) + (rotate_half(k) * sin)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
def sinusoidal_embedding(t: torch.Tensor, dim: int, scale: float = 1e4) -> torch.Tensor:
|
||||
t = t.to(torch.float32)
|
||||
half = dim // 2
|
||||
freq = math.log(scale) / (half - 1)
|
||||
freq = torch.exp(torch.arange(half, dtype=torch.float32, device=t.device) * -freq)
|
||||
emb = t.unsqueeze(-1) * freq
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
if dim % 2 == 1:
|
||||
emb = F.pad(emb, (0, 1))
|
||||
return emb
|
||||
|
||||
|
||||
class Ideogram4MRoPE(nn.Module):
|
||||
inv_freq: torch.Tensor
|
||||
|
||||
def __init__(self, head_dim: int, base: int, mrope_section: tuple[int, ...]) -> None:
|
||||
super().__init__()
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.mrope_section = tuple(mrope_section)
|
||||
self.head_dim = head_dim
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# position_ids: (B, L, 3) of int (t, h, w).
|
||||
assert position_ids.ndim == 3 and position_ids.shape[-1] == 3
|
||||
batch_size, _, _ = position_ids.shape
|
||||
|
||||
pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32) # (3, B, L)
|
||||
inv_freq = self.inv_freq.to(dtype=torch.float32)[None, None, :, None].expand(3, batch_size, -1, 1)
|
||||
freqs = inv_freq @ pos.unsqueeze(2) # (3, B, F, L)
|
||||
freqs = freqs.transpose(2, 3) # (3, B, L, F)
|
||||
|
||||
# interleaved mrope: pull H freqs into idx 1 mod 3, W freqs into idx 2 mod 3.
|
||||
freqs_t = freqs[0].clone()
|
||||
for axis, offset in ((1, 1), (2, 2)):
|
||||
length = self.mrope_section[axis] * 3
|
||||
idx = torch.arange(offset, length, 3, device=freqs_t.device)
|
||||
freqs_t[..., idx] = freqs[axis][..., idx]
|
||||
|
||||
emb = torch.cat((freqs_t, freqs_t), dim=-1)
|
||||
return emb.cos(), emb.sin()
|
||||
|
||||
|
||||
class Ideogram4RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return F.rms_norm(x, self.weight.shape, self.weight, self.eps)
|
||||
|
||||
|
||||
class Ideogram4Attention(nn.Module):
|
||||
def __init__(self, hidden_size: int, num_heads: int, eps: float = 1e-5) -> None:
|
||||
super().__init__()
|
||||
assert hidden_size % num_heads == 0
|
||||
self.hidden_size = hidden_size
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = hidden_size // num_heads
|
||||
|
||||
self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
|
||||
self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps)
|
||||
self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps)
|
||||
self.o = nn.Linear(hidden_size, hidden_size, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor, segment_ids: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, seq_len, _ = x.shape
|
||||
|
||||
qkv = self.qkv(x)
|
||||
qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
|
||||
q, k, v = qkv.unbind(dim=2)
|
||||
|
||||
q = self.norm_q(q)
|
||||
k = self.norm_k(k)
|
||||
|
||||
# SDPA expects (B, num_heads, L, head_dim).
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
q, k = apply_rotary_pos_emb(q, k, cos, sin)
|
||||
|
||||
# Block-diagonal mask from segment ids: (B, 1, L, L), True = attend.
|
||||
attn_mask = (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1)
|
||||
|
||||
out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
|
||||
out = out.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)
|
||||
return self.o(out)
|
||||
|
||||
|
||||
class Ideogram4MLP(nn.Module):
|
||||
def __init__(self, dim: int, hidden_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
||||
|
||||
|
||||
class Ideogram4TransformerBlock(nn.Module):
|
||||
def __init__(self, hidden_size: int, intermediate_size: int, num_heads: int, norm_eps: float, adaln_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.attention = Ideogram4Attention(hidden_size, num_heads, eps=1e-5)
|
||||
self.feed_forward = Ideogram4MLP(hidden_size, intermediate_size)
|
||||
|
||||
self.attention_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.attention_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.ffn_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
|
||||
self.adaln_modulation = nn.Linear(adaln_dim, 4 * hidden_size, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor, segment_ids: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, adaln_input: torch.Tensor) -> torch.Tensor:
|
||||
mod = self.adaln_modulation(adaln_input)
|
||||
scale_msa, gate_msa, scale_mlp, gate_mlp = mod.chunk(4, dim=-1)
|
||||
gate_msa = torch.tanh(gate_msa)
|
||||
gate_mlp = torch.tanh(gate_mlp)
|
||||
scale_msa = 1.0 + scale_msa
|
||||
scale_mlp = 1.0 + scale_mlp
|
||||
|
||||
attn_out = self.attention(self.attention_norm1(x) * scale_msa, segment_ids=segment_ids, cos=cos, sin=sin)
|
||||
x = x + gate_msa * self.attention_norm2(attn_out)
|
||||
x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
|
||||
return x
|
||||
|
||||
|
||||
class Ideogram4EmbedScalar(nn.Module):
|
||||
def __init__(self, dim: int, input_range: tuple[float, float]) -> None:
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.range_min, self.range_max = input_range
|
||||
assert self.range_max > self.range_min
|
||||
self.mlp_in = nn.Linear(dim, dim, bias=True)
|
||||
self.mlp_out = nn.Linear(dim, dim, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# x holds a scalar per token; keep its (float) dtype as the compute dtype so
|
||||
# SDNQ-quantized Linears (whose .weight is int) don't drive an int cast.
|
||||
compute_dtype = x.dtype if torch.is_floating_point(x) else torch.float32
|
||||
x = x.to(torch.float32)
|
||||
scaled = 1e4 * (x - self.range_min) / (self.range_max - self.range_min)
|
||||
emb = sinusoidal_embedding(scaled, self.dim)
|
||||
emb = emb.to(getattr(self.mlp_in, "compute_dtype", None) or compute_dtype)
|
||||
emb = F.silu(self.mlp_in(emb))
|
||||
return self.mlp_out(emb)
|
||||
|
||||
|
||||
class Ideogram4FinalLayer(nn.Module):
|
||||
def __init__(self, hidden_size: int, out_channels: int, adaln_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
|
||||
self.linear = nn.Linear(hidden_size, out_channels, bias=True)
|
||||
self.adaln_modulation = nn.Linear(adaln_dim, hidden_size, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
|
||||
scale = 1.0 + self.adaln_modulation(F.silu(c))
|
||||
return self.linear(self.norm_final(x) * scale)
|
||||
|
||||
|
||||
class Ideogram4Transformer2DModel(ModelMixin, ConfigMixin):
|
||||
"""Ideogram 4 flow-matching transformer (single-stream DiT)."""
|
||||
|
||||
_no_split_modules = ["Ideogram4TransformerBlock"]
|
||||
_supports_gradient_checkpointing = False
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_attention_heads: int = 18,
|
||||
attention_head_dim: int = 256,
|
||||
num_layers: int = 34,
|
||||
intermediate_size: int = 12288,
|
||||
adaln_dim: int = 512,
|
||||
in_channels: int = 128,
|
||||
llm_features_dim: int = 53248,
|
||||
mrope_section: tuple[int, ...] = (24, 20, 20),
|
||||
rope_theta: int = 5_000_000,
|
||||
norm_eps: float = 1e-5,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
emb_dim = num_attention_heads * attention_head_dim
|
||||
head_dim = attention_head_dim
|
||||
self.num_heads = num_attention_heads
|
||||
self.emb_dim = emb_dim
|
||||
|
||||
self.input_proj = nn.Linear(in_channels, emb_dim, bias=True)
|
||||
self.llm_cond_norm = Ideogram4RMSNorm(llm_features_dim, eps=1e-6)
|
||||
self.llm_cond_proj = nn.Linear(llm_features_dim, emb_dim, bias=True)
|
||||
self.t_embedding = Ideogram4EmbedScalar(emb_dim, input_range=(0.0, 1.0))
|
||||
self.adaln_proj = nn.Linear(emb_dim, adaln_dim, bias=True)
|
||||
|
||||
self.embed_image_indicator = nn.Embedding(2, emb_dim)
|
||||
|
||||
self.rotary_emb = Ideogram4MRoPE(head_dim=head_dim, base=rope_theta, mrope_section=tuple(mrope_section))
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
Ideogram4TransformerBlock(
|
||||
hidden_size=emb_dim,
|
||||
intermediate_size=intermediate_size,
|
||||
num_heads=num_attention_heads,
|
||||
norm_eps=norm_eps,
|
||||
adaln_dim=adaln_dim,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = Ideogram4FinalLayer(hidden_size=emb_dim, out_channels=in_channels, adaln_dim=adaln_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
*,
|
||||
llm_features: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
t: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
segment_ids: torch.Tensor,
|
||||
indicator: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Velocity prediction.
|
||||
|
||||
Args:
|
||||
llm_features: (B, L, llm_features_dim) Qwen3-VL conditioning features.
|
||||
x: (B, L, in_channels) noise tokens.
|
||||
t: (B,) or (B, L) flow-matching time in [0, 1].
|
||||
position_ids: (B, L, 3) (t, h, w) positions for MRoPE.
|
||||
segment_ids: (B, L) sample id within a packed batch.
|
||||
indicator: (B, L) per-token role (LLM_TOKEN_INDICATOR / OUTPUT_IMAGE_INDICATOR).
|
||||
|
||||
Returns:
|
||||
(B, L, in_channels) velocity in float32; only OUTPUT_IMAGE_INDICATOR
|
||||
positions are meaningful.
|
||||
"""
|
||||
_, _, in_channels = x.shape
|
||||
assert in_channels == self.config["in_channels"]
|
||||
|
||||
# Compute dtype from a norm weight (never SDNQ-quantized), honoring an
|
||||
# explicit Fp8Linear compute_dtype if one is present.
|
||||
param_dtype = getattr(self.input_proj, "compute_dtype", None)
|
||||
if param_dtype is None:
|
||||
w = self.input_proj.weight
|
||||
param_dtype = w.dtype if torch.is_floating_point(w) else self.llm_cond_norm.weight.dtype
|
||||
|
||||
x = x.to(param_dtype)
|
||||
t = t.to(param_dtype)
|
||||
llm_features = llm_features.to(param_dtype)
|
||||
|
||||
indicator = indicator.to(torch.long)
|
||||
llm_token_mask = (indicator == LLM_TOKEN_INDICATOR).to(x.dtype).unsqueeze(-1)
|
||||
output_image_mask = (indicator == OUTPUT_IMAGE_INDICATOR).to(x.dtype).unsqueeze(-1)
|
||||
|
||||
llm_features = llm_features * llm_token_mask
|
||||
x = x * output_image_mask
|
||||
|
||||
x = self.input_proj(x) * output_image_mask
|
||||
|
||||
# Keep shape (B, 1, ...) when t is per-sample so the adaln projections don't
|
||||
# pay for L identical copies.
|
||||
t_cond = self.t_embedding(t)
|
||||
if t.dim() == 1:
|
||||
t_cond = t_cond.unsqueeze(1)
|
||||
adaln_input = F.silu(self.adaln_proj(t_cond))
|
||||
|
||||
llm_features = self.llm_cond_norm(llm_features)
|
||||
llm_features = self.llm_cond_proj(llm_features) * llm_token_mask
|
||||
|
||||
h = x + llm_features
|
||||
|
||||
image_indicator_embedding = self.embed_image_indicator((indicator == OUTPUT_IMAGE_INDICATOR).to(torch.long))
|
||||
h = h + image_indicator_embedding
|
||||
|
||||
cos, sin = self.rotary_emb(position_ids)
|
||||
cos = cos.to(h.dtype)
|
||||
sin = sin.to(h.dtype)
|
||||
|
||||
for layer in self.layers:
|
||||
h = layer(h, segment_ids=segment_ids, cos=cos, sin=sin, adaln_input=adaln_input)
|
||||
|
||||
out = self.final_layer(h, c=adaln_input)
|
||||
return out.to(torch.float32)
|
||||
@@ -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')
|
||||
|
||||
@@ -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())
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user