| 373 | |
| 374 | |
| 375 | def get_hertz_dev_config(is_split=True, use_pure_audio_ablation=False): |
| 376 | if is_split: |
| 377 | checkpoints = [('inference_care_50000', 'e4ff4fe5c7e9f066410d2a5673b7a935'), ('inference_scion_54000', 'cb8bc484423922747b277ebc2933af5d')] |
| 378 | elif not use_pure_audio_ablation: |
| 379 | checkpoints = [('inference_whip_72000', '5e7cee7316900737d55fc5d44cc7a8f7'), ('inference_caraway_112000', 'fcb8368ef8ebf7712f3e31e6856da580')] |
| 380 | else: |
| 381 | checkpoints = [('inference_whip_72000', '5e7cee7316900737d55fc5d44cc7a8f7'), ('inference_syrup_110000', '353c48f553f1706824c11f3bb6a049e9')] |
| 382 | |
| 383 | quantizer_config=LatentQuantizer.Config( |
| 384 | from_pretrained=('inference_volcano_3', 'd42bf674022c5f84b051d5d7794f6169'), |
| 385 | compressor_config=FSQ.Config( |
| 386 | levels=[8,8,8,8,8], |
| 387 | dim=2048, |
| 388 | num_codebooks=1, |
| 389 | keep_num_codebooks_dim=None, |
| 390 | scale=None, |
| 391 | allowed_dtypes=['float32', 'float64', 'bfloat16'], |
| 392 | channel_first=False, |
| 393 | projection_has_bias=True, |
| 394 | return_indices=True, |
| 395 | force_quantization_f32=True, |
| 396 | use_rms=False |
| 397 | ), |
| 398 | dim=2048, |
| 399 | ff_dim=8192, |
| 400 | input_dim=32 |
| 401 | ) |
| 402 | |
| 403 | resynthesizer_config=TransformerVAE.Config( |
| 404 | io_config=GaussianMixtureIOLayer.Config( |
| 405 | latent_dim=32, |
| 406 | dim=4096, |
| 407 | num_components=8, |
| 408 | ), |
| 409 | stack_config=Stack.Config( |
| 410 | layers=8, |
| 411 | dim=4096, |
| 412 | seq_len=8192, |
| 413 | n_head=16, |
| 414 | ff_dim=11008, |
| 415 | kv_heads=16, |
| 416 | eps=1e-5, |
| 417 | theta=10_000 |
| 418 | ), |
| 419 | quantizer_config=quantizer_config, |
| 420 | plex_layer=None, |
| 421 | plex_roll=1, |
| 422 | split=is_split, |
| 423 | from_pretrained=checkpoints[0], |
| 424 | ) |
| 425 | |
| 426 | return HertzDevModel.Config( |
| 427 | dim=4096, |
| 428 | vocab_size=32_768, |
| 429 | stack_config=Stack.Config( |
| 430 | layers=32, |
| 431 | dim=4096, |
| 432 | seq_len=2048, |