Merge pull request #3140 from bghira/feature/better-music-audio-model-card-data

add XM and NextLat info to model card, plus audio-specific fixes

Author: bghira

simpletuner/helpers/models/common.py

@@ -6913,6 +6913,12 @@ def custom_model_card_schedule_info(self):
         """
         return []
 
+    def custom_model_card_training_mode_info(self, args) -> str:
+        """
+        Override this in a subclass to add model-specific training mode details to model cards.
+        """
+        return ""
+
     def custom_model_card_code_example(self, repo_id: str = None) -> str:
         """
         Override this to provide custom code examples for model cards.

simpletuner/helpers/models/minimaxmusic/model.py

@@ -318,6 +318,18 @@ def __init__(self, config, accelerator):
             self.TEXT_ENCODER_CONFIGURATION = {}
             self.DEFAULT_LORA_TARGET = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
 
+    def custom_model_card_training_mode_info(self, args) -> str:
+        train_component = str(getattr(args, "minimax_music_train_component", "transformer") or "transformer")
+        component_labels = {
+            "language_model": "language_model (global LM / RVQ planner)",
+            "transformer": "transformer (DiT/audio denoiser)",
+        }
+        lines = [f"- MiniMax Music train component: `{component_labels.get(train_component, train_component)}`"]
+        lm_max_frames = getattr(args, "minimax_music_lm_max_frames", None)
+        if lm_max_frames:
+            lines.append(f"- MiniMax Music LM max frames: `{lm_max_frames}`")
+        return "\n".join(lines)
+
     @classmethod
     def max_swappable_blocks(cls, config=None) -> Optional[int]:
         return 35

simpletuner/helpers/publishing/metadata.py

@@ -433,6 +433,57 @@ def model_card_note(args):
     return f"\n**Note:** {note_contents}\n"
 
 
+def _model_card_bool(args, name: str) -> bool:
+    value = getattr(args, name, False)
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, (int, float)):
+        return bool(value)
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return False
+
+
+def _model_card_optional_value(args, name: str):
+    value = getattr(args, name, None)
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    return None
+
+
+def _model_card_training_modes(args, model: Optional[ModelFoundation] = None) -> str:
+    lines = []
+    model_specific_info = ""
+    if model and hasattr(model, "custom_model_card_training_mode_info"):
+        model_specific_info = model.custom_model_card_training_mode_info(args) or ""
+        if not isinstance(model_specific_info, str):
+            model_specific_info = ""
+        model_specific_info = model_specific_info.strip()
+
+    if _model_card_bool(args, "nextlat_enabled"):
+        lines.append("- NextLat: Enabled")
+        lines.append(f"  - Block index: `{_model_card_optional_value(args, 'nextlat_block_index')}`")
+        lines.append(f"  - Weight: `{_model_card_optional_value(args, 'nextlat_weight')}`")
+        lines.append(f"  - State loss: `{_model_card_optional_value(args, 'nextlat_state_loss')}`")
+        lines.append(f"  - KL weight: `{_model_card_optional_value(args, 'nextlat_kl_weight')}`")
+
+    if _model_card_bool(args, "xm_enabled"):
+        lines.append("- XM: Enabled")
+        lines.append(f"  - Candidate count: `{_model_card_optional_value(args, 'xm_candidate_count')}`")
+        lines.append(f"  - Selection scope: `{_model_card_optional_value(args, 'xm_selection_scope')}`")
+        lines.append(f"  - Training target: `{_model_card_optional_value(args, 'xm_training_target')}`")
+        lines.append(f"  - Block size: `{_model_card_optional_value(args, 'xm_block_size')}`")
+
+    if not lines and not model_specific_info:
+        return ""
+    blocks = []
+    if model_specific_info:
+        blocks.append(model_specific_info)
+    if lines:
+        blocks.append("\n".join(lines))
+    return "## Training modes\n\n" + "\n".join(blocks) + "\n\n"
+
+
 def save_metadata_sample(
     image_path: str,
     image: Union[Image.Image, np.ndarray, list, str, torch.Tensor],
@@ -772,6 +823,36 @@ def _add_widget_entries(media, asset_prefix: str):
             gallery_intro = "You can find some example audio samples in the following gallery:"
         else:
             gallery_intro = "You can find some example images in the following gallery:"
+    gallery_section = f"{gallery_intro}\n\n<Gallery />\n" if has_media else ""
+    validation_disabled = _model_card_bool(args, "validation_disable")
+    validation_intro = (
+        "Validation was disabled during training."
+        if validation_disabled
+        else (
+            "The main validation prompt used during training was:"
+            if prompt
+            else (
+                "Validation used ground-truth images as an input for partial denoising (img2img)."
+                if args.validation_using_datasets
+                else "No validation prompt was used during training."
+            )
+        )
+    )
+    validation_prompt_block = f"```\n{prompt}\n```\n" if prompt and not validation_disabled else ""
+    validation_settings = ""
+    if not validation_disabled:
+        validation_settings = f"""## Validation settings
+- CFG: `{StateTracker.get_args().validation_guidance}`
+- CFG Rescale: `{StateTracker.get_args().validation_guidance_rescale}`
+- Steps: `{StateTracker.get_args().validation_num_inference_steps}`
+- Sampler: `{_validation_scheduler_label(model, StateTracker.get_args())}`
+- Seed: `{StateTracker.get_args().validation_seed}`
+- Resolution{'s' if ',' in str(StateTracker.get_args().validation_resolution) else ''}: `{str(StateTracker.get_args().validation_resolution)}`
+{f"- Skip-layer guidance: {_skip_layers(args)}" if args.model_family in ['sd3', 'flux'] else ''}
+
+Note: The validation settings are not necessarily the same as the [training settings](#training-settings).
+
+"""
     sage_usage = getattr(args.sageattention_usage, "value", args.sageattention_usage)
     license_metadata = _license_metadata(model)
     yaml_content = f"""---
@@ -800,26 +881,11 @@ def _add_widget_entries(media, asset_prefix: str):
 
 This is a {model_type(args)} derived from [{base_model}](https://huggingface.co/{base_model}).
 
-{'The main validation prompt used during training was:' if prompt else 'Validation used ground-truth images as an input for partial denoising (img2img).' if args.validation_using_datasets else 'No validation prompt was used during training.'}
-{'```' if prompt else ''}
-{prompt}
-{'```' if prompt else ''}
+{validation_intro}
+{validation_prompt_block}
 
 {model_card_note(args)}
-## Validation settings
-- CFG: `{StateTracker.get_args().validation_guidance}`
-- CFG Rescale: `{StateTracker.get_args().validation_guidance_rescale}`
-- Steps: `{StateTracker.get_args().validation_num_inference_steps}`
-- Sampler: `{_validation_scheduler_label(model, StateTracker.get_args())}`
-- Seed: `{StateTracker.get_args().validation_seed}`
-- Resolution{'s' if ',' in str(StateTracker.get_args().validation_resolution) else ''}: `{str(StateTracker.get_args().validation_resolution)}`
-{f"- Skip-layer guidance: {_skip_layers(args)}" if args.model_family in ['sd3', 'flux'] else ''}
-
-Note: The validation settings are not necessarily the same as the [training settings](#training-settings).
-
-{gallery_intro}\n
-
-<Gallery />
+{validation_settings}{gallery_section}
 
 The text encoder {'**was**' if train_text_encoder else '**was not**'} trained.
 {'You may reuse the base model text encoder for inference.' if not train_text_encoder else 'If the text encoder from this repository is not used at inference time, unexpected or bad results could occur.'}
@@ -849,6 +915,7 @@ def _add_widget_entries(media, asset_prefix: str):
 {('- SLA: Enabled (you MUST use SLA for inference; sla_attention.pt contains attention weights)') if StateTracker.get_args().attention_mechanism == 'sla' else ''}
 {lora_info(args=StateTracker.get_args())}
 
+{_model_card_training_modes(args, model=model)}
 ## Datasets
 
 {datasets_str}

tests/test_model_card.py

@@ -41,6 +41,7 @@ def setUp(self):
         self.args.model_type = "lora"
         self.args.model_family = "sdxl"
         self.args.validation_prompt = "A test prompt"
+        self.args.validation_disable = False
         self.args.validation_negative_prompt = "A negative prompt"
         self.args.validation_num_inference_steps = 50
         self.args.validation_guidance = 7.5
@@ -84,6 +85,18 @@ def setUp(self):
         self.args.t5_padding = "unmodified"
         self.args.enable_xformers_memory_efficient_attention = False
         self.args.attention_mechanism = "diffusers"
+        self.args.minimax_music_train_component = None
+        self.args.minimax_music_lm_max_frames = None
+        self.args.nextlat_enabled = False
+        self.args.nextlat_block_index = -1
+        self.args.nextlat_weight = 0.0
+        self.args.nextlat_state_loss = "smooth_l1"
+        self.args.nextlat_kl_weight = 0.0
+        self.args.xm_enabled = False
+        self.args.xm_candidate_count = 1
+        self.args.xm_selection_scope = "sample"
+        self.args.xm_training_target = "noise"
+        self.args.xm_block_size = 0
         self.mock_model = MagicMock(MODEL_TYPE=MagicMock(value="unet"))
 
     def test_model_imports(self):
@@ -293,6 +306,93 @@ def test_audio_model_families_use_audio_pipeline_tag(self):
                 self.assertEqual(_pipeline_tag(self.args), "text-to-audio")
                 self.assertEqual(_secondary_pipeline_tag(self.args), "audio")
 
+    def test_minimax_music_model_card_reports_modes_and_disabled_validation(self):
+        self.args.model_family = "minimaxmusic"
+        self.args.model_type = "lora"
+        self.args.lora_type = "standard"
+        self.args.peft_lora_mode = "standard"
+        self.args.controlnet = False
+        self.args.control = False
+        self.args.validation_disable = True
+        self.args.model_card_note = ""
+        self.args.minimax_music_train_component = "language_model"
+        self.args.minimax_music_lm_max_frames = 128
+        self.args.nextlat_enabled = True
+        self.args.nextlat_block_index = -1
+        self.args.nextlat_weight = 0.1
+        self.args.nextlat_state_loss = "smooth_l1"
+        self.args.nextlat_kl_weight = 0.0
+        self.args.xm_enabled = True
+        self.args.xm_candidate_count = 2
+        self.args.xm_selection_scope = "block"
+        self.args.xm_training_target = "route"
+        self.args.xm_block_size = 16
+        model = MagicMock(
+            MODEL_LICENSE="other",
+            PREDICTION_TYPE=SimpleNamespace(value="autoregressive_next_token"),
+            gligen=False,
+        )
+        model.validation_audio_sample_rate.return_value = 44100
+        model.custom_model_card_schedule_info.return_value = ""
+        model.custom_model_card_code_example.return_value = "```python\npass\n```"
+        model.custom_model_card_training_mode_info.return_value = (
+            "- MiniMax Music train component: `language_model (global LM / RVQ planner)`\n"
+            "- MiniMax Music LM max frames: `128`"
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with (
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_model_family", return_value="minimaxmusic"),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_data_backends", return_value={}),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_weight_dtype", return_value=torch.bfloat16),
+                patch(
+                    "simpletuner.helpers.publishing.metadata.StateTracker.get_accelerator",
+                    return_value=MagicMock(num_processes=1),
+                ),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_args", return_value=self.args),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_model", return_value=model),
+            ):
+                save_model_card(
+                    repo_id="test-repo",
+                    images=None,
+                    audios=None,
+                    base_model="MiniMaxAI/MiniMax-Music3",
+                    train_text_encoder=False,
+                    prompt="",
+                    validation_prompts=None,
+                    validation_shortnames=None,
+                    repo_folder=tmpdir,
+                    model=model,
+                    global_step=6000,
+                    epoch=250,
+                )
+
+            readme = Path(tmpdir, "README.md").read_text(encoding="utf-8")
+            self.assertIn("Validation was disabled during training.", readme)
+            self.assertNotIn("## Validation settings", readme)
+            self.assertNotIn("<Gallery />", readme)
+            self.assertIn("## Training modes", readme)
+            self.assertIn("- MiniMax Music train component: `language_model (global LM / RVQ planner)`", readme)
+            self.assertIn("- MiniMax Music LM max frames: `128`", readme)
+            self.assertIn("- NextLat: Enabled", readme)
+            self.assertIn("  - Weight: `0.1`", readme)
+            self.assertIn("- XM: Enabled", readme)
+            self.assertIn("  - Candidate count: `2`", readme)
+            self.assertIn("  - Training target: `route`", readme)
+            model.custom_model_card_training_mode_info.assert_called_once_with(self.args)
+
+    def test_minimax_music_model_card_training_mode_info(self):
+        from simpletuner.helpers.models.minimaxmusic.model import MiniMaxMusic
+
+        self.args.minimax_music_train_component = "language_model"
+        self.args.minimax_music_lm_max_frames = 128
+        model = MiniMaxMusic.__new__(MiniMaxMusic)
+
+        details = model.custom_model_card_training_mode_info(self.args)
+
+        self.assertIn("- MiniMax Music train component: `language_model (global LM / RVQ planner)`", details)
+        self.assertIn("- MiniMax Music LM max frames: `128`", details)
+
     def test_hub_commit_message_omits_diffusion_schedule_fields_for_flow_matching(self):
         hub_manager = object.__new__(HubManager)
         hub_manager.collected_data_backend_str = "['alt-embed-cache', 'h3-drift0-anyflow-openvid-39f-480']"