scripts/sample_rule.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.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

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 = 'cond_demo/'
    if 'cond_table/' in args.config_path:
        args.dir = root_dir + os.path.splitext(args.config_path.split('cond_table/')[-1])[0] + f'_cls_{args.class_label}'
    else:
        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)

    target_rules = vars(config.target_rules)
    source = 'given'
    # see if target rules are given, if not, extract from dataset
    for key, val in target_rules.items():
        if val is None:
            source = 'dataset'
        break

    if source == 'dataset':
        if 'vertical_nd' in target_rules.keys():
            # create a new dummy rule name and delete the old names
            target_rules['note_density'] = None
            target_rules.pop('vertical_nd')
            target_rules.pop('horizontal_nd')
        model_kwargs = {"rule": {k: v for k, v in target_rules.items()}}
        logger.log(f"loading midi from test set cls {args.class_label}...")
        val_data = load_data(
            data_dir=args.data_dir + "_test_cls_" + str(args.class_label) + ".csv",
            batch_size=args.batch_size,
            class_cond=True,
            deterministic=args.record or args.deterministic,   # for record, use the same target
            image_size=gen_shape[2] * 8,
            rule=None,
        )
        with th.no_grad():
            gt, extra = next(val_data)
            gt = gt.to(dist_util.dev())
            for rule_name in target_rules.keys():
                target_rule = _extract_rule(rule_name, gt)
                model_kwargs["rule"][rule_name] = target_rule

    else:
        for key, val in target_rules.items():
            if 'vertical_nd' in key:
                # to make vertical and horizontal nd to be of similar scale
                if '_hr_' in key:
                    str_hr_scale = key.split('_hr_')[-1]
                    horizontal_scale = int(str_hr_scale)
                    horizontal_nd = [x / horizontal_scale for x in target_rules[f'horizontal_nd_hr_{str_hr_scale}']]
                    target_rules[f'note_density_hr_{str_hr_scale}'] = target_rules[key] + horizontal_nd
                else:
                    horizontal_scale = 5
                    horizontal_nd = [x / horizontal_scale for x in target_rules['horizontal_nd']]
                    target_rules['note_density'] = target_rules[key] + horizontal_nd

                target_rules.pop(key)
                target_rules.pop(key.replace('vertical', 'horizontal'))
                break

        for key, val in target_rules.items():
            val = th.tensor(val, device=dist_util.dev())
            if key == 'pitch_hist':
                val = val / (th.sum(val) + 1e-12)
            target_rules[key] = val
        model_kwargs = {"rule": {k: v.repeat(args.batch_size, 1) for k, v in target_rules.items()}}

    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
        model_kwargs["y"] = classes

    save_dir = logger.get_dir()
    os.makedirs(os.path.expanduser(save_dir), 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
    )

    logger.log("sampling...")
    count_samples = 0
    all_results = pd.DataFrame()

    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,
            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)
        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)
            else:
                midi_util.save_piano_roll_midi(arr, save_dir, args.fs, save_ind=count_samples)

        # test distance between generated samples and target
        generated_samples = th.from_numpy(arr) / 63.5 - 1
        results = midi_util.eval_rule_loss(generated_samples, model_kwargs["rule"])
        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)

    if args.record:
        import pickle
        with open(os.path.join(save_dir, 'log_probs.pkl'), 'wb') as f:
            pickle.dump(diffusion.log_probs, f)
        with open(os.path.join(save_dir, 'loss_std.pkl'), 'wb') as f:
            pickle.dump(diffusion.loss_std, f)
        with open(os.path.join(save_dir, 'loss_range.pkl'), 'wb') as f:
            pickle.dump(diffusion.loss_range, f)
        with open(os.path.join(save_dir, 'each_loss.pkl'), 'wb') as f:
            pickle.dump(diffusion.each_loss, f)

        midi_util.plot_record(diffusion.log_probs, 'log_prob', save_dir)
        midi_util.plot_record(diffusion.loss_std, 'loss_std', save_dir)
        midi_util.plot_record(diffusion.loss_range, 'loss_range', save_dir)

        if len(diffusion.inter_piano_rolls) > 0:
            diffusion.inter_piano_rolls.append(th.from_numpy(arr))
            inter_piano_rolls = th.concat(diffusion.inter_piano_rolls, dim=0)
            save_dir_inter = os.path.join(save_dir, 'inter')
            os.makedirs(save_dir_inter, exist_ok=True)
            midi_util.save_piano_roll_midi(inter_piano_rolls.numpy(), save_dir=save_dir_inter, fs=args.fs)

    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
        deterministic=False,  # whether to use the same rule everytime
        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()