Function make_tokenizer

tokenizer.py:542–580

(device='cuda')

def make_tokenizer(device='cuda'):
    generator_config = WaveCodec.Config(
        resnet_config=ResNetStack.Config(
            input_channels=1,
            output_channels=1,
            encode_channels=16,
            decode_channel_multiplier=4,
            kernel_size=7,
            bias=True,
            channel_ratios=(4, 8, 16, 16, 16, 16),
            strides=(2, 2, 4, 5, 5, 5),
            mode=None,
        ),
        use_weight_norm=True,

        compressor_config=GaussianZ.Config(
            dim=None,
            latent_dim=32,

            bias=True,
            use_weight_norm=True
        ),

        norm_stddev=0.05,
    )
    checkpoint = load_ckpt("inference_apatosaurus_95000", expected_hash="ba876edb97b988e9196e449dd176ca97")

    tokenizer = generator_config()

    load_result = tokenizer.load_state_dict(checkpoint, strict=False)
    print_colored(f"Loaded tokenizer state dict: {load_result}", "grey")

    tokenizer = tokenizer.eval()
    # Only convert to bfloat16 if using CUDA
    if device == 'cuda':
        tokenizer = tokenizer.bfloat16()
    tokenizer = tokenizer.to(device)
    tokenizer.requires_grad_ = False
    return tokenizer