File size: 1,513 Bytes
9965bf6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import os
import torch
from load_utils import load_model
from guided_diffusion import dist_util
from guided_diffusion.gaussian_diffusion import _encode, _decode
from guided_diffusion.pr_datasets_all import load_data
from tqdm import tqdm
from guided_diffusion.midi_util import visualize_full_piano_roll, save_piano_roll_midi
from music_rule_guidance import music_rules
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")
plt.rcParams["figure.figsize"] = (20,3)
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
MODEL_NAME = 'kl/f8-all-onset'
MODEL_CKPT = 'taming-transformers/checkpoints/all_onset/epoch_14.ckpt'
TOTAL_BATCH = 256
def main():
data = load_data(
data_dir='datasets/all-len-40-gap-16-no-empty_train.csv',
batch_size=32,
class_cond=True,
image_size=1024,
deterministic=False,
fs=100,
)
embed_model = load_model(MODEL_NAME, MODEL_CKPT)
del embed_model.loss
embed_model.to(dist_util.dev())
embed_model.eval()
z_list = []
with torch.no_grad():
for _ in tqdm(range(TOTAL_BATCH)):
batch, cond = next(data)
batch = batch.to(dist_util.dev())
enc = _encode(batch, embed_model, scale_factor=1.)
z_list.append(enc.cpu())
latents = torch.concat(z_list, dim=0)
scale_factor = 1. / latents.flatten().std().item()
print(f"scale_factor: {scale_factor}")
print("done")
if __name__ == "__main__":
main()
|