""" Make model classes Load from checkpoints Test on dummy outputs to see if everything matches """ import os import fire import numpy as np import torch as t import jukebox.utils.dist_adapter as dist from jukebox.hparams import REMOTE_PREFIX, Hyperparams, setup_hparams from jukebox.utils.dist_utils import print_all from jukebox.utils.remote_utils import download from jukebox.utils.torch_utils import freeze_model from jukebox.vqvae.vqvae import calculate_strides MODELS = { "5b": ("vqvae", "upsampler_level_0", "upsampler_level_1", "prior_5b"), "5b_lyrics": ("vqvae", "upsampler_level_0", "upsampler_level_1", "prior_5b_lyrics"), "1b_lyrics": ("vqvae", "upsampler_level_0", "upsampler_level_1", "prior_1b_lyrics"), #'your_model': ("you_vqvae_here", "your_upsampler_here", ..., "you_top_level_prior_here") } def load_checkpoint(path): restore = path if restore.startswith(REMOTE_PREFIX): remote_path = restore cache_dir = os.environ.get("JUKEBOX_CACHE_DIR", "~/.cache") local_path = os.path.join( os.path.expanduser(cache_dir), remote_path[len(REMOTE_PREFIX) :] ) if dist.get_rank() % 8 == 0: print("Downloading from azure") if not os.path.exists(os.path.dirname(local_path)): os.makedirs(os.path.dirname(local_path)) if not os.path.exists(local_path): download(remote_path, local_path) restore = local_path dist.barrier() checkpoint = t.load(restore, map_location=t.device("cpu"), weights_only=False) print("Restored from {}".format(restore)) return checkpoint def save_checkpoint(logger, name, model, opt, metrics, hps): with t.no_grad(): save_hps = {**hps} save_hps = { k: v for k, v in save_hps.items() if k not in [ "metadata_v2", "metadata_v3", "alignments", "lyric_processor", "midi_processor", ] } t.save( { "hps": save_hps, "model": model.state_dict(), # should also save bottleneck k's as buffers "opt": opt.state_dict() if opt is not None else None, "step": logger.iters, **metrics, }, f"{logger.logdir}/checkpoint_{name}.pth.tar", ) return def restore_model(hps, model, checkpoint_path): model.step = 0 if checkpoint_path != "": checkpoint = load_checkpoint(checkpoint_path) # checkpoint_hps = Hyperparams(**checkpoint['hps']) # for k in set(checkpoint_hps.keys()).union(set(hps.keys())): # if checkpoint_hps.get(k, None) != hps.get(k, None): # print(k, "Checkpoint:", checkpoint_hps.get(k, None), "Ours:", hps.get(k, None)) checkpoint["model"] = { k[7:] if k[:7] == "module." else k: v for k, v in checkpoint["model"].items() } model.load_state_dict(checkpoint["model"], strict=False) if "step" in checkpoint: model.step = checkpoint["step"] def restore_opt(opt, shd, checkpoint_path): if not checkpoint_path: return checkpoint = load_checkpoint(checkpoint_path) if "opt" in checkpoint: opt.load_state_dict(checkpoint["opt"]) if "step" in checkpoint: shd.step(checkpoint["step"]) def make_vqvae(hps, device="cuda"): from jukebox.vqvae.vqvae import VQVAE block_kwargs = dict( width=hps.width, depth=hps.depth, m_conv=hps.m_conv, dilation_growth_rate=hps.dilation_growth_rate, dilation_cycle=hps.dilation_cycle, reverse_decoder_dilation=hps.vqvae_reverse_decoder_dilation, ) if not hps.sample_length: assert hps.sample_length_in_seconds != 0 downsamples = calculate_strides(hps.strides_t, hps.downs_t) top_raw_to_tokens = np.prod(downsamples) hps.sample_length = ( hps.sample_length_in_seconds * hps.sr // top_raw_to_tokens ) * top_raw_to_tokens print( f"Setting sample length to {hps.sample_length} (i.e. {hps.sample_length/hps.sr} seconds) to be multiple of {top_raw_to_tokens}" ) vqvae = VQVAE( input_shape=(hps.sample_length, 1), levels=hps.levels, downs_t=hps.downs_t, strides_t=hps.strides_t, emb_width=hps.emb_width, l_bins=hps.l_bins, mu=hps.l_mu, commit=hps.commit, spectral=hps.spectral, multispectral=hps.multispectral, multipliers=hps.hvqvae_multipliers, use_bottleneck=hps.use_bottleneck, **block_kwargs, ) vqvae = vqvae.to(device) restore_model(hps, vqvae, hps.restore_vqvae) if hps.train and not hps.prior: print_all("Loading vqvae in train mode") if hps.restore_vqvae != "": print_all("Reseting bottleneck emas") for level, bottleneck in enumerate(vqvae.bottleneck.level_blocks): num_samples = hps.sample_length downsamples = calculate_strides(hps.strides_t, hps.downs_t) raw_to_tokens = np.prod(downsamples[: level + 1]) num_tokens = (num_samples // raw_to_tokens) * dist.get_world_size() bottleneck.restore_k( num_tokens=num_tokens, threshold=hps.revival_threshold ) else: print_all("Loading vqvae in eval mode") vqvae.eval() freeze_model(vqvae) return vqvae def make_prior(hps, vqvae, device="cuda"): from jukebox.prior.prior import SimplePrior prior_kwargs = dict( input_shape=(hps.n_ctx,), bins=vqvae.l_bins, width=hps.prior_width, depth=hps.prior_depth, heads=hps.heads, attn_order=hps.attn_order, blocks=hps.blocks, spread=hps.spread, attn_dropout=hps.attn_dropout, resid_dropout=hps.resid_dropout, emb_dropout=hps.emb_dropout, zero_out=hps.zero_out, res_scale=hps.res_scale, pos_init=hps.pos_init, init_scale=hps.init_scale, m_attn=hps.m_attn, m_mlp=hps.m_mlp, checkpoint_res=hps.c_res if hps.train else 0, checkpoint_attn=hps.c_attn if hps.train else 0, checkpoint_mlp=hps.c_mlp if hps.train else 0, ) x_cond_kwargs = dict( out_width=hps.prior_width, init_scale=hps.init_scale, width=hps.cond_width, depth=hps.cond_depth, m_conv=hps.cond_m_conv, dilation_growth_rate=hps.cond_dilation_growth_rate, dilation_cycle=hps.cond_dilation_cycle, zero_out=hps.cond_zero_out, res_scale=hps.cond_res_scale, checkpoint_res=hps.cond_c_res, ) # have to keep this else names wrong y_cond_kwargs = dict( out_width=hps.prior_width, init_scale=hps.init_scale, y_bins=hps.y_bins, t_bins=hps.t_bins, sr=hps.sr, min_duration=hps.min_duration, max_duration=hps.max_duration, max_bow_genre_size=hps.max_bow_genre_size, ) if hps.use_tokens and not hps.single_enc_dec: prime_kwargs = dict( use_tokens=hps.use_tokens, prime_loss_fraction=hps.prime_loss_fraction, n_tokens=hps.n_tokens, bins=hps.n_vocab, width=hps.prime_width, depth=hps.prime_depth, heads=hps.prime_heads, attn_order=hps.prime_attn_order, blocks=hps.prime_blocks, spread=hps.prime_spread, attn_dropout=hps.prime_attn_dropout, resid_dropout=hps.prime_resid_dropout, emb_dropout=hps.prime_emb_dropout, zero_out=hps.prime_zero_out, res_scale=hps.prime_res_scale, pos_init=hps.prime_pos_init, init_scale=hps.prime_init_scale, m_attn=hps.prime_m_attn, m_mlp=hps.prime_m_mlp, checkpoint_res=hps.prime_c_res if hps.train else 0, checkpoint_attn=hps.prime_c_attn if hps.train else 0, checkpoint_mlp=hps.prime_c_mlp if hps.train else 0, ) else: prime_kwargs = dict( use_tokens=hps.use_tokens, prime_loss_fraction=hps.prime_loss_fraction, n_tokens=hps.n_tokens, bins=hps.n_vocab, ) # z_shapes for other levels given this level gets n_ctx codes rescale = lambda z_shape: (z_shape[0] * hps.n_ctx // vqvae.z_shapes[hps.level][0],) z_shapes = [rescale(z_shape) for z_shape in vqvae.z_shapes] prior = SimplePrior( z_shapes=z_shapes, l_bins=vqvae.l_bins, encoder=vqvae.encode, decoder=vqvae.decode, level=hps.level, downs_t=vqvae.downs_t, strides_t=vqvae.strides_t, labels=hps.labels, prior_kwargs=prior_kwargs, x_cond_kwargs=x_cond_kwargs, y_cond_kwargs=y_cond_kwargs, prime_kwargs=prime_kwargs, copy_input=hps.copy_input, labels_v3=hps.labels_v3, merged_decoder=hps.merged_decoder, single_enc_dec=hps.single_enc_dec, ) prior.alignment_head = hps.get("alignment_head", None) prior.alignment_layer = hps.get("alignment_layer", None) if hps.fp16_params: print_all("Converting to fp16 params") from jukebox.transformer.ops import _convert_conv_weights_to_fp16 prior.apply(_convert_conv_weights_to_fp16) prior = prior.to(device) restore_model(hps, prior, hps.restore_prior) if hps.train: print_all("Loading prior in train mode") pass else: print_all("Loading prior in eval mode") prior.eval() freeze_model(prior) return prior def make_model(model, device, hps, levels=None): vqvae, *priors = MODELS[model] vqvae = make_vqvae( setup_hparams( vqvae, dict( sample_length=hps.get("sample_length", 0), sample_length_in_seconds=hps.get("sample_length_in_seconds", 0), ), ), device, ) hps.sample_length = vqvae.sample_length if levels is None: levels = range(len(priors)) priors = [ make_prior(setup_hparams(priors[level], dict()), vqvae, "cpu") for level in levels ] return vqvae, priors def save_outputs(model, device, hps): # Check logits if hps.labels_v3: n_ctx = 6144 n_tokens = 384 prime_bins = 79 else: n_ctx = 8192 n_tokens = 512 prime_bins = 80 rng = t.random.manual_seed(0) x = ( 2 * t.rand((1, n_ctx * 8 * 4 * 4, 1), generator=rng, dtype=t.float).cuda() - 1.0 ) # -1 to 1 lyric_tokens = ( t.randint(0, prime_bins, (1, n_tokens), generator=rng, dtype=t.long) .view(-1) .numpy() ) artist_id = 10 genre_ids = [1] total_length = 2 * 2646000 offset = 2646000 vqvae, priors = make_model(model, device, hps) # encode vq_prior = priors[-1] zs = vq_prior.encode(x, start_level=0) x_ds = [ vq_prior.decode(zs[level:], start_level=level) for level in range(0, len(zs)) ] # priors data = dict(zs=zs, x_ds=x_ds) for level in range(len(priors)): print(f"Doing level {level}") if hps.labels_v3 and level != hps.levels - 1: print(f"Skipping level {level}") continue prior = priors[level] prior.cuda() x_in = x[:, : n_ctx * 8 * (4**level)] y_in = ( t.from_numpy( prior.labeller.get_y_from_ids( artist_id, genre_ids, lyric_tokens, total_length, offset ) ) .view(1, -1) .cuda() .long() ) x_out, _, metrics = prior( x_in, y_in, fp16=hps.fp16, get_preds=True, decode=True ) preds = metrics["preds"] data[level] = dict(x=x_in, y=y_in, x_out=x_out, preds=preds) prior.cpu() t.save(data, "data.pth.tar") dist.barrier() print("Saved data") exit() def run(model, port=29500, **kwargs): from jukebox.utils.dist_utils import setup_dist_from_mpi rank, local_rank, device = setup_dist_from_mpi(port=port) hps = Hyperparams(**kwargs) with t.no_grad(): save_outputs(model, device, hps) if __name__ == "__main__": fire.Fire(run)