| { | |
| "model_type": "diffusion_cond", | |
| "sample_size": 524288, | |
| "sample_rate": 44100, | |
| "audio_channels": 2, | |
| "model": { | |
| "pretransform": { | |
| "type": "autoencoder", | |
| "iterate_batch": false, | |
| "model_half": true, | |
| "config": { | |
| "encoder": { | |
| "type": "oobleck", | |
| "requires_grad": false, | |
| "config": { | |
| "in_channels": 2, | |
| "channels": 128, | |
| "c_mults": [1, 2, 4, 8, 16], | |
| "strides": [2, 4, 4, 8, 8], | |
| "latent_dim": 128, | |
| "use_snake": true | |
| } | |
| }, | |
| "decoder": { | |
| "type": "oobleck", | |
| "config": { | |
| "out_channels": 2, | |
| "channels": 128, | |
| "c_mults": [1, 2, 4, 8, 16], | |
| "strides": [2, 4, 4, 8, 8], | |
| "latent_dim": 64, | |
| "use_snake": true, | |
| "final_tanh": false | |
| } | |
| }, | |
| "bottleneck": { | |
| "type": "vae" | |
| }, | |
| "latent_dim": 64, | |
| "downsampling_ratio": 2048, | |
| "io_channels": 2 | |
| } | |
| }, | |
| "conditioning": { | |
| "configs": [ | |
| { | |
| "id": "prompt", | |
| "type": "t5", | |
| "config": { | |
| "t5_model_name": "t5-base", | |
| "max_length": 64 | |
| } | |
| }, | |
| { | |
| "id": "seconds_total", | |
| "type": "number", | |
| "config": { | |
| "min_val": 0, | |
| "max_val": 256 | |
| } | |
| } | |
| ], | |
| "cond_dim": 768 | |
| }, | |
| "diffusion": { | |
| "cross_attention_cond_ids": ["prompt", "seconds_total"], | |
| "global_cond_ids": ["seconds_total"], | |
| "diffusion_objective": "rf_denoiser", | |
| "type": "dit", | |
| "config": { | |
| "io_channels": 64, | |
| "embed_dim": 1024, | |
| "depth": 16, | |
| "num_heads": 8, | |
| "cond_token_dim": 768, | |
| "global_cond_dim": 768, | |
| "transformer_type": "continuous_transformer", | |
| "attn_kwargs": { | |
| "qk_norm": "ln" | |
| } | |
| } | |
| }, | |
| "io_channels": 64 | |
| }, | |
| "training": { | |
| "use_ema": true, | |
| "pre_encoded": false, | |
| "log_loss_info": false, | |
| "timestep_sampler": "trunc_logit_normal", | |
| "clip_grad_norm": 1.0, | |
| "cfg_dropout_prob": 0.0, | |
| "arc": { | |
| "ode_warmup": { | |
| "warmup_steps": 0, | |
| "refresh_rate": 10, | |
| "sampling_steps": 25, | |
| "cfg": 4 | |
| }, | |
| "noise_dist": { | |
| "generator": "logsnr_uniform", | |
| "discriminator": "logit_normal" | |
| }, | |
| "use_model_as_discriminator": true, | |
| "discriminator_base_ckpt": "/path/to/base/rf/model", | |
| "discriminator": { | |
| "type": "convnext", | |
| "dit_hidden_layer": 12, | |
| "weights": { | |
| "generator": 1.0, | |
| "discriminator": 1.0 | |
| }, | |
| "loss_type": "relativistic", | |
| "config": { | |
| "channels": 128, | |
| "strides": [1, 2, 2], | |
| "c_mults": [1, 2, 4], | |
| "num_blocks": [2, 2, 2] | |
| }, | |
| "contrastive": true, | |
| "include_grad_penalties": true | |
| } | |
| }, | |
| "optimizer_configs": { | |
| "diffusion": { | |
| "optimizer": { | |
| "type": "AdamW", | |
| "config": { | |
| "lr": 1e-6, | |
| "betas": [0.9, 0.95], | |
| "eps": 1e-8, | |
| "weight_decay": 0.01, | |
| "foreach": true | |
| } | |
| } | |
| }, | |
| "discriminator": { | |
| "optimizer": { | |
| "type": "AdamW", | |
| "config": { | |
| "lr": 1e-6, | |
| "betas": [0.9, 0.95], | |
| "eps": 1e-8, | |
| "weight_decay": 0.01, | |
| "foreach": true | |
| } | |
| } | |
| } | |
| }, | |
| "demo": { | |
| "demo_every": 250, | |
| "demo_steps": 8, | |
| "num_demos": 8, | |
| "demo_cond": [ | |
| {"prompt": "drum breaks 174 BPM", "seconds_total": 6}, | |
| {"prompt": "People talking in a crowded cafe", "seconds_total": 10}, | |
| {"prompt": "A short, beautiful piano riff in C minor", "seconds_total": 6}, | |
| {"prompt": "Tight Snare Drum", "seconds_total": 1}, | |
| {"prompt": "A dog barking next to a waterfall", "seconds_total": 6}, | |
| {"prompt": "Glitchy bass design", "seconds_total": 4}, | |
| {"prompt": "Synth pluck arp with reverb and delay, 128 BPM", "seconds_total": 6}, | |
| {"prompt": "Birds singing in the forest", "seconds_total": 10} | |
| ], | |
| "demo_cfg_scales": [1] | |
| } | |
| } | |
| } | |