From 6a8d9333b62d90cc1329ed02d0126e241b2f0652 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 4 Jun 2026 13:00:23 -0600 Subject: [PATCH] Improved the prompt handeling of ideogram4 model. Now used advanced prompts class to store them smaller and allow longer prompts --- .../diffusion_models/ideogram4/ideogram4.py | 69 +++++++++---------- .../ideogram4/src/pipeline.py | 36 ++++++++-- toolkit/data_loader.py | 2 +- toolkit/models/base_model.py | 4 ++ toolkit/stable_diffusion_model.py | 4 ++ 5 files changed, 75 insertions(+), 40 deletions(-) diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index a3fdad2a..481d0858 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -9,7 +9,7 @@ from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.models.base_model import BaseModel from toolkit.basic import flush from toolkit.print import print_acc -from toolkit.prompt_utils import PromptEmbeds +from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) @@ -29,6 +29,7 @@ from .src.latent_norm import get_latent_norm from .src.pipeline import ( Ideogram4Pipeline, get_qwen3_vl_features, + pad_text_features, patchify_latents, predict_velocity, unpatchify_latents, @@ -176,15 +177,21 @@ class Ideogram4Model(BaseModel): self.patch_size = 2 self.vae_scale_factor = 8 - # Length the Qwen3-VL caption is padded/truncated to. Fixed so cached and - # batched embeddings always share a sequence length. + # Safety cap on caption token length (truncation only). Captions are stored + # per-sample at their natural length and padded to the batch max at the + # model call, so this is just an upper bound for very long JSON prompts. self.max_text_length = int( - self.model_config.model_kwargs.get("max_text_length", 512) + self.model_config.model_kwargs.get("max_text_length", 3072) ) self._latent_shift = None self._latent_scale = None + @property + def text_embedding_space_version(self): + # we changed the embeddings. invalidate cache. + return self.arch + "_te_v2" + @staticmethod def get_train_scheduler(): return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) @@ -340,8 +347,8 @@ class Ideogram4Model(BaseModel): self, pipeline: Ideogram4Pipeline, gen_config: GenerateImageConfig, - conditional_embeds: PromptEmbeds, - unconditional_embeds: PromptEmbeds, + conditional_embeds: AdvancedPromptEmbeds, + unconditional_embeds: AdvancedPromptEmbeds, generator: torch.Generator, extra: dict, ): @@ -371,7 +378,7 @@ class Ideogram4Model(BaseModel): self, latent_model_input: torch.Tensor, # (B, 128, gh, gw) timestep: torch.Tensor, # 0 to 1000 scale - text_embeddings: PromptEmbeds, + text_embeddings: AdvancedPromptEmbeds, **kwargs, ): if self.model.device == torch.device("cpu"): @@ -383,8 +390,10 @@ class Ideogram4Model(BaseModel): if t01.shape[0] != latent_model_input.shape[0]: t01 = t01.expand(latent_model_input.shape[0]) - llm_features = text_embeddings.text_embeds.to(self.device_torch) - text_mask = text_embeddings.attention_mask.to(self.device_torch) + # Pad the per-sample caption features to the batch max here. + llm_features, text_mask = pad_text_features( + text_embeddings.text_embeds, self.device_torch, self.torch_dtype + ) pred = predict_velocity( self.transformer, @@ -395,7 +404,7 @@ class Ideogram4Model(BaseModel): ) return pred - def get_prompt_embeds(self, prompt) -> PromptEmbeds: + def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds: if isinstance(prompt, str): prompt = [prompt] @@ -403,11 +412,11 @@ class Ideogram4Model(BaseModel): self.text_encoder.to(self.device_torch) device = self.text_encoder.device - pad_id = self.tokenizer.pad_token_id - if pad_id is None: - pad_id = self.tokenizer.eos_token_id or 0 - - token_id_list = [] + # Encode each caption at its natural length (no cross-sample padding) and + # store one feature tensor per batch item. Padding to a common length is + # deferred to the model call, so caching a prompt only stores its real + # length -- important for the long structured (JSON) captions. + features_list = [] for p in prompt: messages = [{"role": "user", "content": [{"type": "text", "text": p}]}] text = self.tokenizer.apply_chat_template( @@ -419,29 +428,19 @@ class Ideogram4Model(BaseModel): truncation=True, max_length=self.max_text_length, )["input_ids"] - token_id_list.append(ids) + if len(ids) == 0: + ids = [self.tokenizer.eos_token_id or 0] - seq_len = self.max_text_length - batch_size = len(token_id_list) - token_ids = torch.full((batch_size, seq_len), pad_id, dtype=torch.long) - attention_mask = torch.zeros((batch_size, seq_len), dtype=torch.long) - for b, ids in enumerate(token_id_list): - n = min(len(ids), seq_len) - # left pad - token_ids[b, seq_len - n :] = torch.tensor(ids[:n], dtype=torch.long) - attention_mask[b, seq_len - n :] = 1 + token_ids = torch.tensor([ids], dtype=torch.long, device=device) + attention_mask = torch.ones_like(token_ids) + pos_2d = (attention_mask.cumsum(dim=-1) - 1).clamp(min=0).to(torch.long) - token_ids = token_ids.to(device) - attention_mask = attention_mask.to(device) - pos_2d = (attention_mask.cumsum(dim=-1) - 1).clamp(min=0).to(torch.long) + features = get_qwen3_vl_features( + self.text_encoder, token_ids, attention_mask, pos_2d + ) # (1, Lt, D) + features_list.append(features[0].to(self.torch_dtype)) - features = get_qwen3_vl_features( - self.text_encoder, token_ids, attention_mask, pos_2d - ) - - pe = PromptEmbeds(features.to(self.torch_dtype)) - pe.attention_mask = attention_mask - return pe + return AdvancedPromptEmbeds(text_embeds=features_list) def get_model_has_grad(self): return False diff --git a/extensions_built_in/diffusion_models/ideogram4/src/pipeline.py b/extensions_built_in/diffusion_models/ideogram4/src/pipeline.py index 4e57ab41..c68e2d4f 100644 --- a/extensions_built_in/diffusion_models/ideogram4/src/pipeline.py +++ b/extensions_built_in/diffusion_models/ideogram4/src/pipeline.py @@ -123,6 +123,32 @@ def get_qwen3_vl_features( # --------------------------------------------------------------------------- +def pad_text_features( + features_list: List[torch.Tensor], + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Right-pad a list of per-sample (Lt_i, D) features into a batch. + + Captions are stored at their natural length (one tensor per batch item) and + only padded to the batch max here, right before the model call. Returns + ``(features (B, Lt, D), attention_mask (B, Lt))``; the mask is 1 for real + tokens and 0 for padding (which the transformer masks out anyway). + """ + lengths = [f.shape[0] for f in features_list] + max_len = max(lengths) + dim = features_list[0].shape[-1] + batch_size = len(features_list) + + features = torch.zeros(batch_size, max_len, dim, device=device, dtype=dtype) + mask = torch.zeros(batch_size, max_len, dtype=torch.long, device=device) + for i, f in enumerate(features_list): + n = f.shape[0] + features[i, :n] = f.to(device, dtype) + mask[i, :n] = 1 + return features, mask + + def predict_velocity( transformer: Ideogram4Transformer2DModel, latents: torch.Tensor, # (B, 128, gh, gw) @@ -281,11 +307,13 @@ class Ideogram4Pipeline: latents = latents.to(device, dtype=torch.float32) latents = latents * scheduler.init_noise_sigma - cond_feats = conditional_embeds.text_embeds.to(device, dtype=dtype) - cond_mask = conditional_embeds.attention_mask.to(device) + cond_feats, cond_mask = pad_text_features( + conditional_embeds.text_embeds, device, dtype + ) if do_cfg: - uncond_feats = unconditional_embeds.text_embeds.to(device, dtype=dtype) - uncond_mask = unconditional_embeds.attention_mask.to(device) + uncond_feats, uncond_mask = pad_text_features( + unconditional_embeds.text_embeds, device, dtype + ) for t in timesteps: t01 = (t / 1000.0).to(device).expand(latents.shape[0]) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 1d5536e0..9fd5f865 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -533,7 +533,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin size_database=self.size_database, dataset_root=dataset_folder, encode_control_in_text_embeddings=self.sd.encode_control_in_text_embeddings if self.sd else False, - text_embedding_space_version=self.sd.model_config.arch if self.sd else "sd1", + text_embedding_space_version=self.sd.text_embedding_space_version if self.sd else "sd1", te_padding_side=self.sd.te_padding_side if self.sd else "right", latent_space_version=latent_space_version, temporal_compression=temporal_compression, diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index d75c5ebf..8dd6c502 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -255,6 +255,10 @@ class BaseModel: @property def is_lumina2(self): return self.arch == 'lumina2' + + @property + def text_embedding_space_version(self): + return self.arch def get_bucket_divisibility(self): if self.vae is None: diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 2f1030c2..7d9d46f6 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -269,6 +269,10 @@ class StableDiffusion: def is_lumina2(self): return self.arch == 'lumina2' + @property + def text_embedding_space_version(self): + return self.arch + @property def unet_unwrapped(self): return unwrap_model(self.unet)