Merge pull request #3085 from bghira/bugfix/anyflow-uncond-batch-aliases

anyflow: swap unconditional aliases instead of popping them

Author: bghira

simpletuner/helpers/distillation/anyflow/distiller.py

@@ -950,10 +950,24 @@ def _unconditional_batch(prepared_batch: Dict[str, Any]) -> Dict[str, Any]:
         batch.pop("encoder_attention_mask", None)
         if torch.is_tensor(negative_mask):
             batch["encoder_attention_mask"] = negative_mask
-        # Some families (e.g. Ideogram) read model-specific conditioning aliases ahead of the
-        # generic keys; drop them so the swapped unconditional embeds/mask take effect.
-        for alias in ("prompt_embeds", "attention_mask", "attention_masks"):
-            batch.pop(alias, None)
+        # Some families read model-specific conditioning aliases ahead of (or instead of) the
+        # generic keys: Ideogram prefers `prompt_embeds` when present, and Flux requires it.
+        # Swap the aliases to the unconditional tensors rather than popping them.
+        if "prompt_embeds" in batch:
+            batch["prompt_embeds"] = negative
+        for alias in ("attention_mask", "attention_masks"):
+            if alias in batch:
+                if torch.is_tensor(negative_mask):
+                    batch[alias] = negative_mask
+                else:
+                    batch.pop(alias)
+        negative_pooled = prepared_batch.get("negative_add_text_embeds")
+        if torch.is_tensor(negative_pooled):
+            if "add_text_embeds" in batch:
+                batch["add_text_embeds"] = negative_pooled
+            added_cond_kwargs = batch.get("added_cond_kwargs")
+            if isinstance(added_cond_kwargs, dict) and "text_embeds" in added_cond_kwargs:
+                batch["added_cond_kwargs"] = {**added_cond_kwargs, "text_embeds": negative_pooled}
         # Lets families with a dedicated unconditional model (e.g. Ideogram) dispatch to it.
         batch["is_unconditional_pass"] = True
         return batch

simpletuner/helpers/models/common.py

@@ -5456,6 +5456,9 @@ def prepare_batch(self, batch: dict, state: dict) -> dict:
         negative_attention_mask = batch.get("negative_encoder_attention_mask")
         if negative_attention_mask is not None and hasattr(negative_attention_mask, "to"):
             batch["negative_encoder_attention_mask"] = negative_attention_mask.to(device=self.accelerator.device)
+        negative_pooled_embeds = batch.get("negative_add_text_embeds")
+        if negative_pooled_embeds is not None and hasattr(negative_pooled_embeds, "to"):
+            batch["negative_add_text_embeds"] = negative_pooled_embeds.to(**target_device_kwargs)
 
         # Process additional conditioning if provided
         pooled_embeds = batch.get("add_text_embeds")

simpletuner/helpers/training/collate.py

@@ -1335,6 +1335,7 @@ def resolve_local_source_audio_path(path: str | None) -> str | None:
         "negative_prompt_embeds": unconditional_text_encoder_outputs.get("prompt_embeds"),
         "negative_text_token_tags": unconditional_text_encoder_outputs.get("text_token_tags"),
         "negative_encoder_attention_mask": unconditional_text_encoder_outputs.get("attention_masks"),
+        "negative_add_text_embeds": unconditional_text_encoder_outputs.get("pooled_prompt_embeds"),
         "add_text_embeds": all_text_encoder_outputs.get("pooled_prompt_embeds"),
         "t5xxl_ids": all_text_encoder_outputs.get("t5xxl_ids"),
         "t5xxl_weights": all_text_encoder_outputs.get("t5xxl_weights"),

tests/helpers/distillation/test_anyflow_distiller.py

@@ -751,18 +751,30 @@ def test_unconditional_batch_drops_positive_attention_mask_when_negative_has_non
 
         self.assertNotIn("encoder_attention_mask", unconditional_batch)
 
-    def test_unconditional_batch_drops_model_specific_conditioning_aliases(self):
+    def test_unconditional_batch_swaps_model_specific_conditioning_aliases(self):
         batch = _prepared_batch()
         batch["prompt_embeds"] = torch.ones(2, 7, 4)
         batch["attention_mask"] = torch.ones(2, 7, dtype=torch.bool)
         batch["attention_masks"] = torch.ones(2, 7, dtype=torch.bool)
 
         unconditional_batch = AnyFlowDistiller._unconditional_batch(batch)
 
-        self.assertNotIn("prompt_embeds", unconditional_batch)
+        self.assertIs(unconditional_batch["prompt_embeds"], batch["negative_encoder_hidden_states"])
         self.assertNotIn("attention_mask", unconditional_batch)
         self.assertNotIn("attention_masks", unconditional_batch)
 
+    def test_unconditional_batch_swaps_alias_masks_when_negative_mask_present(self):
+        batch = _prepared_batch()
+        batch["prompt_embeds"] = torch.ones(2, 7, 4)
+        batch["attention_mask"] = torch.ones(2, 7, dtype=torch.bool)
+        batch["attention_masks"] = torch.ones(2, 7, dtype=torch.bool)
+        batch["negative_encoder_attention_mask"] = torch.zeros(2, 5, dtype=torch.bool)
+
+        unconditional_batch = AnyFlowDistiller._unconditional_batch(batch)
+
+        self.assertIs(unconditional_batch["attention_mask"], batch["negative_encoder_attention_mask"])
+        self.assertIs(unconditional_batch["attention_masks"], batch["negative_encoder_attention_mask"])
+
     def test_onpolicy_initializes_separate_discriminator_adapter_and_optimizer(self):
         model = _FlowModel()
         distiller = AnyFlowDistiller(
@@ -1071,7 +1083,7 @@ def test_get_scheduler_wraps_conditional_transformer_pipeline(self):
 
 
 class AnyFlowUnconditionalBatchTests(unittest.TestCase):
-    def test_unconditional_batch_drops_stale_mask_and_model_specific_aliases(self):
+    def test_unconditional_batch_drops_stale_mask_and_swaps_model_specific_aliases(self):
         batch = _prepared_batch()
         batch["encoder_attention_mask"] = torch.ones(2, 5, dtype=torch.long)
         batch["prompt_embeds"] = torch.ones(2, 7, 4)
@@ -1082,7 +1094,7 @@ def test_unconditional_batch_drops_stale_mask_and_model_specific_aliases(self):
 
         self.assertIs(unconditional_batch["encoder_hidden_states"], batch["negative_encoder_hidden_states"])
         self.assertNotIn("encoder_attention_mask", unconditional_batch)
-        self.assertNotIn("prompt_embeds", unconditional_batch)
+        self.assertIs(unconditional_batch["prompt_embeds"], batch["negative_encoder_hidden_states"])
         self.assertNotIn("attention_mask", unconditional_batch)
         self.assertNotIn("attention_masks", unconditional_batch)
 
@@ -1162,3 +1174,16 @@ def __init__(self):
         )
         with tempfile.TemporaryDirectory() as temp_dir:
             save_file(state, str(Path(temp_dir) / "adapter.safetensors"))
+
+    def test_unconditional_batch_swaps_negative_pooled_embeds(self):
+        batch = _prepared_batch()
+        batch["add_text_embeds"] = torch.ones(2, 8)
+        batch["added_cond_kwargs"] = {"text_embeds": batch["add_text_embeds"], "time_ids": torch.zeros(2, 6)}
+        batch["negative_add_text_embeds"] = torch.zeros(2, 8)
+
+        unconditional_batch = AnyFlowDistiller._unconditional_batch(batch)
+
+        self.assertIs(unconditional_batch["add_text_embeds"], batch["negative_add_text_embeds"])
+        self.assertIs(unconditional_batch["added_cond_kwargs"]["text_embeds"], batch["negative_add_text_embeds"])
+        self.assertIs(unconditional_batch["added_cond_kwargs"]["time_ids"], batch["added_cond_kwargs"]["time_ids"])
+        self.assertIs(batch["added_cond_kwargs"]["text_embeds"], batch["add_text_embeds"])