Add unconditional lora support so Ideogram 4 inference will more closely resemble the full pipeline results. I pushed a finetuned unconditional lora to the hub as an adapter.

extensions_built_in/diffusion_models/ideogram4/ideogram4.py

@@ -5,8 +5,9 @@ import torch
 import yaml
 from safetensors.torch import load_file, save_file
 
-from toolkit.config_modules import GenerateImageConfig, ModelConfig
+from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig
 from toolkit.models.base_model import BaseModel
+from toolkit.lora_special import LoRASpecialNetwork
 from toolkit.basic import flush
 from toolkit.print import print_acc
 from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
@@ -189,6 +190,11 @@ class Ideogram4Model(BaseModel):
         self._latent_shift = None
         self._latent_scale = None
 
+        # Optional LoRA that is only switched on during the unconditional (negative)
+        # CFG pass. Loaded from model_config.unconditional_lora_path if set; stays
+        # inactive everywhere else (training, conditional pass).
+        self.unconditional_lora: Optional[LoRASpecialNetwork] = None
+
     @property
     def text_embedding_space_version(self):
         # we changed the embeddings. invalidate cache.
@@ -267,6 +273,88 @@ class Ideogram4Model(BaseModel):
         vae.requires_grad_(False)
         return vae
 
+    def load_unconditional_lora(self, transformer: Ideogram4Transformer2DModel):
+        """Load the unconditional-pass LoRA and leave it applied but inactive.
+
+        The adapter is wired into the transformer via ``apply_to`` (no merge) so
+        the pipeline can flip ``is_active`` on for the unconditional CFG pass only.
+        It never affects the conditional pass or training, where it stays inactive.
+        """
+        lora_path = self.model_config.unconditional_lora_path
+        self.print_and_status_update(f"Loading unconditional LoRA from {lora_path}")
+
+        if not os.path.exists(lora_path):
+            # assume it is a "repo/owner/filename.safetensors" hub path
+            lora_splits = lora_path.split("/")
+            if len(lora_splits) != 3:
+                raise ValueError(
+                    f"Unconditional LoRA path {lora_path} is not a valid local path "
+                    "or hub path."
+                )
+            repo_id = "/".join(lora_splits[:2])
+            filename = lora_splits[2]
+            try:
+                lora_path = huggingface_hub.hf_hub_download(
+                    repo_id=repo_id, filename=filename, token=HF_TOKEN
+                )
+                self.model_config.unconditional_lora_path = lora_path
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to download unconditional LoRA from {lora_path}: {e}"
+                )
+
+        # Detect the LoRA rank from the first down-projection weight in the file.
+        lora_state_dict = load_file(lora_path)
+        lora_dim = None
+        for key, value in lora_state_dict.items():
+            if key.endswith("lora_A.weight") or key.endswith("lora_down.weight"):
+                lora_dim = int(value.shape[0])
+                break
+        if lora_dim is None:
+            raise ValueError(
+                f"Could not determine LoRA rank from {lora_path}: no lora_A/lora_down "
+                "weights found."
+            )
+
+        # transformer_only=False so every nn.Linear in the model is targeted (not
+        # just the transformer blocks) -- the extraction script factors all linears,
+        # so the adapter must wrap all of them to load every key.
+        network_config = NetworkConfig(
+            type="lora",
+            linear=lora_dim,
+            linear_alpha=lora_dim,
+            transformer_only=False,
+        )
+        network = LoRASpecialNetwork(
+            text_encoder=None,
+            unet=transformer,
+            lora_dim=lora_dim,
+            multiplier=1.0,
+            alpha=lora_dim,
+            # train_unet just gates module creation here; the network is applied,
+            # kept inactive, and never trained (the pipeline only toggles is_active).
+            train_unet=True,
+            train_text_encoder=False,
+            network_config=network_config,
+            network_type="lora",
+            transformer_only=False,
+            is_transformer=True,
+            target_lin_modules=self.target_lora_modules,
+            # base_model_ref lets load_weights run convert_lora_weights_before_load
+            # so saved "diffusion_model." keys map back to "transformer.".
+            base_model=self,
+        )
+        network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True)
+        network.force_to(self.device_torch, dtype=self.torch_dtype)
+        network._update_torch_multiplier()
+        network.load_weights(lora_path)
+        network.eval()
+
+        # Inactive by default; the pipeline flips this on only for the uncond pass.
+        network.is_active = False
+        self.unconditional_lora = network
+        self.print_and_status_update("Unconditional LoRA loaded (inactive)")
+
     def load_model(self):
         dtype = self.torch_dtype
         self.print_and_status_update("Loading Ideogram4 model")
@@ -290,7 +378,11 @@ class Ideogram4Model(BaseModel):
                 transformer,
                 self.device_torch,
                 offload_percent=self.model_config.layer_offloading_transformer_percent,
-                ignore_modules=[transformer.rotary_emb.inv_freq, transformer.input_proj, transformer.llm_cond_proj],
+                ignore_modules=[
+                    transformer.rotary_emb.inv_freq,
+                    transformer.input_proj,
+                    transformer.llm_cond_proj,
+                ],
             )
         elif self.model_config.low_vram:
             self.print_and_status_update("Moving transformer to CPU")
@@ -337,6 +429,10 @@ class Ideogram4Model(BaseModel):
         self.tokenizer = tokenizer
         self.model = transformer
         self.pipeline = Ideogram4Pipeline(self)
+
+        if self.model_config.unconditional_lora_path is not None:
+            self.load_unconditional_lora(transformer)
+
         self.print_and_status_update("Model Loaded")
 
     # ------------------------------------------------------------------

extensions_built_in/diffusion_models/ideogram4/src/pipeline.py

@@ -325,15 +325,28 @@ class Ideogram4Pipeline:
             )
             uncond_mask = torch.zeros(batch_size, 0, dtype=torch.long, device=device)
 
+        # The unconditional LoRA (if present) must be active *only* on the
+        # unconditional pass. We force it off before each conditional pass since the
+        # outer sampling context (``with network:``) may switch it on globally.
+        uncond_lora = getattr(model, "unconditional_lora", None)
+
         for t in timesteps:
             t01 = (t / 1000.0).to(device).expand(latents.shape[0])
+            if uncond_lora is not None:
+                uncond_lora.is_active = False
             v_cond = predict_velocity(
                 transformer, latents.to(dtype), t01, cond_feats, cond_mask
             )
             if do_cfg:
-                v_uncond = predict_velocity(
+                if uncond_lora is not None:
-                    transformer, latents.to(dtype), t01, uncond_feats, uncond_mask
+                    uncond_lora.is_active = True
-                )
+                try:
+                    v_uncond = predict_velocity(
+                        transformer, latents.to(dtype), t01, uncond_feats, uncond_mask
+                    )
+                finally:
+                    if uncond_lora is not None:
+                        uncond_lora.is_active = False
                 v = v_uncond + guidance_scale * (v_cond - v_uncond)
             else:
                 v = v_cond

toolkit/config_modules.py

@@ -620,6 +620,9 @@ class ModelConfig:
         # mainly for decompression loras for distilled models
         self.assistant_lora_path = kwargs.get('assistant_lora_path', None)
         self.inference_lora_path = kwargs.get('inference_lora_path', None)
+        # a lora that stays inactive except during the unconditional (negative)
+        # CFG pass -- used to learn the unconditional branch without a second model
+        self.unconditional_lora_path = kwargs.get('unconditional_lora_path', None)
         self.latent_space_version = kwargs.get('latent_space_version', None)
 
         # only for SDXL models for now

ui/src/app/jobs/new/SimpleJob.tsx

@@ -307,6 +307,20 @@ export default function SimpleJob({
                 placeholder=""
               />
             )}
+            {modelArch?.additionalSections?.includes('model.unconditional_lora_path') && (
+              <TextInput
+                label="Unconditional Adapter Path"
+                value={jobConfig.config.process[0].model.unconditional_lora_path ?? ''}
+                docKey="config.process[0].model.unconditional_lora_path"
+                onChange={(value: string | undefined) => {
+                  if (value?.trim() === '') {
+                    value = undefined;
+                  }
+                  setJobConfig(value, 'config.process[0].model.unconditional_lora_path');
+                }}
+                placeholder=""
+              />
+            )}
             {modelArch?.additionalSections?.includes('model.low_vram') && (
               <FormGroup label="Options">
                 <Checkbox

ui/src/app/jobs/new/options.ts

@@ -32,6 +32,7 @@ type AdditionalSections =
   | 'model.low_vram'
   | 'model.qie.match_target_res'
   | 'model.assistant_lora_path'
+  | 'model.unconditional_lora_path'
   | 'ideogram_4_prompt';
 
 type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio';
@@ -993,6 +994,10 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].network.conv_alpha': [undefined, 16],
       'config.process[0].model.low_vram': [true, false],
       'config.process[0].sample': [defaultIdeogramSamplesConfig, defaultSampleConfig],
+      'config.process[0].model.unconditional_lora_path': [
+        'ostris/ideogram_4_unconditional_lora/ideogram_4_unconditional_lora_r16.safetensors',
+        undefined,
+      ],
     },
     disableSections: [
       'network.conv',
@@ -1001,6 +1006,7 @@ export const modelArchs: ModelArch[] = [
       'model.low_vram',
       'model.layer_offloading',
       'ideogram_4_prompt',
+      'model.unconditional_lora_path',
     ],
     hasMultiLinePrompts: true,
   },

ui/src/types.ts

@@ -171,6 +171,7 @@ export interface ModelConfig {
   layer_offloading_transformer_percent?: number;
   layer_offloading_text_encoder_percent?: number;
   assistant_lora_path?: string;
+  unconditional_lora_path?: string;
   compile?: boolean;
   block_compile?: boolean;
   compile_mode?: 'default' | 'max-autotune' | 'fastest';

version.py

@@ -1 +1 @@
-VERSION = "0.10.12"
+VERSION = "0.10.13"