vall-e/vall_e/models/__init__.py

def get_model(cfg, training=True):
	name = cfg.name

	if not cfg.experimental:
		from .ar_nar import AR_NAR
		model = AR_NAR(
			n_tokens=cfg.tokens,
			d_model=cfg.dim,
			n_heads=cfg.heads,
			n_layers=cfg.layers,
			n_experts=cfg.experts,
			
			p_dropout=cfg.dropout,
			
			l_padding = cfg.input_alignment,
			
			training = training,
			config = cfg,
		)
		model._cfg = cfg
	else:
		from .experimental import Model as Experimental
		model = Experimental(
			d_model=cfg.dim,
			n_layers=cfg.layers,
			n_heads=cfg.heads,
			p_dropout=cfg.dropout,

			config = cfg,
		)

	print(f"{name} ({next(model.parameters()).dtype}): {sum(p.numel() for p in model.parameters() if p.requires_grad)} parameters")

	return model

def get_models(models, training=True):
	return { model.full_name: get_model(model, training=training) for model in models }