Merge remote-tracking branch 'origin/main' into script/train-minimax-music-rvq-encoder

Author: bghira

simpletuner/helpers/caching/text_embeds.py

@@ -678,10 +678,13 @@ def compute_prompt_embeddings_with_model(
                                 f"\n-> error: {e}"
                                 f"\n-> id: {self.id}, data_backend id: {self.data_backend.id}"
                             )
-                            raise Exception(
-                                "Cache retrieval for text embed file failed. Ensure your dataloader config value for "
-                                "skip_file_discovery does not contain 'text', and that preserve_data_backend_cache is "
-                                "disabled or unset."
+                            raise RuntimeError(
+                                "Cache retrieval for text embed file failed. For Webshart or other multi-caption "
+                                "datasets, this can mean the existing text embedding cache was generated without every "
+                                "caption variant that training can request. Set text_cache_ondemand: true to compute "
+                                "missing embeds during training, or clear the text embedding cache and rerun precache. "
+                                "Also ensure skip_file_discovery does not contain 'text' and preserve_data_backend_cache "
+                                "is disabled or unset."
                             ) from e
                         if self.model.requires_text_embed_image_context() and not record.get("metadata"):
                             raise ValueError(
@@ -715,6 +718,8 @@ def compute_prompt_embeddings_with_model(
                         prompt_contexts=prompt_contexts,
                         is_validation=is_validation,
                     )
+                    text_encoder_output = self.model.pack_text_embeddings_for_cache(text_encoder_output)
+                    text_encoder_output = self.model.unpack_text_embeddings_from_cache(text_encoder_output)
                     logger.debug(
                         f"Filename {filename} prompt embeds: {gather_dict_of_tensors_shapes(tensors=text_encoder_output)}, keys: {text_encoder_output.keys()}"
                     )

simpletuner/helpers/data_backend/webshart.py

@@ -338,14 +338,18 @@ def _sample_index_for_filename(self, shard_idx: int, filename: str) -> Optional[
             }
         return self._shard_sample_index_cache[shard_idx].get(str(filename))
 
-    def get_caption(self, image_path: str) -> Optional[str]:
+    def get_caption(self, image_path: str) -> Optional[Union[str, List[str], dict]]:
         if not self.is_sample_id(image_path):
             return None
 
         sample_ref = self.parse_sample_id(image_path)
         sample_metadata = self.get_shard_metadata(sample_ref.shard_idx).get(sample_ref.filename, {}) or {}
         caption = sample_metadata.get("captions")
         if caption:
+            if isinstance(caption, dict):
+                return caption
+            if isinstance(caption, list):
+                return [str(item).strip() for item in caption if item is not None and str(item).strip()]
             return str(caption).strip()
 
         caption_filename = Path(sample_ref.filename).with_suffix(".txt").name

simpletuner/helpers/metadata/backends/webshart.py

@@ -91,7 +91,7 @@ def __init__(
         if self.dataset_type not in {DatasetType.IMAGE, DatasetType.VIDEO, DatasetType.CONDITIONING, DatasetType.EVAL}:
             raise ValueError("WebshartMetadataBackend supports image, video, conditioning, and eval datasets only.")
 
-        self.caption_cache: Dict[str, Union[str, List[str]]] = {}
+        self.caption_cache: Dict[str, Union[str, List[str], dict]] = {}
 
         context = accelerator.main_process_first() if hasattr(accelerator, "main_process_first") else nullcontext()
         with context:
@@ -119,7 +119,7 @@ def _sync_image_files_with_buckets(self) -> None:
             return
         StateTracker.set_image_files([("", [], sample_ids)], data_backend_id=self.data_backend.id)
 
-    def caption_cache_entry(self, index: str) -> Optional[Union[str, List[str]]]:
+    def caption_cache_entry(self, index: str) -> Optional[Union[str, List[str], dict]]:
         index = self.data_backend.normalize_sample_id(index)
         caption = self.caption_cache.get(index, None)
         if caption is not None:

simpletuner/helpers/models/ideogram/model.py

@@ -524,11 +524,12 @@ def collate_prompt_embeds(self, text_encoder_output: list[dict]) -> dict:
                 attention_mask = item.get("attention_masks")
             if prompt_embeds.dim() == 3:
                 prompt_embeds = prompt_embeds.squeeze(0)
+            prompt_embeds = prompt_embeds.to("cpu")
             if attention_mask is None:
-                attention_mask = torch.ones(prompt_embeds.shape[0], dtype=torch.bool, device=prompt_embeds.device)
+                attention_mask = torch.ones(prompt_embeds.shape[0], dtype=torch.bool, device="cpu")
             elif attention_mask.dim() == 2:
                 attention_mask = attention_mask.squeeze(0)
-            attention_mask = attention_mask.to(dtype=torch.bool)
+            attention_mask = attention_mask.to(device="cpu", dtype=torch.bool)
             length = int(attention_mask.sum().item())
             prompt_embeds = prompt_embeds[:length]
             attention_mask = attention_mask[:length]

simpletuner/helpers/models/minimaxh3/model.py

@@ -1592,6 +1592,18 @@ def model_predict(self, prepared_batch):
             batch_size,
             self.accelerator.device,
         )
+        if text_valid_mask is not None and bool(text_valid_mask.any()):
+            # Trailing text padding forces a dense attention mask that the context-parallel
+            # attention dispatch cannot accept; trim it away and drop the mask when nothing
+            # in the batch is padded after trimming.
+            keep_len = int(text_valid_mask.any(dim=0).nonzero().max().item()) + 1
+            if keep_len < text_seq_len:
+                encoder_hidden_states = encoder_hidden_states[:, :keep_len]
+                text_token_tags = text_token_tags[:, :keep_len]
+                text_valid_mask = text_valid_mask[:, :keep_len]
+                text_seq_len = keep_len
+            if bool(text_valid_mask.all()):
+                text_valid_mask = None
         packed_target_video = patchify_video_latents(noisy_latents, patch_size).view(
             batch_size, -1, channels * patch_product
         )

simpletuner/helpers/prompts.py

@@ -266,6 +266,38 @@ def __init__(
         self.text_encoders = text_encoders
         self.tokenizers = tokenizers
 
+    @staticmethod
+    def _caption_payload_is_multi_value(caption) -> bool:
+        return isinstance(caption, (list, tuple, dict, numpy.ndarray, pd.Series))
+
+    @staticmethod
+    def _normalize_caption_payload(caption) -> list[str]:
+        if caption is None:
+            return []
+        if isinstance(caption, bytes):
+            caption = caption.decode("utf-8")
+        if isinstance(caption, str):
+            caption = caption.strip()
+            return [caption] if caption else []
+        if isinstance(caption, dict):
+            captions = []
+            for value in caption.values():
+                captions.extend(PromptHandler._normalize_caption_payload(value))
+            return captions
+        if isinstance(caption, (list, tuple, numpy.ndarray, pd.Series)):
+            captions = []
+            for value in caption:
+                captions.extend(PromptHandler._normalize_caption_payload(value))
+            return captions
+        caption = str(caption).strip()
+        return [caption] if caption else []
+
+    @staticmethod
+    def _restore_caption_payload_shape(caption, caption_values: list[str]):
+        if PromptHandler._caption_payload_is_multi_value(caption):
+            return caption_values
+        return caption_values[0] if caption_values else ""
+
     @staticmethod
     def retrieve_prompt_column_from_parquet(
         sampler_backend_id: str,
@@ -338,18 +370,10 @@ def prepare_instance_prompt_from_parquet(
             raise CaptionNotFoundError(
                 f"Could not locate caption for image {image_path} in sampler_backend {sampler_backend_id} with filename column {filename_column}, caption column {caption_column}, and a parquet database with {len(parquet_db)} entries."
             )
-        if type(image_caption) == bytes:
-            image_caption = image_caption.decode("utf-8")
-        if type(image_caption) == str:
-            image_caption = image_caption.strip()
-        if type(image_caption) in (list, tuple, numpy.ndarray, pd.Series):
-            image_caption = [str(item).strip() for item in image_caption if item is not None]
+        caption_values = PromptHandler._normalize_caption_payload(image_caption)
         if prepend_instance_prompt:
-            if type(image_caption) == list:
-                image_caption = [instance_prompt + " " + x for x in image_caption]
-            else:
-                image_caption = instance_prompt + " " + image_caption
-        return image_caption
+            caption_values = [instance_prompt + " " + x for x in caption_values]
+        return PromptHandler._restore_caption_payload_shape(image_caption, caption_values)
 
     @staticmethod
     def prepare_instance_prompt_from_filename(
@@ -460,22 +484,13 @@ def prepare_instance_prompt_from_huggingface(
             if caption is None:
                 raise CaptionNotFoundError(f"Could not find caption for {image_path} in HuggingFace dataset")
 
-        # Process the caption
-        if isinstance(caption, bytes):
-            caption = caption.decode("utf-8")
-        if isinstance(caption, str):
-            caption = caption.strip()
-        if isinstance(caption, (list, tuple, numpy.ndarray, pd.Series)):
-            caption = [str(item).strip() for item in caption if item is not None]
+        caption_values = PromptHandler._normalize_caption_payload(caption)
 
         # Prepend instance prompt if requested
         if prepend_instance_prompt and instance_prompt:
-            if isinstance(caption, list):
-                caption = [instance_prompt + " " + c for c in caption]
-            else:
-                caption = instance_prompt + " " + caption
+            caption_values = [instance_prompt + " " + c for c in caption_values]
 
-        return caption
+        return PromptHandler._restore_caption_payload_shape(caption, caption_values)
 
     @staticmethod
     def prepare_instance_prompt_from_webshart(
@@ -503,20 +518,12 @@ def prepare_instance_prompt_from_webshart(
         if caption is None:
             raise CaptionNotFoundError(f"Could not find caption for {image_path} in Webshart dataset")
 
-        if isinstance(caption, bytes):
-            caption = caption.decode("utf-8")
-        if isinstance(caption, str):
-            caption = caption.strip()
-        if isinstance(caption, (list, tuple, numpy.ndarray, pd.Series)):
-            caption = [str(item).strip() for item in caption if item is not None]
+        caption_values = PromptHandler._normalize_caption_payload(caption)
 
         if prepend_instance_prompt and instance_prompt:
-            if isinstance(caption, list):
-                caption = [instance_prompt + " " + c for c in caption]
-            else:
-                caption = instance_prompt + " " + caption
+            caption_values = [instance_prompt + " " + c for c in caption_values]
 
-        return caption
+        return PromptHandler._restore_caption_payload_shape(caption, caption_values)
 
     @staticmethod
     def magic_prompt(
@@ -625,7 +632,7 @@ def magic_prompt(
 
         # Apply shuffle expansion if enabled
         if shuffle_enabled and instance_prompt:
-            caption_values = instance_prompt if isinstance(instance_prompt, list) else [instance_prompt]
+            caption_values = PromptHandler._normalize_caption_payload(instance_prompt)
             expanded_captions = []
             for value in caption_values:
                 shuffled = CaptionShuffler.expand_with_shuffles(value, caption_shuffle_config)
@@ -795,7 +802,7 @@ def get_all_captions(
                     logger.error(f"Could not load caption for image {image_path}: {e}")
                     images_missing_captions.append(image_path)
                 else:
-                    caption_values = caption if isinstance(caption, (tuple, list, dict)) else [caption]
+                    caption_values = PromptHandler._normalize_caption_payload(caption)
 
                     # Apply shuffle expansion if enabled
                     if shuffle_enabled:

tests/test_ideogram4_prompting.py

@@ -109,6 +109,28 @@ def test_text_embed_conversion_and_collation_pad_at_batch_time(self):
         self.assertEqual(converted["prompt_embeds"].shape, (1, 2, 4))
         self.assertEqual(negative["negative_prompt_embeds"].shape, (1, 2, 4))
 
+    def test_text_embed_collation_normalizes_cache_batch_to_cpu(self):
+        model = Ideogram4.__new__(Ideogram4)
+        device = torch.device("cpu")
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        elif torch.backends.mps.is_available():
+            device = torch.device("mps")
+        first = {
+            "prompt_embeds": torch.ones(1, 2, 4, device=device),
+            "attention_mask": torch.ones(1, 2, dtype=torch.bool, device=device),
+        }
+        second = {
+            "prompt_embeds": torch.ones(1, 1, 4),
+            "attention_masks": torch.ones(1, 1, dtype=torch.bool),
+        }
+
+        collated = model.collate_prompt_embeds([first, second])
+
+        self.assertEqual(collated["prompt_embeds"].device.type, "cpu")
+        self.assertEqual(collated["attention_masks"].device.type, "cpu")
+        self.assertEqual(collated["prompt_embeds"].shape, (2, 2, 4))
+
     def test_text_embed_cache_projection_uses_projection_component(self):
         model = Ideogram4.__new__(Ideogram4)
         model.config = types.SimpleNamespace(

tests/test_prompts.py

@@ -13,6 +13,14 @@ def list_files(self, instance_data_dir=None, file_extensions=None):
         return self._files
 
 
+class _DummyMetadataBackend:
+    def __init__(self, captions):
+        self._captions = captions
+
+    def caption_cache_entry(self, image_path):
+        return self._captions.get(image_path)
+
+
 class PromptHandlerTests(unittest.TestCase):
     def test_instanceprompt_returns_entry_per_image(self):
         backend = _DummyBackend(["a.jpg", "b.jpg", "c.jpg"])
@@ -33,6 +41,50 @@ def test_instanceprompt_returns_entry_per_image(self):
         self.assertEqual(captions, ["minecraft", "minecraft", "minecraft"])
         self.assertEqual(paths, ["a.jpg", "b.jpg", "c.jpg"])
 
+    def test_webshart_get_all_captions_expands_structured_caption_variants(self):
+        backend = _DummyBackend(["webshart://0/1/first.jpg", "webshart://0/2/second.jpg"])
+        metadata_backend = _DummyMetadataBackend(
+            {
+                "webshart://0/1/first.jpg": ["first primary", "first alternate"],
+                "webshart://0/2/second.jpg": {
+                    "primary": "second primary",
+                    "alternates": ["second alternate"],
+                },
+            }
+        )
+
+        with (
+            patch("simpletuner.helpers.prompts.StateTracker.get_data_backend_config", return_value={}),
+            patch("simpletuner.helpers.prompts.StateTracker.get_image_files", return_value=None),
+            patch(
+                "simpletuner.helpers.prompts.StateTracker.get_data_backend",
+                return_value={"metadata_backend": metadata_backend},
+            ),
+        ):
+            captions, missing, paths = PromptHandler.get_all_captions(
+                instance_data_dir="",
+                use_captions=True,
+                prepend_instance_prompt=False,
+                data_backend=backend,
+                caption_strategy="webshart",
+                return_image_paths=True,
+            )
+
+        self.assertEqual(missing, [])
+        self.assertEqual(
+            captions,
+            ["first primary", "first alternate", "second primary", "second alternate"],
+        )
+        self.assertEqual(
+            paths,
+            [
+                "webshart://0/1/first.jpg",
+                "webshart://0/1/first.jpg",
+                "webshart://0/2/second.jpg",
+                "webshart://0/2/second.jpg",
+            ],
+        )
+
 
 if __name__ == "__main__":
     unittest.main()