refactor model-specific registry fields into the models subdir instead

simpletuner/helpers/models/field_registry/__init__.py

@@ -0,0 +1,20 @@
+"""Model-specific field registry discovery."""
+
+import importlib
+import pkgutil
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from simpletuner.simpletuner_sdk.server.services.field_registry.registry import FieldRegistry
+
+
+def register_model_field_registries(registry: "FieldRegistry") -> None:
+    package = importlib.import_module(__name__)
+
+    for module_info in sorted(pkgutil.iter_modules(package.__path__), key=lambda item: item.name):
+        if module_info.ispkg or module_info.name.startswith("_"):
+            continue
+
+        module = importlib.import_module(f"{__name__}.{module_info.name}")
+        register_fields = getattr(module, "register_fields")
+        register_fields(registry)

simpletuner/helpers/models/field_registry/ace_step.py

@@ -0,0 +1,171 @@
+from simpletuner.helpers.training.optimizer_param import available_optimizer_keys as _available_optimizer_keys
+from simpletuner.simpletuner_sdk.server.services.field_registry.types import (
+    ConfigField,
+    FieldDependency,
+    FieldType,
+    ImportanceLevel,
+    ValidationRule,
+    ValidationRuleType,
+)
+
+
+def register_fields(registry) -> None:
+    registry._add_field(
+        ConfigField(
+            name="validation_lyrics",
+            arg_name="--validation_lyrics",
+            ui_label="Validation Lyrics",
+            field_type=FieldType.TEXTAREA,
+            tab="validation",
+            section="prompt_management",
+            placeholder="Enter lyrics for audio validation",
+            help_text="Lyrics to use for audio validation",
+            tooltip="Provide lyrics for music generation validation. Only used by audio models.",
+            importance=ImportanceLevel.ADVANCED,
+            order=2,
+            allow_empty=True,
+            model_specific=["ace_step"],
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="validation_audio_duration",
+            arg_name="--validation_audio_duration",
+            ui_label="Validation Audio Duration",
+            field_type=FieldType.NUMBER,
+            tab="validation",
+            section="validation_schedule",
+            default_value=30.0,
+            validation_rules=[
+                ValidationRule(ValidationRuleType.MIN, value=1.0, message="Duration must be at least 1 second"),
+                ValidationRule(ValidationRuleType.MAX, value=300.0, message="Duration recommended to be under 300s"),
+            ],
+            help_text="Duration of generated audio for validation (seconds)",
+            tooltip="Length of the audio clip to generate during validation runs.",
+            importance=ImportanceLevel.ADVANCED,
+            order=6,
+            model_specific=["ace_step"],
+        )
+    )
+
+    optimizer_choices = _available_optimizer_keys()
+    if not optimizer_choices:
+        raise RuntimeError("No optimizers available for the current environment.")
+    lr_scheduler_choices = [
+        "linear",
+        "sine",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "constant",
+        "constant_with_warmup",
+    ]
+
+    registry._add_field(
+        ConfigField(
+            name="lyrics_embedder_train",
+            arg_name="--lyrics_embedder_train",
+            ui_label="Train Lyrics Embedder",
+            field_type=FieldType.CHECKBOX,
+            tab="training",
+            section="lyrics_embedder",
+            default_value=False,
+            help_text="Enable fine-tuning of the ACE-Step lyrics embedder components.",
+            tooltip="Unlock lyric embedding layers for training. Recommended for ACE-Step only.",
+            importance=ImportanceLevel.ADVANCED,
+            model_specific=["ace_step"],
+            order=1,
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="lyrics_embedder_optimizer",
+            arg_name="--lyrics_embedder_optimizer",
+            ui_label="Lyrics Embedder Optimizer",
+            field_type=FieldType.SELECT,
+            tab="training",
+            section="lyrics_embedder",
+            default_value=None,
+            choices=[{"value": opt, "label": opt} for opt in optimizer_choices],
+            dynamic_choices=True,
+            validation_rules=[ValidationRule(ValidationRuleType.CHOICES, value=optimizer_choices)],
+            dependencies=[FieldDependency(field="lyrics_embedder_train", operator="equals", value=True, action="show")],
+            help_text="Optional optimizer override for the lyrics embedder (leave empty to reuse the main optimizer).",
+            tooltip="Pick a different optimizer just for the lyrics embedder, or leave blank to share the primary one.",
+            importance=ImportanceLevel.EXPERIMENTAL,
+            model_specific=["ace_step"],
+            allow_empty=True,
+            order=2,
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="lyrics_embedder_lr",
+            arg_name="--lyrics_embedder_lr",
+            ui_label="Lyrics Embedder Learning Rate",
+            field_type=FieldType.NUMBER,
+            tab="training",
+            section="lyrics_embedder",
+            default_value=None,
+            validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")],
+            dependencies=[FieldDependency(field="lyrics_embedder_train", operator="equals", value=True, action="show")],
+            help_text="Optional learning rate override for the lyrics embedder.",
+            tooltip="Leave empty to share the main learning rate. Set a value to use a dedicated rate.",
+            importance=ImportanceLevel.ADVANCED,
+            model_specific=["ace_step"],
+            allow_empty=True,
+            order=3,
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="lyrics_embedder_lr_scheduler",
+            arg_name="--lyrics_embedder_lr_scheduler",
+            ui_label="Lyrics Embedder LR Scheduler",
+            field_type=FieldType.SELECT,
+            tab="training",
+            section="lyrics_embedder",
+            default_value=None,
+            choices=[{"value": s, "label": s.replace("_", " ").title()} for s in lr_scheduler_choices],
+            validation_rules=[ValidationRule(ValidationRuleType.CHOICES, value=lr_scheduler_choices)],
+            dependencies=[FieldDependency(field="lyrics_embedder_train", operator="equals", value=True, action="show")],
+            help_text="Select a scheduler for the lyrics embedder (leave empty to mirror the main scheduler).",
+            tooltip="Use a distinct scheduler for lyric embeddings if needed, or leave blank to follow the primary plan.",
+            importance=ImportanceLevel.EXPERIMENTAL,
+            model_specific=["ace_step"],
+            allow_empty=True,
+            order=4,
+        )
+    )
+
+    acestep_targets = [
+        "attn_qkv",
+        "attn_qkv+linear_qkv",
+        "attn_qkv+linear_qkv+speech_embedder",
+    ]
+    registry._add_field(
+        ConfigField(
+            name="acestep_lora_target",
+            arg_name="--acestep_lora_target",
+            ui_label="ACE-Step LoRA Target Layers",
+            field_type=FieldType.SELECT,
+            tab="model",
+            section="lora_config",
+            subsection="model_specific",
+            default_value="attn_qkv+linear_qkv",
+            choices=[{"value": t, "label": t} for t in acestep_targets],
+            dependencies=[
+                FieldDependency(field="model_type", value="lora"),
+                FieldDependency(field="model_family", value="ace_step"),
+            ],
+            help_text="Which layers to train in ACE-Step models",
+            tooltip="'attn_qkv+linear_qkv' is default. '+speech_embedder' adds speaker embedding. 'attn_qkv' is minimal.",
+            importance=ImportanceLevel.ADVANCED,
+            model_specific=["ace_step"],
+            order=11,
+        )
+    )

simpletuner/helpers/models/field_registry/deepfloyd.py

@@ -0,0 +1,203 @@
+from simpletuner.simpletuner_sdk.server.services.field_registry.types import (
+    ConfigField,
+    FieldType,
+    ImportanceLevel,
+    ValidationRule,
+    ValidationRuleType,
+)
+
+
+def register_fields(registry) -> None:
+    registry._add_field(
+        ConfigField(
+            name="deepfloyd_validation_pipeline_mode",
+            arg_name="--deepfloyd_validation_pipeline_mode",
+            ui_label="DeepFloyd Validation Pipeline",
+            field_type=FieldType.SELECT,
+            tab="validation",
+            section="validation_options",
+            default_value="auto",
+            choices=[
+                {"value": "auto", "label": "Auto"},
+                {"value": "trained-stage", "label": "Trained Stage Only"},
+                {"value": "full-pipeline", "label": "Full Pipeline"},
+            ],
+            help_text="Choose whether DeepFloyd validation runs only the trained stage or chains fixed peer stages.",
+            tooltip="Auto uses the full DeepFloyd pipeline for prompt validation and the trained stage for dataset image validation.",
+            importance=ImportanceLevel.ADVANCED,
+            order=21,
+            subsection="advanced",
+            model_specific=["deepfloyd"],
+            documentation="OPTIONS.md#--deepfloyd_validation_pipeline_mode",
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="deepfloyd_validation_stage1_model",
+            arg_name="--deepfloyd_validation_stage1_model",
+            ui_label="DeepFloyd Stage I Model",
+            field_type=FieldType.TEXT,
+            tab="validation",
+            section="validation_options",
+            default_value=None,
+            placeholder="DeepFloyd/IF-I-XL-v1.0",
+            help_text="Fixed DeepFloyd stage I model used when validating a trained stage II model through the full pipeline.",
+            tooltip="Leave blank to use DeepFloyd/IF-I-XL-v1.0.",
+            importance=ImportanceLevel.ADVANCED,
+            order=22,
+            subsection="advanced",
+            model_specific=["deepfloyd"],
+            documentation="OPTIONS.md#--deepfloyd_validation_stage1_model",
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="deepfloyd_validation_stage2_model",
+            arg_name="--deepfloyd_validation_stage2_model",
+            ui_label="DeepFloyd Stage II Model",
+            field_type=FieldType.TEXT,
+            tab="validation",
+            section="validation_options",
+            default_value=None,
+            placeholder="DeepFloyd/IF-II-M-v1.0",
+            help_text="Fixed DeepFloyd stage II model used when validating a trained stage I model through the full pipeline.",
+            tooltip="Leave blank to use DeepFloyd/IF-II-M-v1.0.",
+            importance=ImportanceLevel.ADVANCED,
+            order=23,
+            subsection="advanced",
+            model_specific=["deepfloyd"],
+            documentation="OPTIONS.md#--deepfloyd_validation_stage2_model",
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="deepfloyd_validation_stage3_mode",
+            arg_name="--deepfloyd_validation_stage3_mode",
+            ui_label="DeepFloyd Stage III Mode",
+            field_type=FieldType.SELECT,
+            tab="validation",
+            section="validation_options",
+            default_value="none",
+            choices=[
+                {"value": "none", "label": "None"},
+                {"value": "sd-x4-upscaler", "label": "Stable Diffusion x4 Upscaler"},
+            ],
+            help_text="Optional terminal DeepFloyd validation upscaler after stage II.",
+            tooltip="Stage III is not a released DeepFloyd model; this option can use the era-compatible SD x4 upscaler.",
+            importance=ImportanceLevel.ADVANCED,
+            order=24,
+            subsection="advanced",
+            model_specific=["deepfloyd"],
+            documentation="OPTIONS.md#--deepfloyd_validation_stage3_mode",
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="deepfloyd_validation_stage3_model",
+            arg_name="--deepfloyd_validation_stage3_model",
+            ui_label="DeepFloyd Stage III Model",
+            field_type=FieldType.TEXT,
+            tab="validation",
+            section="validation_options",
+            default_value=None,
+            placeholder="stabilityai/stable-diffusion-x4-upscaler",
+            help_text="Model repository used when DeepFloyd stage III mode is the SD x4 upscaler.",
+            tooltip="Leave blank to use stabilityai/stable-diffusion-x4-upscaler.",
+            importance=ImportanceLevel.ADVANCED,
+            order=25,
+            subsection="advanced",
+            model_specific=["deepfloyd"],
+            documentation="OPTIONS.md#--deepfloyd_validation_stage3_model",
+        )
+    )
+
+    for field_name, label, arg_name, order in [
+        (
+            "deepfloyd_validation_stage1_num_inference_steps",
+            "DeepFloyd Stage I Steps",
+            "--deepfloyd_validation_stage1_num_inference_steps",
+            26,
+        ),
+        (
+            "deepfloyd_validation_stage2_num_inference_steps",
+            "DeepFloyd Stage II Steps",
+            "--deepfloyd_validation_stage2_num_inference_steps",
+            27,
+        ),
+    ]:
+        registry._add_field(
+            ConfigField(
+                name=field_name,
+                arg_name=arg_name,
+                ui_label=label,
+                field_type=FieldType.NUMBER,
+                tab="validation",
+                section="validation_options",
+                default_value=None,
+                validation_rules=[ValidationRule(ValidationRuleType.MIN, value=1, message="Must be at least 1")],
+                help_text="Override the DeepFloyd per-stage validation step count.",
+                tooltip="Leave blank to use the normal validation step count.",
+                importance=ImportanceLevel.ADVANCED,
+                order=order,
+                subsection="advanced",
+                model_specific=["deepfloyd"],
+            )
+        )
+
+    for field_name, label, arg_name, order in [
+        ("deepfloyd_validation_stage1_guidance", "DeepFloyd Stage I Guidance", "--deepfloyd_validation_stage1_guidance", 28),
+        (
+            "deepfloyd_validation_stage2_guidance",
+            "DeepFloyd Stage II Guidance",
+            "--deepfloyd_validation_stage2_guidance",
+            29,
+        ),
+        (
+            "deepfloyd_validation_stage3_guidance",
+            "DeepFloyd Stage III Guidance",
+            "--deepfloyd_validation_stage3_guidance",
+            30,
+        ),
+    ]:
+        registry._add_field(
+            ConfigField(
+                name=field_name,
+                arg_name=arg_name,
+                ui_label=label,
+                field_type=FieldType.NUMBER,
+                tab="validation",
+                section="validation_options",
+                default_value=None,
+                validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")],
+                help_text="Override the DeepFloyd per-stage validation guidance scale.",
+                tooltip="Leave blank to use the normal validation guidance value.",
+                importance=ImportanceLevel.ADVANCED,
+                order=order,
+                subsection="advanced",
+                model_specific=["deepfloyd"],
+            )
+        )
+
+    registry._add_field(
+        ConfigField(
+            name="deepfloyd_validation_stage3_noise_level",
+            arg_name="--deepfloyd_validation_stage3_noise_level",
+            ui_label="DeepFloyd Stage III Noise",
+            field_type=FieldType.NUMBER,
+            tab="validation",
+            section="validation_options",
+            default_value=100,
+            validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")],
+            help_text="Noise level passed to the SD x4 upscaler during DeepFloyd validation.",
+            tooltip="Only used when DeepFloyd stage III mode is the SD x4 upscaler.",
+            importance=ImportanceLevel.ADVANCED,
+            order=31,
+            subsection="advanced",
+            model_specific=["deepfloyd"],
+            documentation="OPTIONS.md#--deepfloyd_validation_stage3_noise_level",
+        )
+    )