scripts/edit.py

"""
Generate a large batch of image samples from a model and save them as a large
numpy array. This can be used to produce samples for FID evaluation.
"""

import argparse
import os

import numpy as np
import torch as th
import pandas as pd
import torch.distributed as dist
import torch.nn.functional as F
import multiprocessing

from guided_diffusion import dist_util, midi_util, logger
from guided_diffusion.midi_util import (
    VERTICAL_ND_BOUNDS, VERTICAL_ND_CENTER, HORIZONTAL_ND_BOUNDS, HORIZONTAL_ND_CENTER,
    get_full_piano_roll
)
from guided_diffusion.dit import DiT_models
from guided_diffusion.script_util import (
    NUM_CLASSES,
    model_and_diffusion_defaults,
    create_diffusion,
    create_model_and_diffusion,
    add_dict_to_argparser,
    args_to_dict,
)
from guided_diffusion.gaussian_diffusion import _encode, _extract_rule
from guided_diffusion.pr_datasets_all import load_data
from load_utils import load_model
import diff_collage as dc
from guided_diffusion.condition_functions import (
    model_fn, dc_model_fn, composite_nn_zt, composite_rule)
from functools import partial
import matplotlib.pyplot as plt
import pretty_midi

plt.rcParams["figure.figsize"] = (20, 3)
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300


def main():
    args = create_argparser().parse_args()
    root_dir = 'edit_demo/'
    args.dir = root_dir + os.path.splitext(args.config_path.split(root_dir)[-1])[0] + f'_cls_{args.class_label}'

    comm = dist_util.setup_dist(port=args.port)
    logger.configure(args=args, comm=comm)
    config = midi_util.load_config(args.config_path)
    if config.sampling.use_ddim:
        args.timestep_respacing = config.sampling.timestep_respacing

    logger.log("creating model and diffusion...")
    model = DiT_models[args.model](
        input_size=args.image_size,
        in_channels=args.in_channels,
        num_classes=args.num_classes,
        learn_sigma=args.learn_sigma,
    )
    diffusion = create_diffusion(
        learn_sigma=args.learn_sigma,
        diffusion_steps=args.diffusion_steps,
        noise_schedule=args.noise_schedule,
        timestep_respacing=args.timestep_respacing,
        use_kl=args.use_kl,
        predict_xstart=args.predict_xstart,
        rescale_timesteps=args.rescale_timesteps,
        rescale_learned_sigmas=args.rescale_learned_sigmas,
    )
    model.load_state_dict(
        dist_util.load_state_dict(args.model_path, map_location="cpu"), strict=False
    )
    model.to(dist_util.dev())
    if args.use_fp16:
        model.convert_to_fp16()
    model.eval()

    # create embed model
    if args.vae is not None:
        embed_model = load_model(args.vae, args.vae_path)
        embed_model.to(dist_util.dev())
        embed_model.eval()
    else:
        embed_model = None

    cond_fn_config = config.guidance.cond_fn
    if config.guidance.nn:
        logger.log("loading classifier...")
        classifier_config = cond_fn_config.classifiers
        num_classifiers = len(classifier_config.names)
        classifiers = []
        for i in range(num_classifiers):
            classifier = DiT_models[classifier_config.names[i]](
                # classifier trained on latents, so has the same img size as diffusion
                input_size=args.image_size,
                in_channels=args.in_channels,
                num_classes=classifier_config.num_classes[i],
            )
            classifier.load_state_dict(
                dist_util.load_state_dict(classifier_config.paths[i], map_location="cpu")
            )
            classifier.to(dist_util.dev())
            classifier.eval()
            classifiers.append(classifier)

    if cond_fn_config is not None:
        if config.guidance.nn:
            cond_fn_used = partial(composite_nn_zt, fns=cond_fn_config.fns,
                                   classifier_scales=cond_fn_config.classifier_scales,
                                   classifiers=classifiers, rule_names=cond_fn_config.rule_names)
        else:
            cond_fn_used = partial(composite_rule, fns=cond_fn_config.fns,
                                   classifier_scales=cond_fn_config.classifier_scales,
                                   rule_names=cond_fn_config.rule_names)
    else:
        cond_fn_used = None

    if config.sampling.diff_collage:
        def eps_fn(x, t, y=None):
            # since our backbone takes 128x16 as input
            return model(x.permute(0, 1, 3, 2), t, y=y).permute(0, 1, 3, 2)

        # circle need one more num_img than linear
        img_shape = (args.in_channels, args.image_size[1], args.image_size[0])  # 4 x 16 x 128
        if config.dc.type == 'circle':
            worker = dc.CondIndCircle(img_shape, eps_fn, config.dc.num_img + 1, overlap_size=config.dc.overlap_size)
        else:
            worker = dc.CondIndSimple(img_shape, eps_fn, config.dc.num_img, overlap_size=config.dc.overlap_size)
        model_long_fn = worker.eps_scalar_t_fn
        gen_shape = (args.batch_size, worker.shape[0], worker.shape[2], worker.shape[1])
        model_fn_used = partial(dc_model_fn, model=model_long_fn, num_classes=args.num_classes,
                                class_cond=args.class_cond, cfg=args.cfg, w=args.w)
    else:
        gen_shape = (args.batch_size, args.in_channels, args.image_size[0], args.image_size[1])
        model_fn_used = partial(model_fn, model=model, num_classes=args.num_classes,
                                class_cond=args.class_cond, cfg=args.cfg, w=args.w)

    if args.class_cond:
        # only generate one class
        classes = th.ones(size=(args.batch_size,), device=dist_util.dev(), dtype=th.int) * args.class_label

    save_dir = logger.get_dir()
    save_dir_gt = os.path.join(save_dir, 'gt')
    os.makedirs(os.path.expanduser(save_dir), exist_ok=True)
    os.makedirs(os.path.expanduser(save_dir_gt), exist_ok=True)

    ddim_stochastic = partial(diffusion.ddim_sample_loop, eta=1.)
    sample_fn = (
        diffusion.p_sample_loop if not config.sampling.use_ddim else ddim_stochastic
    )

    edit_kwargs = vars(config.edit)
    edit_kwargs["l_start_pix"] = edit_kwargs["l_start"] * 8
    edit_kwargs["l_end_pix"] = edit_kwargs["l_end"] * 8
    source = getattr(config.edit, 'source', None)
    if source == 'dataset':
        logger.log(f"loading midi from test set cls {args.class_label} to edit...")
        val_data = load_data(
            data_dir=args.data_dir + "_test_cls_" + str(args.class_label) + ".csv",
            batch_size=args.batch_size,
            class_cond=True,
            image_size=gen_shape[2] * 8,
            rule=None,
        )
        gt, extra = next(val_data)
        gt = gt.to(dist_util.dev())
    else:
        midi_data = pretty_midi.PrettyMIDI(source)
        gt = get_full_piano_roll(midi_data, fs=args.fs)
        gt = th.from_numpy(gt).float()[None] / 63.5 - 1
        gt = F.pad(gt, (0, gen_shape[2] * 8 - gt.shape[3]), "constant", -1)
        gt = gt.to(dist_util.dev())

    gt_latent = _encode(gt, embed_model, scale_factor=args.scale_factor)
    mask = th.ones_like(gt_latent)
    mask[:, :, edit_kwargs["l_start"]:edit_kwargs["l_end"], :] = 0.
    edit_kwargs["gt"] = gt_latent
    edit_kwargs["mask"] = mask

    logger.log("sampling...")

    with th.no_grad():
        model_kwargs = {"rule": {}}
        target_rules = vars(config.target_rules)
        gt_partial = gt[:, :, :, edit_kwargs["l_start"]*8:edit_kwargs["l_end"]*8]
        for rule_name, val in target_rules.items():
            if 'horizontal' in rule_name:
                continue
            # generate a different target for nd, generate the same target for chord
            elif 'vertical' in rule_name:
                hr_nd = target_rules[rule_name.replace('vertical', 'horizontal')]
                if '_hr_' in rule_name:
                    str_hr_scale = rule_name.split('_hr_')[-1]
                    horizontal_scale = int(str_hr_scale)
                    rule_name = f'note_density_hr_{str_hr_scale}'
                else:
                    horizontal_scale = 5
                    rule_name = 'note_density'
                # need orig_rule for all cases because want to record orig_rule
                orig_rule = _extract_rule(rule_name, gt_partial)
                if len(orig_rule.shape) == 1:
                    # unsqueeze the first dimension of batch_size = 1
                    orig_rule = orig_rule.reshape(1, -1)

                # if not given target or target is to shift extracted nd
                if isinstance(val, int) or val is None:
                    # need to compute class to shift
                    vt_bounds = th.tensor(VERTICAL_ND_BOUNDS).to(dist_util.dev())
                    hr_bounds = th.tensor(HORIZONTAL_ND_BOUNDS).to(dist_util.dev()) / horizontal_scale
                    vt_center = th.tensor(VERTICAL_ND_CENTER).to(dist_util.dev())
                    hr_center = th.tensor(HORIZONTAL_ND_CENTER).to(dist_util.dev()) / horizontal_scale
                    if isinstance(val, int):
                        vertical_rand = val
                        horizontal_rand = hr_nd
                    else:
                        # randomly shift nd
                        vertical_range = 1
                        horizontal_range = 1
                        vertical_rand = th.randint(-vertical_range, vertical_range + 1,
                                                   size=(orig_rule.shape[0], 1), device=orig_rule.device)
                        horizontal_rand = th.randint(-horizontal_range, horizontal_range + 1,
                                                     size=(orig_rule.shape[0], 1), device=orig_rule.device)
                    total_length = orig_rule.shape[-1]
                    vt_nd_classes = (th.bucketize(orig_rule[:, :total_length // 2], vt_bounds) + vertical_rand)
                    hr_nd_classes = (th.bucketize(orig_rule[:, total_length // 2:], hr_bounds) + horizontal_rand)
                    vt_nd_val = vt_center[vt_nd_classes.clamp_(min=0, max=7)]
                    hr_nd_val = hr_center[hr_nd_classes.clamp_(min=0, max=7)]
                    target_rule = th.concat((vt_nd_val, hr_nd_val), dim=-1)
                else:
                    # use given nd
                    hr_nd_rescale = [x / horizontal_scale for x in hr_nd]
                    nd_val = val + hr_nd_rescale
                    target_rule = th.tensor(nd_val, device=dist_util.dev())
            elif 'pitch' in rule_name and val is not None:
                orig_rule = _extract_rule(rule_name, gt_partial)
                val = th.tensor(val, device=dist_util.dev())
                target_rule = val / (th.sum(val) + 1e-12)
            else:
                orig_rule = _extract_rule(rule_name, gt_partial)
                if val is not None:
                    target_rule = th.tensor(val, device=dist_util.dev())
                else:
                    target_rule = _extract_rule(rule_name, gt_partial)
            if source == 'dataset':
                if len(target_rule.shape) == 1:
                    target_rule = target_rule.reshape(1, -1).repeat(args.batch_size, 1)
                model_kwargs["rule"][rule_name] = target_rule
            else:
                # if given only one source, generate multiple variations
                model_kwargs["rule"][rule_name] = target_rule.repeat(args.batch_size, 1)

    if args.class_cond:
        model_kwargs["y"] = classes

    all_results = pd.DataFrame()

    count_samples = 0
    while count_samples < args.num_samples:
        sample = sample_fn(
            model_fn_used,
            gen_shape,
            clip_denoised=args.clip_denoised,
            model_kwargs=model_kwargs,
            device=dist_util.dev(),
            cond_fn=cond_fn_used,
            # None for NN(z_0), embed_model for rule(decoder(z_0))
            embed_model=embed_model if config.guidance.vae else None,
            scale_factor=args.scale_factor,
            guidance_kwargs=config.guidance,
            scg_kwargs=vars(config.scg) if config.guidance.scg else None,
            edit_kwargs=edit_kwargs,
            t_end=config.sampling.t_end,
            record=args.record,
            progress=True
        )
        sample = midi_util.decode_sample_for_midi(sample, embed_model=embed_model,
                                                  scale_factor=args.scale_factor, threshold=-0.95)
        arr = sample.cpu().numpy()
        arr = arr.transpose(0, 3, 1, 2)
        gt = ((gt + 1) * 63.5).clamp(0, 127).to(th.uint8)
        arr_gt = gt.cpu().numpy()
        if args.save_files:
            if args.class_cond:
                label_arr = classes.cpu().numpy()
                midi_util.save_piano_roll_midi(arr, save_dir, args.fs, y=label_arr, save_ind=count_samples)
                midi_util.save_piano_roll_midi(arr_gt, save_dir_gt, args.fs, y=label_arr, save_ind=count_samples)
            else:
                midi_util.save_piano_roll_midi(arr, save_dir, args.fs, save_ind=count_samples)
                midi_util.save_piano_roll_midi(arr_gt, save_dir_gt, args.fs, save_ind=count_samples)

        # test distance between generated samples and target
        generated_samples = th.from_numpy(arr) / 63.5 - 1
        # only take editable part to compute rule loss
        generated_samples = generated_samples[:, :, :, edit_kwargs["l_start_pix"]:edit_kwargs["l_end_pix"]]
        results = midi_util.eval_rule_loss(generated_samples, model_kwargs["rule"])
        # save original rules
        orig_rule_dict = {}
        for rule_name in model_kwargs["rule"].keys():
            orig_rule_dict[rule_name + '.orig_rule'] = orig_rule.cpu().tolist()
        orig_rule_df = pd.DataFrame(orig_rule_dict)
        results = pd.concat([orig_rule_df, results], axis=1)
        all_results = pd.concat([all_results, results], ignore_index=True)
        # save every step
        if args.save_files:
            all_results.to_csv(os.path.join(save_dir, 'results.csv'), index=False)
        count_samples += args.batch_size

    if args.save_files:
        all_results.to_csv(os.path.join(save_dir, 'results.csv'), index=False)
        # Create the DataFrame for loss_stats
        loss_columns = [col for col in all_results.columns if '.loss' in col]
        rows = []
        for col in loss_columns:
            rows.append({'Attr': col, 'Mean': all_results[col].mean(), 'Std': all_results[col].std()})
        loss_stats = pd.DataFrame(rows, columns=['Attr', 'Mean', 'Std'])
        loss_stats.to_csv(os.path.join(save_dir, 'summary.csv'))
        print(loss_stats)

    logger.log("sampling complete")


def create_argparser():
    defaults = dict(
        project="music-sampling",
        dir="",
        data_dir="",   # use to load in val data to extract rule
        config_path="",
        model="DiTRotary_XL_8",  # DiT model names
        model_path="",
        vae="kl/f8-all-onset",
        vae_path="taming-transformers/checkpoints/all_onset/epoch_14.ckpt",
        clip_denoised=False,
        num_samples=128,
        batch_size=16,
        scale_factor=1.,
        fs=100,
        num_classes=0,
        class_label=1,  # class to generate
        cfg=False,
        w=4.,  # for cfg
        classifier_scale=1.0,
        record=False,
        save_files=True,
        training=False,  # not training, so don't need to create more folders than needed
        port=None,
    )
    defaults.update(model_and_diffusion_defaults())
    parser = argparse.ArgumentParser()
    add_dict_to_argparser(parser, defaults)
    return parser


if __name__ == "__main__":
    multiprocessing.set_start_method('spawn', force=True)
    main()