Function build_model

questions/inference_server/models.py:331–377

(path, device)

def build_model(path, device):
    config = Path(__file__).parent / "config.json"
    assert config.exists(), f"Config path incorrect: config.json not found at {config}"
    with open(config, "r") as r:
        args = recursive_munch(json.load(r))
    assert args.decoder.type == "istftnet", f"Unknown decoder type: {args.decoder.type}"
    decoder = Decoder(
        dim_in=args.hidden_dim,
        style_dim=args.style_dim,
        dim_out=args.n_mels,
        resblock_kernel_sizes=args.decoder.resblock_kernel_sizes,
        upsample_rates=args.decoder.upsample_rates,
        upsample_initial_channel=args.decoder.upsample_initial_channel,
        resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
        upsample_kernel_sizes=args.decoder.upsample_kernel_sizes,
        gen_istft_n_fft=args.decoder.gen_istft_n_fft,
        gen_istft_hop_size=args.decoder.gen_istft_hop_size,
    )
    text_encoder = TextEncoder(channels=args.hidden_dim, kernel_size=5, depth=args.n_layer, n_symbols=args.n_token)
    predictor = ProsodyPredictor(
        style_dim=args.style_dim,
        d_hid=args.hidden_dim,
        nlayers=args.n_layer,
        max_dur=args.max_dur,
        dropout=args.dropout,
    )
    bert = load_plbert()
    bert_encoder = nn.Linear(bert.config.hidden_size, args.hidden_dim)
    for parent in [bert, bert_encoder, predictor, decoder, text_encoder]:
        for child in parent.children():
            if isinstance(child, nn.RNNBase):
                child.flatten_parameters()
    model = Munch(
        bert=bert.to(device).eval(),
        bert_encoder=bert_encoder.to(device).eval(),
        predictor=predictor.to(device).eval(),
        decoder=decoder.to(device).eval(),
        text_encoder=text_encoder.to(device).eval(),
    )
    for key, state_dict in torch.load(path, map_location="cpu", weights_only=True)["net"].items():
        assert key in model, key
        try:
            model[key].load_state_dict(state_dict)
        except:
            state_dict = {k[7:]: v for k, v in state_dict.items()}
            model[key].load_state_dict(state_dict, strict=False)
    return model