| { |
| "model_type": "diffusion_cond_inpaint", |
| "sample_size": 5324800, |
| "sample_rate": 44100, |
| "audio_channels": 2, |
| "model": { |
| "pretransform": { |
| "type": "autoencoder", |
| "iterate_batch": false, |
| "chunked": true, |
| "enable_grad": false, |
| "config": { |
| "pretransform": { |
| "type": "patched", |
| "enable_grad": false, |
| "config": { |
| "patch_size": 256, |
| "channels": 2 |
| } |
| }, |
| "encoder": { |
| "type": "taae_v2", |
| "requires_grad": false, |
| "config": { |
| "in_channels": 512, |
| "channels": 128, |
| "c_mults": [ |
| 6 |
| ], |
| "strides": [ |
| 16 |
| ], |
| "latent_dim": 256, |
| "transformer_depths": [ |
| 6 |
| ], |
| "checkpointing": false, |
| "differential": true, |
| "mapping_style": "none", |
| "dim_heads": 64, |
| "variable_stride": true, |
| "use_flash": true, |
| "dyt": true, |
| "chunk_size": 32, |
| "chunk_midpoint_shift": true, |
| "conv_mapping": false, |
| "mask_noise": 0.0 |
| } |
| }, |
| "decoder": { |
| "type": "taae_v2", |
| "requires_grad": false, |
| "config": { |
| "out_channels": 512, |
| "channels": 128, |
| "c_mults": [ |
| 6 |
| ], |
| "strides": [ |
| 16 |
| ], |
| "latent_dim": 256, |
| "transformer_depths": [ |
| 6 |
| ], |
| "sinusoidal_blocks": [ |
| 0 |
| ], |
| "checkpointing": false, |
| "differential": true, |
| "mapping_style": "none", |
| "dim_heads": 64, |
| "variable_stride": true, |
| "use_flash": true, |
| "dyt": true, |
| "chunk_size": 32, |
| "conv_mapping": true, |
| "chunk_midpoint_shift": true, |
| "freeze_backbone": false, |
| "mask_noise": 0.01 |
| } |
| }, |
| "bottleneck": { |
| "type": "softnorm", |
| "requires_grad": false, |
| "config": { |
| "dim": 256, |
| "noise_augment_dim": 0, |
| "noise_regularize": true, |
| "auto_scale": true, |
| "freeze": true |
| } |
| }, |
| "latent_dim": 256, |
| "downsampling_ratio": 4096, |
| "io_channels": 2 |
| } |
| }, |
| "conditioning": { |
| "configs": [ |
| { |
| "id": "prompt", |
| "type": "t5gemma", |
| "config": { |
| "max_length": 256, |
| "padding_mode": "learned", |
| "repo_id": "stabilityai/stable-audio-3-small-music", |
| "subfolder": "t5gemma-b-b-ul2" |
| } |
| }, |
| { |
| "id": "seconds_total", |
| "type": "number", |
| "config": { |
| "min_val": 0, |
| "max_val": 384, |
| "fourier_features_type": "expo" |
| } |
| } |
| ], |
| "cond_dim": 768 |
| }, |
| "diffusion": { |
| "cross_attention_cond_ids": [ |
| "prompt", |
| "seconds_total" |
| ], |
| "global_cond_ids": [ |
| "seconds_total" |
| ], |
| "local_add_cond_ids": [ |
| "inpaint_mask", |
| "inpaint_masked_input" |
| ], |
| "type": "dit", |
| "diffusion_objective": "rectified_flow", |
| "use_effective_length_for_schedule": true, |
| "mask_padding_attention": true, |
| "distribution_shift_options": { |
| "type": "full", |
| "min_length": 256, |
| "max_length": 4096 |
| }, |
| "config": { |
| "io_channels": 256, |
| "embed_dim": 1024, |
| "depth": 20, |
| "num_heads": 16, |
| "cond_token_dim": 768, |
| "global_cond_dim": 768, |
| "local_add_cond_dim": 257, |
| "global_cond_type": "adaLN", |
| "timestep_features_type": "expo", |
| "timestep_features_logsnr": false, |
| "attn_kwargs": { |
| "qk_norm": "rms", |
| "differential": false |
| }, |
| "norm_type": "rms_norm", |
| "norm_kwargs": { |
| "force_fp32": true |
| }, |
| "ff_kwargs": { |
| "mult": 4.0 |
| }, |
| "num_memory_tokens": 64 |
| } |
| }, |
| "io_channels": 256 |
| }, |
| "training": { |
| "use_ema": true, |
| "log_loss_info": false, |
| "pre_encoded": true, |
| "timestep_sampler": "trunc_logit_normal", |
| "mask_loss_weight": 1.0, |
| "log_every_n_steps": 100, |
| "silence_extension_scale_seconds": 4.0, |
| "ot_coupling": true, |
| "inpainting": { |
| "mask_kwargs": { |
| "mask_type_probabilities": [ |
| 0.2, |
| 0.6, |
| 0.2 |
| ] |
| } |
| }, |
| "optimizer_configs": { |
| "diffusion": { |
| "optimizer": { |
| "type": "MuonAdamW", |
| "config": { |
| "muon_lr": 0.001, |
| "muon_momentum": 0.95, |
| "fused_layer_patterns": [ |
| "*.to_qkv.*", |
| "*.to_kv.*", |
| "*.to_q.*", |
| "*.ff.*.proj.*" |
| ], |
| "adam_lr": 5e-05, |
| "adam_betas": [ |
| 0.9, |
| 0.95 |
| ], |
| "adam_weight_decay": 0.01 |
| } |
| }, |
| "scheduler": { |
| "type": "InverseLR", |
| "config": { |
| "inv_gamma": 1000000, |
| "power": 0.5, |
| "warmup": 0.995 |
| } |
| } |
| } |
| }, |
| "demo": { |
| "demo_every": 500, |
| "demo_steps": 50, |
| "demo_cfg_scales": [ |
| 2, |
| 4, |
| 7 |
| ], |
| "demo_cond": [ |
| { |
| "prompt": "A beautiful piano arpeggio grows into a grand cinematic climax", |
| "seconds_total": 119 |
| }, |
| { |
| "prompt": "Elegant and sophisticated Latin jazz piece with a Cuban base and a whispered melodic female voice", |
| "seconds_total": 100 |
| }, |
| { |
| "prompt": "Amen break 174 BPM", |
| "seconds_total": 9 |
| }, |
| { |
| "prompt": "lofi house loop", |
| "seconds_total": 35 |
| } |
| ], |
| "inpaint_demos": { |
| "num_random_spans": 2, |
| "num_causal": 1 |
| }, |
| "log_snr_sampling": true |
| } |
| } |
| } |
