{
  "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": "rectified_flow",
      "distribution_shift_options": { "min_length": 256, "max_length": 4096 },
      "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,
        "log_loss_info": false,
        "pre_encoded": false,
        "timestep_sampler": "trunc_logit_normal",
        "optimizer_configs": {
          "diffusion": {
            "optimizer": {
              "type": "AdamW",
              "config": {
                "lr": 5e-5,
                "betas": [0.9, 0.95],
                "eps": 1e-8,
                "weight_decay": 0.01,
                "foreach": true
              }
            },
            "scheduler": {
              "type": "InverseLR",
              "config": { "inv_gamma": 500000, "power": 0.5, "warmup": 0.999 }
            }
          }
        },
    "demo": {
      "demo_every": 2000,
      "demo_steps": 50,
      "num_demos": 8,
      "demo_cond": [
        {"prompt": "Amen break 174 BPM", "seconds_total": 6},
        {"prompt": "People talking in a crowded cafe", "seconds_total": 10},
        {"prompt": "chillhop 91 bpm", "seconds_total": 6},
        {"prompt": "trap 120bpm", "seconds_total": 12},
        {"prompt": "A dog barking next to a waterfall", "seconds_total": 6},
        {"prompt": "Glitchy bass design, I used Serum for this", "seconds_total": 4},
        {"prompt": "chillhop 132 bpm", "seconds_total": 12},
        {"prompt": "Birds singing in the forest", "seconds_total": 10}
      ],
      "demo_cfg_scales": [1, 4, 7]
    }
  }
}