Improved the prompt handeling of ideogram4 model. Now used advanced prompts class to store them smaller and allow longer prompts

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

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])

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,

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:

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)