diff --git a/simpletuner/helpers/models/field_registry/__init__.py b/simpletuner/helpers/models/field_registry/__init__.py new file mode 100644 index 000000000..ef4dadaf3 --- /dev/null +++ b/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) diff --git a/simpletuner/helpers/models/field_registry/ace_step.py b/simpletuner/helpers/models/field_registry/ace_step.py new file mode 100644 index 000000000..27a7e10e7 --- /dev/null +++ b/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, + ) + ) diff --git a/simpletuner/helpers/models/field_registry/deepfloyd.py b/simpletuner/helpers/models/field_registry/deepfloyd.py new file mode 100644 index 000000000..b51ea134e --- /dev/null +++ b/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", + ) + ) diff --git a/simpletuner/helpers/models/field_registry/flux.py b/simpletuner/helpers/models/field_registry/flux.py new file mode 100644 index 000000000..3ff8270ef --- /dev/null +++ b/simpletuner/helpers/models/field_registry/flux.py @@ -0,0 +1,229 @@ +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_guidance_real", + arg_name="--validation_guidance_real", + ui_label="Real CFG (Distilled Models)", + field_type=FieldType.NUMBER, + tab="validation", + section="validation_guidance", + default_value=1.0, + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=1.0, message="Must be at least 1.0")], + help_text="CFG value for distilled models (e.g., FLUX schnell)", + tooltip="Use 1.0 for no CFG (distilled models). Higher values for real CFG sampling.", + importance=ImportanceLevel.ADVANCED, + order=2, + model_specific=["flux", "flux2"], + ) + ) + + registry._add_field( + ConfigField( + name="validation_no_cfg_until_timestep", + arg_name="--validation_no_cfg_until_timestep", + ui_label="Skip CFG Until Timestep", + field_type=FieldType.NUMBER, + tab="validation", + section="validation_guidance", + default_value=2, + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], + help_text="Skip CFG for initial timesteps (Flux only)", + tooltip="For Flux real CFG: skip CFG on these initial timesteps. Default: 2", + importance=ImportanceLevel.ADVANCED, + order=3, + model_specific=["flux", "flux2"], + ) + ) + + flux_targets = [ + "mmdit", + "context", + "context+ffs", + "all", + "all+ffs", + "ai-toolkit", + "tiny", + "nano", + "controlnet", + "all+ffs+embedder", + "all+ffs+embedder+controlnet", + ] + registry._add_field( + ConfigField( + name="flux_lora_target", + arg_name="--flux_lora_target", + ui_label="Flux LoRA Target Layers", + field_type=FieldType.SELECT, + tab="model", + section="lora_config", + subsection="model_specific", + default_value="all", + choices=[{"value": t, "label": t} for t in flux_targets], + dependencies=[ + FieldDependency(field="model_type", value="lora"), + FieldDependency(field="model_family", value="flux"), + ], + help_text="Which layers to train in Flux models", + tooltip="'all' trains all attention layers. 'context' only trains text layers. '+ffs' includes feed-forward layers.", + importance=ImportanceLevel.ADVANCED, + model_specific=["flux"], + order=10, + ) + ) + + registry._add_field( + ConfigField( + name="flux_guidance_mode", + arg_name="--flux_guidance_mode", + ui_label="Flux Guidance Mode", + field_type=FieldType.SELECT, + tab="model", + section="model_specific", + default_value="constant", + choices=[{"value": "constant", "label": "Constant"}, {"value": "random-range", "label": "Random Range"}], + help_text="Guidance mode for Flux training", + tooltip="Constant uses same guidance for all samples. Random Range varies guidance per sample.", + importance=ImportanceLevel.ADVANCED, + order=40, + dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], + ) + ) + + registry._add_field( + ConfigField( + name="flux_attention_masked_training", + arg_name="--flux_attention_masked_training", + ui_label="Attention Masked Training", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + default_value=False, + model_specific=["flux"], + help_text="Enable attention masked training for Flux models", + tooltip="Experimental feature for Flux models that masks certain attention patterns during training.", + importance=ImportanceLevel.EXPERIMENTAL, + order=10, + ) + ) + + registry._add_field( + ConfigField( + name="flux_fast_schedule", + arg_name="--flux_fast_schedule", + ui_label="Fast Training Schedule", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + default_value=False, + model_specific=["flux"], + help_text="Use experimental fast schedule for Flux training", + tooltip="Experimental feature that may speed up Flux.1S training at the cost of quality.", + importance=ImportanceLevel.EXPERIMENTAL, + order=11, + ) + ) + + registry._add_field( + ConfigField( + name="flux_guidance_value", + arg_name="--flux_guidance_value", + ui_label="Flux Guidance Value", + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + default_value=1.0, + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], + help_text="Guidance value for constant mode", + tooltip="1.0 preserves CFG distillation. Higher values require CFG at inference.", + importance=ImportanceLevel.ADVANCED, + order=41, + dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], + ) + ) + + registry._add_field( + ConfigField( + name="flux_guidance_min", + arg_name="--flux_guidance_min", + ui_label="Flux Guidance Min", + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + default_value=0.0, + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], + help_text="Minimum guidance value for random-range mode", + tooltip="Lower bound of guidance range when using random-range mode.", + importance=ImportanceLevel.ADVANCED, + order=42, + dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], + ) + ) + + registry._add_field( + ConfigField( + name="flux_guidance_max", + arg_name="--flux_guidance_max", + ui_label="Flux Guidance Max", + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + default_value=4.0, + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], + help_text="Maximum guidance value for random-range mode", + tooltip="Upper bound of guidance range when using random-range mode.", + importance=ImportanceLevel.ADVANCED, + order=43, + dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], + ) + ) + + registry._add_field( + ConfigField( + name="fused_qkv_projections", + arg_name="--fuse_qkv_projections", + ui_label="Fused QKV Projections", + field_type=FieldType.CHECKBOX, + tab="model", + section="architecture", + default_value=False, + dependencies=[FieldDependency(field="model_family", operator="equals", value="flux")], + help_text="Enables Flash Attention 3 when supported; otherwise falls back to PyTorch SDPA.", + tooltip="Improves attention efficiency on modern NVIDIA GPUs. Uses native SDPA when Flash Attention 3 is unavailable.", + importance=ImportanceLevel.EXPERIMENTAL, + model_specific=["flux"], + order=19, + aliases=["--fused_qkv_projections"], + documentation="OPTIONS.md#--fuse_qkv_projections", + ) + ) + + registry._add_field( + ConfigField( + name="custom_text_encoder_intermediary_layers", + arg_name="--custom_text_encoder_intermediary_layers", + ui_label="Custom Text Encoder Layers", + field_type=FieldType.TEXT, + tab="model", + section="architecture", + subsection="advanced", + default_value=None, + placeholder="[10, 20, 30]", + dependencies=[FieldDependency(field="model_family", operator="in", values=["flux2"])], + help_text="Override which hidden state layers to extract from the text encoder. Provide as JSON array (e.g., [10, 20, 30]). Leave blank to use model defaults.", + tooltip="FLUX.2-dev uses layers [10, 20, 30] from Mistral-3, Klein models use [9, 18, 27] from Qwen3. Override for experimentation.", + importance=ImportanceLevel.EXPERIMENTAL, + model_specific=["flux2"], + order=34, + documentation="OPTIONS.md#--custom_text_encoder_intermediary_layers", + ) + ) diff --git a/simpletuner/helpers/models/field_registry/hidream.py b/simpletuner/helpers/models/field_registry/hidream.py new file mode 100644 index 000000000..bd042ee57 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/hidream.py @@ -0,0 +1,44 @@ +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="hidream_use_load_balancing_loss", + arg_name="--hidream_use_load_balancing_loss", + ui_label="Enable HiDream Load Balancing Loss", + field_type=FieldType.CHECKBOX, + tab="training", + section="loss_functions", + default_value=False, + dependencies=[FieldDependency(field="model_family", operator="equals", value="hidream")], + help_text="Apply experimental load balancing loss when training HiDream models.", + tooltip="Balances expert contributions during HiDream training. Only available for HiDream model family.", + importance=ImportanceLevel.EXPERIMENTAL, + order=6, + ) + ) + + registry._add_field( + ConfigField( + name="hidream_load_balancing_loss_weight", + arg_name="--hidream_load_balancing_loss_weight", + ui_label="HiDream Load Balancing Weight", + field_type=FieldType.NUMBER, + tab="training", + section="loss_functions", + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be non-negative")], + dependencies=[FieldDependency(field="hidream_use_load_balancing_loss", operator="equals", value=True)], + help_text="Strength multiplier for HiDream load balancing loss.", + tooltip="Adjust if you need stronger balancing between experts. Leave blank to use the trainer default.", + importance=ImportanceLevel.EXPERIMENTAL, + order=7, + ) + ) diff --git a/simpletuner/helpers/models/field_registry/ideogram.py b/simpletuner/helpers/models/field_registry/ideogram.py new file mode 100644 index 000000000..ec8657111 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/ideogram.py @@ -0,0 +1,116 @@ +from simpletuner.simpletuner_sdk.server.services.field_registry.types import ( + ConfigField, + FieldDependency, + FieldType, + ImportanceLevel, +) + + +def register_fields(registry) -> None: + registry._add_field( + ConfigField( + name="ideogram_auto_json", + arg_name="--ideogram_auto_json", + ui_label="Ideogram Auto JSON", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + model_specific=["ideogram"], + default_value=True, + help_text="Convert non-JSON Ideogram 4 prompts into the structured JSON caption schema.", + tooltip="When enabled, plain validation prompts are wrapped into Ideogram 4's JSON caption format.", + importance=ImportanceLevel.ADVANCED, + dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], + order=33, + ) + ) + + registry._add_field( + ConfigField( + name="ideogram_validation", + arg_name="--ideogram_validation", + ui_label="Ideogram Validation", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + model_specific=["ideogram"], + default_value=False, + help_text="Temporarily enable Ideogram validation by reusing the conditional transformer for CFG's unconditional pass.", + tooltip="Validation is off by default for Ideogram until separate unconditional-model handling is implemented.", + importance=ImportanceLevel.ADVANCED, + dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], + order=34, + ) + ) + + registry._add_field( + ConfigField( + name="ideogram_prompt_upsample", + arg_name="--ideogram_prompt_upsample", + ui_label="Ideogram Prompt Upsample", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + model_specific=["ideogram"], + default_value=False, + help_text="Use Ideogram 4's prompt enhancer to rewrite prompts before text embedding cache generation.", + tooltip="When enabled and supported by the loaded pipeline, captions are expanded with Ideogram's prompt upsampler before JSON conversion and encoding.", + importance=ImportanceLevel.ADVANCED, + dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], + order=35, + ) + ) + + registry._add_field( + ConfigField( + name="ideogram_prompt_enhancer_head_id", + arg_name="--ideogram_prompt_enhancer_head_id", + ui_label="Ideogram Prompt Enhancer Head", + field_type=FieldType.TEXT, + tab="model", + section="model_specific", + model_specific=["ideogram"], + default_value="diffusers/qwen3-vl-8b-instruct-lm-head", + help_text="Hugging Face repo id for Ideogram 4's prompt upsampling LM head.", + tooltip="Used when --ideogram_prompt_upsample is enabled to rewrite prompts into Ideogram's structured JSON caption schema.", + importance=ImportanceLevel.ADVANCED, + dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], + order=36, + ) + ) + + registry._add_field( + ConfigField( + name="ideogram_schedule_mu", + arg_name="--ideogram_schedule_mu", + ui_label="Ideogram Schedule Mu", + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + model_specific=["ideogram"], + default_value=0.0, + help_text="Base mean used by Ideogram 4's resolution-aware logit-normal training timestep schedule.", + tooltip="Matches the default mu used by the vendored Ideogram validation pipeline.", + importance=ImportanceLevel.ADVANCED, + dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], + order=37, + ) + ) + + registry._add_field( + ConfigField( + name="ideogram_schedule_std", + arg_name="--ideogram_schedule_std", + ui_label="Ideogram Schedule Std", + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + model_specific=["ideogram"], + default_value=1.5, + help_text="Standard deviation used by Ideogram 4's logit-normal training timestep schedule.", + tooltip="Matches the default std used by the vendored Ideogram validation pipeline.", + importance=ImportanceLevel.ADVANCED, + dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], + order=38, + ) + ) diff --git a/simpletuner/helpers/models/field_registry/krea2.py b/simpletuner/helpers/models/field_registry/krea2.py new file mode 100644 index 000000000..fd81b7b54 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/krea2.py @@ -0,0 +1,21 @@ +from simpletuner.simpletuner_sdk.server.services.field_registry.types import ConfigField, FieldType, ImportanceLevel + + +def register_fields(registry) -> None: + registry._add_field( + ConfigField( + name="krea2_reference_latents", + arg_name="--krea2_reference_latents", + ui_label="Krea 2 Reference Latents", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + model_specific=["krea2"], + default_value=False, + help_text="Enable Krea 2 reference-dataset training with image-context prompt embeds and clean reference latents.", + tooltip="When enabled, Krea 2 requires paired conditioning data, encodes prompts with the reference image through Qwen3VL, and appends the clean reference latents to the transformer input.", + importance=ImportanceLevel.ADVANCED, + order=36, + documentation="OPTIONS.md#--krea2_reference_latents", + ) + ) diff --git a/simpletuner/helpers/models/field_registry/ltx.py b/simpletuner/helpers/models/field_registry/ltx.py new file mode 100644 index 000000000..2022619a8 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/ltx.py @@ -0,0 +1,89 @@ +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="ltx_train_mode", + arg_name="--ltx_train_mode", + ui_label="LTX Train Mode", + field_type=FieldType.SELECT, + tab="model", + section="model_specific", + model_specific=["ltx"], + default_value="i2v", + choices=[ + {"value": "t2v", "label": "Text-to-Video"}, + {"value": "i2v", "label": "Image-to-Video"}, + ], + help_text="Training mode for LTX models", + tooltip="Choose whether datasets default to text-to-video (t2v) or image-to-video (i2v) processing.", + importance=ImportanceLevel.ADVANCED, + order=22, + ) + ) + + registry._add_field( + ConfigField( + name="ltx_i2v_prob", + arg_name="--ltx_i2v_prob", + ui_label="LTX Image-to-Video Probability", + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + model_specific=["ltx"], + default_value=0.1, + validation_rules=[ + ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be between 0 and 1"), + ValidationRule(ValidationRuleType.MAX, value=1.0, message="Must be between 0 and 1"), + ], + help_text="Probability of using image-to-video training for LTX", + tooltip="Fraction of training that uses image-to-video instead of video-to-video.", + importance=ImportanceLevel.ADVANCED, + order=23, + ) + ) + + registry._add_field( + ConfigField( + name="ltx_partial_noise_fraction", + arg_name="--ltx_partial_noise_fraction", + ui_label="LTX Partial Noise Fraction", + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + model_specific=["ltx"], + default_value=0.05, + validation_rules=[ + ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be between 0 and 1"), + ValidationRule(ValidationRuleType.MAX, value=1.0, message="Must be between 0 and 1"), + ], + help_text="Fraction of noise to add for LTX partial training", + tooltip="Controls how much noise to add during partial training steps.", + importance=ImportanceLevel.ADVANCED, + order=24, + ) + ) + + registry._add_field( + ConfigField( + name="ltx_protect_first_frame", + arg_name="--ltx_protect_first_frame", + ui_label="LTX Protect First Frame", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + model_specific=["ltx"], + default_value=False, + help_text="Protect the first frame from noise in LTX training", + tooltip="When enabled, first frame is kept clean while subsequent frames get noise.", + importance=ImportanceLevel.ADVANCED, + order=25, + ) + ) diff --git a/simpletuner/helpers/models/field_registry/ltxvideo2.py b/simpletuner/helpers/models/field_registry/ltxvideo2.py new file mode 100644 index 000000000..f87c3df93 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/ltxvideo2.py @@ -0,0 +1,135 @@ +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="ltx2_intrinsic_conditioning", + arg_name="--ltx2_intrinsic_conditioning", + ui_label="LTX-2 Intrinsic Conditioning", + field_type=FieldType.TEXT_JSON, + tab="model", + section="model_specific", + model_specific=["ltxvideo2"], + default_value=None, + help_text="JSON array of LTX-2 intrinsic clean-token conditioning objects", + tooltip="Advanced LTX-2 training config. Supported condition types: first_frame, prefix, suffix, spatial_crop, mask.", + importance=ImportanceLevel.ADVANCED, + order=26, + documentation="OPTIONS.md#ltx-2-intrinsic-and-reference-conditioning", + ) + ) + + for offset, field_name, label, help_text in ( + ( + 0, + "ltx2_first_frame_conditioning_probability", + "LTX-2 First Frame Conditioning Probability", + "Probability of replacing first-frame target tokens with clean latents and removing them from video loss", + ), + ( + 1, + "ltx2_prefix_conditioning_probability", + "LTX-2 Prefix Conditioning Probability", + "Probability of replacing prefix target tokens with clean latents and removing them from video loss", + ), + ( + 2, + "ltx2_suffix_conditioning_probability", + "LTX-2 Suffix Conditioning Probability", + "Probability of replacing suffix target tokens with clean latents and removing them from video loss", + ), + ( + 3, + "ltx2_mask_conditioning_probability", + "LTX-2 Mask Conditioning Probability", + "Probability of using mask=1 target tokens as clean conditioning with no video loss", + ), + ): + registry._add_field( + ConfigField( + name=field_name, + arg_name=f"--{field_name}", + ui_label=label, + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + model_specific=["ltxvideo2"], + default_value=0.0, + validation_rules=[ + ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be between 0 and 1"), + ValidationRule(ValidationRuleType.MAX, value=1.0, message="Must be between 0 and 1"), + ], + help_text=help_text, + tooltip=help_text, + importance=ImportanceLevel.ADVANCED, + order=27 + offset, + documentation="OPTIONS.md#ltx-2-intrinsic-and-reference-conditioning", + ) + ) + + for offset, field_name, label, default_value in ( + (0, "ltx2_prefix_conditioning_frames", "LTX-2 Prefix Conditioning Frames", 1), + (1, "ltx2_suffix_conditioning_frames", "LTX-2 Suffix Conditioning Frames", 1), + (2, "ltx2_reference_spatial_scale_factor", "LTX-2 Reference Spatial Scale Factor", None), + (3, "ltx2_reference_temporal_scale_factor", "LTX-2 Reference Temporal Scale Factor", 1), + ): + registry._add_field( + ConfigField( + name=field_name, + arg_name=f"--{field_name}", + ui_label=label, + field_type=FieldType.NUMBER, + tab="model", + section="model_specific", + model_specific=["ltxvideo2"], + default_value=default_value, + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=1, message="Must be at least 1")], + help_text=f"Advanced LTX-2 conditioning setting: {label.lower()}", + tooltip=f"Advanced LTX-2 conditioning setting: {label.lower()}", + importance=ImportanceLevel.ADVANCED, + order=31 + offset, + documentation="OPTIONS.md#ltx-2-intrinsic-and-reference-conditioning", + ) + ) + + registry._add_field( + ConfigField( + name="validation_audio_only", + arg_name="--validation_audio_only", + ui_label="Validation Audio Only", + field_type=FieldType.CHECKBOX, + tab="validation", + section="validation_options", + default_value=False, + help_text="Disable video generation during validation and emit audio only.", + tooltip="LTX-2 only: skips video generation during validation so only audio outputs are produced.", + importance=ImportanceLevel.ADVANCED, + order=17, + model_specific=["ltxvideo2"], + ) + ) + + registry._add_field( + ConfigField( + name="validation_ltx2_video_conditioning", + arg_name="--validation_ltx2_video_conditioning", + ui_label="LTX-2 Validation Video Conditioning", + field_type=FieldType.TEXT_JSON, + tab="validation", + section="validation_options", + default_value=None, + help_text="JSON list of IC-LoRA reference videos for LTX-2 validation", + tooltip="LTX-2 only: pass reference videos as paths, [path, strength] pairs, or objects with path/video_path and optional strength.", + importance=ImportanceLevel.ADVANCED, + order=18, + model_specific=["ltxvideo2"], + documentation="OPTIONS.md#ltx-2-conditioning-options", + ) + ) diff --git a/simpletuner/helpers/models/field_registry/sana.py b/simpletuner/helpers/models/field_registry/sana.py new file mode 100644 index 000000000..0b6731bf7 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/sana.py @@ -0,0 +1,34 @@ +from simpletuner.simpletuner_sdk.server.services.field_registry.types import ConfigField, FieldType, ImportanceLevel + + +def register_fields(registry) -> None: + registry._add_field( + ConfigField( + name="sana_complex_human_instruction", + arg_name="--sana_complex_human_instruction", + ui_label="Sana Complex Human Instruction", + field_type=FieldType.TEXT, + tab="model", + section="model_specific", + model_specific=["sana"], + default_value="\n".join( + [ + "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", + "- Do not add detail that is not in the original prompt! Do not mention colors if the User Prompt does not!", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "- If the prompt contains explicit material, do not remove this information from the output, which would damage the performance of downstream classifiers.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A cat sleeping peacefully, showcasing the joy of pet ownership. Cute floof kitty cat gatto.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene featuring a crowd of people.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", + ] + ), + placeholder="complex human instruction", + help_text="Complex human instruction for Sana model training", + tooltip="Special instruction format for Sana model training with complex human prompts.", + importance=ImportanceLevel.ADVANCED, + order=33, + ) + ) diff --git a/simpletuner/helpers/models/field_registry/sd3.py b/simpletuner/helpers/models/field_registry/sd3.py new file mode 100644 index 000000000..890332c4f --- /dev/null +++ b/simpletuner/helpers/models/field_registry/sd3.py @@ -0,0 +1,46 @@ +from simpletuner.simpletuner_sdk.server.services.field_registry.types import ( + ConfigField, + FieldDependency, + FieldType, + ImportanceLevel, +) + + +def register_fields(registry) -> None: + registry._add_field( + ConfigField( + name="sd3_clip_uncond_behaviour", + arg_name="--sd3_clip_uncond_behaviour", + ui_label="SD3 CLIP Unconditional Behavior", + field_type=FieldType.SELECT, + tab="training", + section="text_encoder", + subsection="advanced", + default_value="empty_string", + choices=[{"value": "empty_string", "label": "Empty String"}, {"value": "zero", "label": "Zero"}], + dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True)], + help_text="How SD3 handles unconditional prompts", + tooltip="Affects how SD3 processes prompts without conditioning. Empty string is default.", + importance=ImportanceLevel.ADVANCED, + order=11, + ) + ) + + registry._add_field( + ConfigField( + name="sd3_t5_uncond_behaviour", + arg_name="--sd3_t5_uncond_behaviour", + ui_label="SD3 T5 Unconditional Behavior", + field_type=FieldType.SELECT, + tab="training", + section="text_encoder", + subsection="advanced", + default_value=None, + choices=[{"value": "empty_string", "label": "Empty String"}, {"value": "zero", "label": "Zero"}], + dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True)], + help_text="How SD3 T5 handles unconditional prompts", + tooltip="Overrides CLIP behavior for T5. If not set, follows CLIP setting.", + importance=ImportanceLevel.ADVANCED, + order=12, + ) + ) diff --git a/simpletuner/helpers/models/field_registry/sdxl.py b/simpletuner/helpers/models/field_registry/sdxl.py new file mode 100644 index 000000000..ec0ed02c2 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/sdxl.py @@ -0,0 +1,88 @@ +from simpletuner.simpletuner_sdk.server.services.field_registry.types import ( + ConfigField, + FieldDependency, + FieldType, + ImportanceLevel, +) + + +def register_fields(registry) -> None: + registry._add_field( + ConfigField( + name="sdxl_refiner_uses_full_range", + arg_name="--sdxl_refiner_uses_full_range", + ui_label="SDXL Refiner Full Range", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + default_value=False, + help_text="Use full timestep range for SDXL refiner", + tooltip="When enabled, refiner uses full timestep range instead of just high timesteps.", + importance=ImportanceLevel.ADVANCED, + dependencies=[FieldDependency(field="model_family", operator="equals", value="sdxl")], + order=32, + ) + ) + + registry._add_field( + ConfigField( + name="sdxl_validation_pipeline_mode", + arg_name="--sdxl_validation_pipeline_mode", + ui_label="SDXL Validation Pipeline", + field_type=FieldType.SELECT, + tab="validation", + section="validation_options", + default_value="trained-stage", + choices=[ + {"value": "trained-stage", "label": "Trained Stage Only"}, + {"value": "full-pipeline", "label": "Full Pipeline"}, + ], + help_text="Choose whether SDXL validation runs only the trained stage or chains the base/refiner split pipeline.", + tooltip="Full pipeline runs stage 1 to the refiner boundary as latents, then resumes through stage 2.", + importance=ImportanceLevel.ADVANCED, + order=33, + subsection="advanced", + model_specific=["sdxl"], + documentation="OPTIONS.md#--sdxl_validation_pipeline_mode", + ) + ) + + registry._add_field( + ConfigField( + name="sdxl_validation_stage1_model", + arg_name="--sdxl_validation_stage1_model", + ui_label="SDXL Stage 1 Model", + field_type=FieldType.TEXT, + tab="validation", + section="validation_options", + default_value=None, + placeholder="stabilityai/stable-diffusion-xl-base-1.0", + help_text="Fixed SDXL stage 1 model used when validating a trained refiner through the full pipeline.", + tooltip="Leave blank to infer the matching SDXL base model from the selected flavour.", + importance=ImportanceLevel.ADVANCED, + order=34, + subsection="advanced", + model_specific=["sdxl"], + documentation="OPTIONS.md#--sdxl_validation_stage1_model", + ) + ) + + registry._add_field( + ConfigField( + name="sdxl_validation_stage2_model", + arg_name="--sdxl_validation_stage2_model", + ui_label="SDXL Stage 2 Model", + field_type=FieldType.TEXT, + tab="validation", + section="validation_options", + default_value=None, + placeholder="stabilityai/stable-diffusion-xl-refiner-1.0", + help_text="Fixed SDXL refiner model used when validating a trained stage 1 model through the full pipeline.", + tooltip="Leave blank to infer the matching SDXL refiner from the selected flavour.", + importance=ImportanceLevel.ADVANCED, + order=35, + subsection="advanced", + model_specific=["sdxl"], + documentation="OPTIONS.md#--sdxl_validation_stage2_model", + ) + ) diff --git a/simpletuner/helpers/models/field_registry/wan.py b/simpletuner/helpers/models/field_registry/wan.py new file mode 100644 index 000000000..9905fd530 --- /dev/null +++ b/simpletuner/helpers/models/field_registry/wan.py @@ -0,0 +1,85 @@ +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="enable_chunked_feed_forward", + arg_name="--enable_chunked_feed_forward", + ui_label="Enable Feed-Forward Chunking", + field_type=FieldType.CHECKBOX, + tab="model", + section="memory_optimization", + default_value=False, + help_text="Split Wan feed-forward layers into smaller chunks to reduce peak VRAM usage.", + tooltip="Available for Wan models. Breaks long MLPs into mini-batches so checkpoint recomputes allocate less memory.", + importance=ImportanceLevel.ADVANCED, + model_specific=["wan", "wan_s2v"], + order=8, + ) + ) + + registry._add_field( + ConfigField( + name="feed_forward_chunk_size", + arg_name="--feed_forward_chunk_size", + ui_label="Feed-Forward Chunk Size", + field_type=FieldType.NUMBER, + tab="model", + section="memory_optimization", + default_value=None, + validation_rules=[ValidationRule(ValidationRuleType.MIN, value=1, message="Chunk size must be at least 1")], + help_text="Number of samples processed per chunk when feed-forward chunking is enabled.", + tooltip="Leave blank for auto. Lower values reduce memory further but increase wall-clock time.", + importance=ImportanceLevel.ADVANCED, + model_specific=["wan", "wan_s2v"], + order=9, + dependencies=[ + FieldDependency(field="enable_chunked_feed_forward", operator="equals", value=True, action="show") + ], + allow_empty=True, + ) + ) + + registry._add_field( + ConfigField( + name="wan_force_2_1_time_embedding", + arg_name="--wan_force_2_1_time_embedding", + ui_label="Force Wan 2.1 Time Embedding", + field_type=FieldType.CHECKBOX, + tab="model", + section="model_specific", + default_value=False, + dependencies=[FieldDependency(field="model_family", operator="equals", value="wan", action="show")], + help_text="Use Wan 2.1 style time embeddings even when running Wan 2.2 checkpoints.", + tooltip="Enable this if Wan 2.2 checkpoints report shape mismatches in the time embedding layers.", + importance=ImportanceLevel.ADVANCED, + order=30, + ) + ) + + registry._add_field( + ConfigField( + name="wan_validation_load_other_stage", + arg_name="--wan_validation_load_other_stage", + ui_label="Wan Paired-Stage Validation", + field_type=FieldType.CHECKBOX, + tab="validation", + section="validation_options", + default_value=False, + help_text="Load the opposite Wan 2.2 stage during validation so the pipeline can switch denoisers at the stage boundary.", + tooltip="For Wan 2.2 and compatible staged flavours such as AnimeGen, this loads the fixed peer stage alongside the trained stage for validation renders.", + importance=ImportanceLevel.ADVANCED, + order=32, + subsection="advanced", + model_specific=["wan"], + documentation="OPTIONS.md#--wan_validation_load_other_stage", + ) + ) diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/__init__.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/__init__.py index f216c6488..92e701da7 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/__init__.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/__init__.py @@ -2,6 +2,8 @@ from typing import TYPE_CHECKING, Callable, Iterable +from simpletuner.helpers.models.field_registry import register_model_field_registries + from . import advanced, data, logging_fields, lora, loss, memory, model, optimizer, publishing, training, validation if TYPE_CHECKING: @@ -28,3 +30,4 @@ def register_all_sections(registry: "FieldRegistry") -> None: for registrar in _REGISTRARS: registrar(registry) + register_model_field_registries(registry) diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py index ea456e031..5b41b3aa1 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py @@ -412,22 +412,6 @@ def register_advanced_fields(registry: "FieldRegistry") -> None: ) ) - registry._add_field( - ConfigField( - name="flux_fast_schedule", - arg_name="--flux_fast_schedule", - ui_label="Flow Fast Schedule", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - default_value=False, - help_text="Use experimental fast schedule for Flux.1S training", - tooltip="Experimental feature that may improve training speed for Flux.1S models.", - importance=ImportanceLevel.EXPERIMENTAL, - order=11, - ) - ) - registry._add_field( ConfigField( name="flow_use_uniform_schedule", @@ -571,115 +555,6 @@ def register_advanced_fields(registry: "FieldRegistry") -> None: ) ) - # Flux Guidance Configuration - registry._add_field( - ConfigField( - name="flux_guidance_mode", - arg_name="--flux_guidance_mode", - ui_label="Flux Guidance Mode", - field_type=FieldType.SELECT, - tab="model", - section="model_specific", - default_value="constant", - choices=[{"value": "constant", "label": "Constant"}, {"value": "random-range", "label": "Random Range"}], - help_text="Guidance mode for Flux training", - tooltip="Constant uses same guidance for all samples. Random Range varies guidance per sample.", - importance=ImportanceLevel.ADVANCED, - order=40, - dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], - ) - ) - - # Flux Attention Masked Training - registry._add_field( - ConfigField( - name="flux_attention_masked_training", - arg_name="--flux_attention_masked_training", - ui_label="Attention Masked Training", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - default_value=False, - model_specific=["flux"], - help_text="Enable attention masked training for Flux models", - tooltip="Experimental feature for Flux models that masks certain attention patterns during training.", - importance=ImportanceLevel.EXPERIMENTAL, - order=10, - ) - ) - - # Flux Fast Schedule - registry._add_field( - ConfigField( - name="flux_fast_schedule", - arg_name="--flux_fast_schedule", - ui_label="Fast Training Schedule", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - default_value=False, - model_specific=["flux"], - help_text="Use experimental fast schedule for Flux training", - tooltip="Experimental feature that may speed up Flux.1S training at the cost of quality.", - importance=ImportanceLevel.EXPERIMENTAL, - order=11, - ) - ) - - registry._add_field( - ConfigField( - name="flux_guidance_value", - arg_name="--flux_guidance_value", - ui_label="Flux Guidance Value", - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - default_value=1.0, - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], - help_text="Guidance value for constant mode", - tooltip="1.0 preserves CFG distillation. Higher values require CFG at inference.", - importance=ImportanceLevel.ADVANCED, - order=41, - dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], - ) - ) - - registry._add_field( - ConfigField( - name="flux_guidance_min", - arg_name="--flux_guidance_min", - ui_label="Flux Guidance Min", - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - default_value=0.0, - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], - help_text="Minimum guidance value for random-range mode", - tooltip="Lower bound of guidance range when using random-range mode.", - importance=ImportanceLevel.ADVANCED, - order=42, - dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], - ) - ) - - registry._add_field( - ConfigField( - name="flux_guidance_max", - arg_name="--flux_guidance_max", - ui_label="Flux Guidance Max", - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - default_value=4.0, - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], - help_text="Maximum guidance value for random-range mode", - tooltip="Upper bound of guidance range when using random-range mode.", - importance=ImportanceLevel.ADVANCED, - order=43, - dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True, action="show")], - ) - ) - # T5 Configuration registry._add_field( ConfigField( @@ -700,44 +575,6 @@ def register_advanced_fields(registry: "FieldRegistry") -> None: ) ) - registry._add_field( - ConfigField( - name="sd3_clip_uncond_behaviour", - arg_name="--sd3_clip_uncond_behaviour", - ui_label="SD3 CLIP Unconditional Behavior", - field_type=FieldType.SELECT, - tab="training", - section="text_encoder", - subsection="advanced", - default_value="empty_string", - choices=[{"value": "empty_string", "label": "Empty String"}, {"value": "zero", "label": "Zero"}], - dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True)], - help_text="How SD3 handles unconditional prompts", - tooltip="Affects how SD3 processes prompts without conditioning. Empty string is default.", - importance=ImportanceLevel.ADVANCED, - order=11, - ) - ) - - registry._add_field( - ConfigField( - name="sd3_t5_uncond_behaviour", - arg_name="--sd3_t5_uncond_behaviour", - ui_label="SD3 T5 Unconditional Behavior", - field_type=FieldType.SELECT, - tab="training", - section="text_encoder", - subsection="advanced", - default_value=None, - choices=[{"value": "empty_string", "label": "Empty String"}, {"value": "zero", "label": "Zero"}], - dependencies=[FieldDependency(field="i_know_what_i_am_doing", operator="equals", value=True)], - help_text="How SD3 T5 handles unconditional prompts", - tooltip="Overrides CLIP behavior for T5. If not set, follows CLIP setting.", - importance=ImportanceLevel.ADVANCED, - order=12, - ) - ) - # Soft Min SNR Configuration registry._add_field( ConfigField( diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/logging_fields.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/logging_fields.py index 74bcdb44f..d68d739ee 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/logging_fields.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/logging_fields.py @@ -488,196 +488,6 @@ def register_logging_fields(registry: "FieldRegistry") -> None: ) ) - # LTX Model Configuration - registry._add_field( - ConfigField( - name="ltx_train_mode", - arg_name="--ltx_train_mode", - ui_label="LTX Train Mode", - field_type=FieldType.SELECT, - tab="model", - section="model_specific", - model_specific=["ltx"], - default_value="i2v", - choices=[ - {"value": "t2v", "label": "Text-to-Video"}, - {"value": "i2v", "label": "Image-to-Video"}, - ], - help_text="Training mode for LTX models", - tooltip="Choose whether datasets default to text-to-video (t2v) or image-to-video (i2v) processing.", - importance=ImportanceLevel.ADVANCED, - order=22, - ) - ) - - registry._add_field( - ConfigField( - name="ltx_i2v_prob", - arg_name="--ltx_i2v_prob", - ui_label="LTX Image-to-Video Probability", - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - model_specific=["ltx"], - default_value=0.1, - validation_rules=[ - ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be between 0 and 1"), - ValidationRule(ValidationRuleType.MAX, value=1.0, message="Must be between 0 and 1"), - ], - help_text="Probability of using image-to-video training for LTX", - tooltip="Fraction of training that uses image-to-video instead of video-to-video.", - importance=ImportanceLevel.ADVANCED, - order=23, - ) - ) - - registry._add_field( - ConfigField( - name="ltx_partial_noise_fraction", - arg_name="--ltx_partial_noise_fraction", - ui_label="LTX Partial Noise Fraction", - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - model_specific=["ltx"], - default_value=0.05, - validation_rules=[ - ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be between 0 and 1"), - ValidationRule(ValidationRuleType.MAX, value=1.0, message="Must be between 0 and 1"), - ], - help_text="Fraction of noise to add for LTX partial training", - tooltip="Controls how much noise to add during partial training steps.", - importance=ImportanceLevel.ADVANCED, - order=24, - ) - ) - - registry._add_field( - ConfigField( - name="ltx_protect_first_frame", - arg_name="--ltx_protect_first_frame", - ui_label="LTX Protect First Frame", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - model_specific=["ltx"], - default_value=False, - help_text="Protect the first frame from noise in LTX training", - tooltip="When enabled, first frame is kept clean while subsequent frames get noise.", - importance=ImportanceLevel.ADVANCED, - order=25, - ) - ) - - registry._add_field( - ConfigField( - name="ltx2_intrinsic_conditioning", - arg_name="--ltx2_intrinsic_conditioning", - ui_label="LTX-2 Intrinsic Conditioning", - field_type=FieldType.TEXT_JSON, - tab="model", - section="model_specific", - model_specific=["ltxvideo2"], - default_value=None, - help_text="JSON array of LTX-2 intrinsic clean-token conditioning objects", - tooltip="Advanced LTX-2 training config. Supported condition types: first_frame, prefix, suffix, spatial_crop, mask.", - importance=ImportanceLevel.ADVANCED, - order=26, - documentation="OPTIONS.md#ltx-2-intrinsic-and-reference-conditioning", - ) - ) - - for offset, field_name, label, help_text in ( - ( - 0, - "ltx2_first_frame_conditioning_probability", - "LTX-2 First Frame Conditioning Probability", - "Probability of replacing first-frame target tokens with clean latents and removing them from video loss", - ), - ( - 1, - "ltx2_prefix_conditioning_probability", - "LTX-2 Prefix Conditioning Probability", - "Probability of replacing prefix target tokens with clean latents and removing them from video loss", - ), - ( - 2, - "ltx2_suffix_conditioning_probability", - "LTX-2 Suffix Conditioning Probability", - "Probability of replacing suffix target tokens with clean latents and removing them from video loss", - ), - ( - 3, - "ltx2_mask_conditioning_probability", - "LTX-2 Mask Conditioning Probability", - "Probability of using mask=1 target tokens as clean conditioning with no video loss", - ), - ): - registry._add_field( - ConfigField( - name=field_name, - arg_name=f"--{field_name}", - ui_label=label, - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - model_specific=["ltxvideo2"], - default_value=0.0, - validation_rules=[ - ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be between 0 and 1"), - ValidationRule(ValidationRuleType.MAX, value=1.0, message="Must be between 0 and 1"), - ], - help_text=help_text, - tooltip=help_text, - importance=ImportanceLevel.ADVANCED, - order=27 + offset, - documentation="OPTIONS.md#ltx-2-intrinsic-and-reference-conditioning", - ) - ) - - for offset, field_name, label, default_value in ( - (0, "ltx2_prefix_conditioning_frames", "LTX-2 Prefix Conditioning Frames", 1), - (1, "ltx2_suffix_conditioning_frames", "LTX-2 Suffix Conditioning Frames", 1), - (2, "ltx2_reference_spatial_scale_factor", "LTX-2 Reference Spatial Scale Factor", None), - (3, "ltx2_reference_temporal_scale_factor", "LTX-2 Reference Temporal Scale Factor", 1), - ): - registry._add_field( - ConfigField( - name=field_name, - arg_name=f"--{field_name}", - ui_label=label, - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - model_specific=["ltxvideo2"], - default_value=default_value, - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=1, message="Must be at least 1")], - help_text=f"Advanced LTX-2 conditioning setting: {label.lower()}", - tooltip=f"Advanced LTX-2 conditioning setting: {label.lower()}", - importance=ImportanceLevel.ADVANCED, - order=31 + offset, - documentation="OPTIONS.md#ltx-2-intrinsic-and-reference-conditioning", - ) - ) - - registry._add_field( - ConfigField( - name="krea2_reference_latents", - arg_name="--krea2_reference_latents", - ui_label="Krea 2 Reference Latents", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - model_specific=["krea2"], - default_value=False, - help_text="Enable Krea 2 reference-dataset training with image-context prompt embeds and clean reference latents.", - tooltip="When enabled, Krea 2 requires paired conditioning data, encodes prompts with the reference image through Qwen3VL, and appends the clean reference latents to the transformer input.", - importance=ImportanceLevel.ADVANCED, - order=36, - documentation="OPTIONS.md#--krea2_reference_latents", - ) - ) - # Offload Parameter Path registry._add_field( ConfigField( @@ -787,162 +597,3 @@ def register_logging_fields(registry: "FieldRegistry") -> None: order=31, ) ) - - # SDXL Refiner Full Range - registry._add_field( - ConfigField( - name="sdxl_refiner_uses_full_range", - arg_name="--sdxl_refiner_uses_full_range", - ui_label="SDXL Refiner Full Range", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - default_value=False, - help_text="Use full timestep range for SDXL refiner", - tooltip="When enabled, refiner uses full timestep range instead of just high timesteps.", - importance=ImportanceLevel.ADVANCED, - dependencies=[FieldDependency(field="model_family", operator="equals", value="sdxl")], - order=32, - ) - ) - - # Ideogram 4 structured caption conversion - registry._add_field( - ConfigField( - name="ideogram_auto_json", - arg_name="--ideogram_auto_json", - ui_label="Ideogram Auto JSON", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - model_specific=["ideogram"], - default_value=True, - help_text="Convert non-JSON Ideogram 4 prompts into the structured JSON caption schema.", - tooltip="When enabled, plain validation prompts are wrapped into Ideogram 4's JSON caption format.", - importance=ImportanceLevel.ADVANCED, - dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], - order=33, - ) - ) - - registry._add_field( - ConfigField( - name="ideogram_validation", - arg_name="--ideogram_validation", - ui_label="Ideogram Validation", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - model_specific=["ideogram"], - default_value=False, - help_text="Temporarily enable Ideogram validation by reusing the conditional transformer for CFG's unconditional pass.", - tooltip="Validation is off by default for Ideogram until separate unconditional-model handling is implemented.", - importance=ImportanceLevel.ADVANCED, - dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], - order=34, - ) - ) - - registry._add_field( - ConfigField( - name="ideogram_prompt_upsample", - arg_name="--ideogram_prompt_upsample", - ui_label="Ideogram Prompt Upsample", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - model_specific=["ideogram"], - default_value=False, - help_text="Use Ideogram 4's prompt enhancer to rewrite prompts before text embedding cache generation.", - tooltip="When enabled and supported by the loaded pipeline, captions are expanded with Ideogram's prompt upsampler before JSON conversion and encoding.", - importance=ImportanceLevel.ADVANCED, - dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], - order=35, - ) - ) - - registry._add_field( - ConfigField( - name="ideogram_prompt_enhancer_head_id", - arg_name="--ideogram_prompt_enhancer_head_id", - ui_label="Ideogram Prompt Enhancer Head", - field_type=FieldType.TEXT, - tab="model", - section="model_specific", - model_specific=["ideogram"], - default_value="diffusers/qwen3-vl-8b-instruct-lm-head", - help_text="Hugging Face repo id for Ideogram 4's prompt upsampling LM head.", - tooltip="Used when --ideogram_prompt_upsample is enabled to rewrite prompts into Ideogram's structured JSON caption schema.", - importance=ImportanceLevel.ADVANCED, - dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], - order=36, - ) - ) - - registry._add_field( - ConfigField( - name="ideogram_schedule_mu", - arg_name="--ideogram_schedule_mu", - ui_label="Ideogram Schedule Mu", - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - model_specific=["ideogram"], - default_value=0.0, - help_text="Base mean used by Ideogram 4's resolution-aware logit-normal training timestep schedule.", - tooltip="Matches the default mu used by the vendored Ideogram validation pipeline.", - importance=ImportanceLevel.ADVANCED, - dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], - order=37, - ) - ) - - registry._add_field( - ConfigField( - name="ideogram_schedule_std", - arg_name="--ideogram_schedule_std", - ui_label="Ideogram Schedule Std", - field_type=FieldType.NUMBER, - tab="model", - section="model_specific", - model_specific=["ideogram"], - default_value=1.5, - help_text="Standard deviation used by Ideogram 4's logit-normal training timestep schedule.", - tooltip="Matches the default std used by the vendored Ideogram validation pipeline.", - importance=ImportanceLevel.ADVANCED, - dependencies=[FieldDependency(field="model_family", operator="equals", value="ideogram")], - order=38, - ) - ) - - # Sana Complex Human Instruction - registry._add_field( - ConfigField( - name="sana_complex_human_instruction", - arg_name="--sana_complex_human_instruction", - ui_label="Sana Complex Human Instruction", - field_type=FieldType.TEXT, - tab="model", - section="model_specific", - model_specific=["sana"], - default_value="\n".join( - [ - "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", - "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", - "- Do not add detail that is not in the original prompt! Do not mention colors if the User Prompt does not!", - "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", - "- If the prompt contains explicit material, do not remove this information from the output, which would damage the performance of downstream classifiers.", - "Here are examples of how to transform or refine prompts:", - "- User Prompt: A cat sleeping -> Enhanced: A cat sleeping peacefully, showcasing the joy of pet ownership. Cute floof kitty cat gatto.", - "- User Prompt: A busy city street -> Enhanced: A bustling city street scene featuring a crowd of people.", - "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", - "User Prompt: ", - ] - ), - placeholder="complex human instruction", - help_text="Complex human instruction for Sana model training", - tooltip="Special instruction format for Sana model training with complex human prompts.", - importance=ImportanceLevel.ADVANCED, - order=33, - ) - ) diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/lora.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/lora.py index f4f9f9a1b..394b6c248 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/lora.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/lora.py @@ -396,72 +396,6 @@ def register_lora_fields(registry: "FieldRegistry") -> None: ) ) - # Flux LoRA Target - flux_targets = [ - "mmdit", - "context", - "context+ffs", - "all", - "all+ffs", - "ai-toolkit", - "tiny", - "nano", - "controlnet", - "all+ffs+embedder", - "all+ffs+embedder+controlnet", - ] - registry._add_field( - ConfigField( - name="flux_lora_target", - arg_name="--flux_lora_target", - ui_label="Flux LoRA Target Layers", - field_type=FieldType.SELECT, - tab="model", - section="lora_config", - subsection="model_specific", - default_value="all", - choices=[{"value": t, "label": t} for t in flux_targets], - dependencies=[ - FieldDependency(field="model_type", value="lora"), - FieldDependency(field="model_family", value="flux"), - ], - help_text="Which layers to train in Flux models", - tooltip="'all' trains all attention layers. 'context' only trains text layers. '+ffs' includes feed-forward layers.", - importance=ImportanceLevel.ADVANCED, - model_specific=["flux"], - order=10, - ) - ) - - # ACE-Step LoRA Target - 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, - ) - ) - # Use DoRA registry._add_field( ConfigField( diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/loss.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/loss.py index c44b19198..bdeb8a056 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/loss.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/loss.py @@ -117,42 +117,6 @@ def register_loss_fields(registry: "FieldRegistry") -> None: ) ) - # HiDream Load Balancing Loss Toggle - registry._add_field( - ConfigField( - name="hidream_use_load_balancing_loss", - arg_name="--hidream_use_load_balancing_loss", - ui_label="Enable HiDream Load Balancing Loss", - field_type=FieldType.CHECKBOX, - tab="training", - section="loss_functions", - default_value=False, - dependencies=[FieldDependency(field="model_family", operator="equals", value="hidream")], - help_text="Apply experimental load balancing loss when training HiDream models.", - tooltip="Balances expert contributions during HiDream training. Only available for HiDream model family.", - importance=ImportanceLevel.EXPERIMENTAL, - order=6, - ) - ) - - # HiDream Load Balancing Weight - registry._add_field( - ConfigField( - name="hidream_load_balancing_loss_weight", - arg_name="--hidream_load_balancing_loss_weight", - ui_label="HiDream Load Balancing Weight", - field_type=FieldType.NUMBER, - tab="training", - section="loss_functions", - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0.0, message="Must be non-negative")], - dependencies=[FieldDependency(field="hidream_use_load_balancing_loss", operator="equals", value=True)], - help_text="Strength multiplier for HiDream load balancing loss.", - tooltip="Adjust if you need stronger balancing between experts. Leave blank to use the trainer default.", - importance=ImportanceLevel.EXPERIMENTAL, - order=7, - ) - ) - registry._add_field( ConfigField( name="crepa_enabled", diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py index 6af694fc0..dafe9f727 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py @@ -864,44 +864,6 @@ def _quant_label(value: str) -> str: ) ) - registry._add_field( - ConfigField( - name="wan_force_2_1_time_embedding", - arg_name="--wan_force_2_1_time_embedding", - ui_label="Force Wan 2.1 Time Embedding", - field_type=FieldType.CHECKBOX, - tab="model", - section="model_specific", - default_value=False, - dependencies=[FieldDependency(field="model_family", operator="equals", value="wan", action="show")], - help_text="Use Wan 2.1 style time embeddings even when running Wan 2.2 checkpoints.", - tooltip="Enable this if Wan 2.2 checkpoints report shape mismatches in the time embedding layers.", - importance=ImportanceLevel.ADVANCED, - order=30, - ) - ) - - # Fused QKV Projections - registry._add_field( - ConfigField( - name="fused_qkv_projections", - arg_name="--fuse_qkv_projections", - ui_label="Fused QKV Projections", - field_type=FieldType.CHECKBOX, - tab="model", - section="architecture", - default_value=False, - dependencies=[FieldDependency(field="model_family", operator="equals", value="flux")], - help_text="Enables Flash Attention 3 when supported; otherwise falls back to PyTorch SDPA.", - tooltip="Improves attention efficiency on modern NVIDIA GPUs. Uses native SDPA when Flash Attention 3 is unavailable.", - importance=ImportanceLevel.EXPERIMENTAL, - model_specific=["flux"], - order=19, - aliases=["--fused_qkv_projections"], - documentation="OPTIONS.md#--fuse_qkv_projections", - ) - ) - registry._add_field( ConfigField( name="rescale_betas_zero_snr", @@ -1250,28 +1212,6 @@ def _quant_label(value: str) -> str: ) ) - # Custom Text Encoder Intermediary Layers - registry._add_field( - ConfigField( - name="custom_text_encoder_intermediary_layers", - arg_name="--custom_text_encoder_intermediary_layers", - ui_label="Custom Text Encoder Layers", - field_type=FieldType.TEXT, - tab="model", - section="architecture", - subsection="advanced", - default_value=None, - placeholder="[10, 20, 30]", - dependencies=[FieldDependency(field="model_family", operator="in", values=["flux2"])], - help_text="Override which hidden state layers to extract from the text encoder. Provide as JSON array (e.g., [10, 20, 30]). Leave blank to use model defaults.", - tooltip="FLUX.2-dev uses layers [10, 20, 30] from Mistral-3, Klein models use [9, 18, 27] from Qwen3. Override for experimentation.", - importance=ImportanceLevel.EXPERIMENTAL, - model_specific=["flux2"], - order=34, - documentation="OPTIONS.md#--custom_text_encoder_intermediary_layers", - ) - ) - # Grounding: max entities per image registry._add_field( ConfigField( diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/training.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/training.py index de60a0bdd..1e7c3327b 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/training.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/training.py @@ -653,46 +653,6 @@ def register_training_fields(registry: "FieldRegistry") -> None: ) ) - # Feed-forward chunking (Wan) - registry._add_field( - ConfigField( - name="enable_chunked_feed_forward", - arg_name="--enable_chunked_feed_forward", - ui_label="Enable Feed-Forward Chunking", - field_type=FieldType.CHECKBOX, - tab="model", - section="memory_optimization", - default_value=False, - help_text="Split Wan feed-forward layers into smaller chunks to reduce peak VRAM usage.", - tooltip="Available for Wan models. Breaks long MLPs into mini-batches so checkpoint recomputes allocate less memory.", - importance=ImportanceLevel.ADVANCED, - model_specific=["wan", "wan_s2v"], - order=8, - ) - ) - - registry._add_field( - ConfigField( - name="feed_forward_chunk_size", - arg_name="--feed_forward_chunk_size", - ui_label="Feed-Forward Chunk Size", - field_type=FieldType.NUMBER, - tab="model", - section="memory_optimization", - default_value=None, - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=1, message="Chunk size must be at least 1")], - help_text="Number of samples processed per chunk when feed-forward chunking is enabled.", - tooltip="Leave blank for auto. Lower values reduce memory further but increase wall-clock time.", - importance=ImportanceLevel.ADVANCED, - model_specific=["wan", "wan_s2v"], - order=9, - dependencies=[ - FieldDependency(field="enable_chunked_feed_forward", operator="equals", value=True, action="show") - ], - allow_empty=True, - ) - ) - # Train Text Encoder registry._add_field( ConfigField( @@ -728,84 +688,6 @@ def register_training_fields(registry: "FieldRegistry") -> None: ) ) - # Lyrics Embedder Training (ACE-Step) - 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="Unlocks 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, - ) - ) - # LR Number of Cycles registry._add_field( ConfigField( diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/validation.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/validation.py index be28e2241..9bf28ff69 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/validation.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/validation.py @@ -126,47 +126,6 @@ def register_validation_fields(registry: "FieldRegistry") -> None: ) ) - # Validation Lyrics - 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"], - ) - ) - - # Validation Audio Duration - 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"], - ) - ) - # Number of Validation Images registry._add_field( ConfigField( @@ -479,44 +438,6 @@ def register_validation_fields(registry: "FieldRegistry") -> None: ) ) - # Validation Guidance Real - registry._add_field( - ConfigField( - name="validation_guidance_real", - arg_name="--validation_guidance_real", - ui_label="Real CFG (Distilled Models)", - field_type=FieldType.NUMBER, - tab="validation", - section="validation_guidance", - default_value=1.0, - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=1.0, message="Must be at least 1.0")], - help_text="CFG value for distilled models (e.g., FLUX schnell)", - tooltip="Use 1.0 for no CFG (distilled models). Higher values for real CFG sampling.", - importance=ImportanceLevel.ADVANCED, - order=2, - model_specific=["flux", "flux2"], - ) - ) - - # Validation No CFG Until Timestep - registry._add_field( - ConfigField( - name="validation_no_cfg_until_timestep", - arg_name="--validation_no_cfg_until_timestep", - ui_label="Skip CFG Until Timestep", - field_type=FieldType.NUMBER, - tab="validation", - section="validation_guidance", - default_value=2, - validation_rules=[ValidationRule(ValidationRuleType.MIN, value=0, message="Must be non-negative")], - help_text="Skip CFG for initial timesteps (Flux only)", - tooltip="For Flux real CFG: skip CFG on these initial timesteps. Default: 2", - importance=ImportanceLevel.ADVANCED, - order=3, - model_specific=["flux", "flux2"], - ) - ) - # Validation Negative Prompt registry._add_field( ConfigField( @@ -879,42 +800,6 @@ def register_validation_fields(registry: "FieldRegistry") -> None: ) ) - # Validation Audio Only (LTX-2) - registry._add_field( - ConfigField( - name="validation_audio_only", - arg_name="--validation_audio_only", - ui_label="Validation Audio Only", - field_type=FieldType.CHECKBOX, - tab="validation", - section="validation_options", - default_value=False, - help_text="Disable video generation during validation and emit audio only.", - tooltip="LTX-2 only: skips video generation during validation so only audio outputs are produced.", - importance=ImportanceLevel.ADVANCED, - order=17, - model_specific=["ltxvideo2"], - ) - ) - - registry._add_field( - ConfigField( - name="validation_ltx2_video_conditioning", - arg_name="--validation_ltx2_video_conditioning", - ui_label="LTX-2 Validation Video Conditioning", - field_type=FieldType.TEXT_JSON, - tab="validation", - section="validation_options", - default_value=None, - help_text="JSON list of IC-LoRA reference videos for LTX-2 validation", - tooltip="LTX-2 only: pass reference videos as paths, [path, strength] pairs, or objects with path/video_path and optional strength.", - importance=ImportanceLevel.ADVANCED, - order=18, - model_specific=["ltxvideo2"], - documentation="OPTIONS.md#ltx-2-conditioning-options", - ) - ) - # Validation Resolution registry._add_field( ConfigField( @@ -955,282 +840,6 @@ def register_validation_fields(registry: "FieldRegistry") -> 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", - ) - ) - - registry._add_field( - ConfigField( - name="wan_validation_load_other_stage", - arg_name="--wan_validation_load_other_stage", - ui_label="Wan Paired-Stage Validation", - field_type=FieldType.CHECKBOX, - tab="validation", - section="validation_options", - default_value=False, - help_text="Load the opposite Wan 2.2 stage during validation so the pipeline can switch denoisers at the stage boundary.", - tooltip="For Wan 2.2 and compatible staged flavours such as AnimeGen, this loads the fixed peer stage alongside the trained stage for validation renders.", - importance=ImportanceLevel.ADVANCED, - order=32, - subsection="advanced", - model_specific=["wan"], - documentation="OPTIONS.md#--wan_validation_load_other_stage", - ) - ) - - registry._add_field( - ConfigField( - name="sdxl_validation_pipeline_mode", - arg_name="--sdxl_validation_pipeline_mode", - ui_label="SDXL Validation Pipeline", - field_type=FieldType.SELECT, - tab="validation", - section="validation_options", - default_value="trained-stage", - choices=[ - {"value": "trained-stage", "label": "Trained Stage Only"}, - {"value": "full-pipeline", "label": "Full Pipeline"}, - ], - help_text="Choose whether SDXL validation runs only the trained stage or chains the base/refiner split pipeline.", - tooltip="Full pipeline runs stage 1 to the refiner boundary as latents, then resumes through stage 2.", - importance=ImportanceLevel.ADVANCED, - order=33, - subsection="advanced", - model_specific=["sdxl"], - documentation="OPTIONS.md#--sdxl_validation_pipeline_mode", - ) - ) - - registry._add_field( - ConfigField( - name="sdxl_validation_stage1_model", - arg_name="--sdxl_validation_stage1_model", - ui_label="SDXL Stage 1 Model", - field_type=FieldType.TEXT, - tab="validation", - section="validation_options", - default_value=None, - placeholder="stabilityai/stable-diffusion-xl-base-1.0", - help_text="Fixed SDXL stage 1 model used when validating a trained refiner through the full pipeline.", - tooltip="Leave blank to infer the matching SDXL base model from the selected flavour.", - importance=ImportanceLevel.ADVANCED, - order=34, - subsection="advanced", - model_specific=["sdxl"], - documentation="OPTIONS.md#--sdxl_validation_stage1_model", - ) - ) - - registry._add_field( - ConfigField( - name="sdxl_validation_stage2_model", - arg_name="--sdxl_validation_stage2_model", - ui_label="SDXL Stage 2 Model", - field_type=FieldType.TEXT, - tab="validation", - section="validation_options", - default_value=None, - placeholder="stabilityai/stable-diffusion-xl-refiner-1.0", - help_text="Fixed SDXL refiner model used when validating a trained stage 1 model through the full pipeline.", - tooltip="Leave blank to infer the matching SDXL refiner from the selected flavour.", - importance=ImportanceLevel.ADVANCED, - order=35, - subsection="advanced", - model_specific=["sdxl"], - documentation="OPTIONS.md#--sdxl_validation_stage2_model", - ) - ) - registry._add_field( ConfigField( name="validation_adapter_path", diff --git a/tests/test_model_field_registry.py b/tests/test_model_field_registry.py new file mode 100644 index 000000000..e84456328 --- /dev/null +++ b/tests/test_model_field_registry.py @@ -0,0 +1,62 @@ +"""Tests for model-owned field registry modules.""" + +import unittest + +from simpletuner.simpletuner_sdk.server.services.field_registry.registry import FieldRegistry + + +class TestModelFieldRegistry(unittest.TestCase): + def setUp(self): + self.registry = FieldRegistry() + + def test_model_registry_modules_are_discovered(self): + expected_fields = [ + "deepfloyd_validation_pipeline_mode", + "wan_validation_load_other_stage", + "sdxl_validation_pipeline_mode", + "validation_lyrics", + "flux_lora_target", + "ltx_train_mode", + "ltx2_intrinsic_conditioning", + "ideogram_auto_json", + "sana_complex_human_instruction", + "hidream_use_load_balancing_loss", + "sd3_clip_uncond_behaviour", + "krea2_reference_latents", + ] + + for field_name in expected_fields: + with self.subTest(field=field_name): + self.assertIsNotNone(self.registry.get_field(field_name)) + + def test_model_specific_fields_keep_context_filtering(self): + flux_fields = { + field.name + for field in self.registry.get_fields_for_tab( + "model", + context={ + "model_family": "flux", + "model_type": "lora", + "i_know_what_i_am_doing": True, + }, + ) + } + self.assertIn("flux_lora_target", flux_fields) + self.assertNotIn("acestep_lora_target", flux_fields) + + ace_fields = { + field.name + for field in self.registry.get_fields_for_tab( + "model", + context={ + "model_family": "ace_step", + "model_type": "lora", + }, + ) + } + self.assertIn("acestep_lora_target", ace_fields) + self.assertNotIn("flux_lora_target", ace_fields) + + +if __name__ == "__main__": + unittest.main()