scripts/ternary_search.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 torch.distributed as dist
from functools import cache
from mpi4py import MPI

from cm import dist_util, logger
from cm.script_util import (
    NUM_CLASSES,
    model_and_diffusion_defaults,
    create_model_and_diffusion,
    add_dict_to_argparser,
    args_to_dict,
)
from cm.random_util import get_generator
from cm.karras_diffusion import stochastic_iterative_sampler
from evaluations.th_evaluator import FIDAndIS


def main():
    args = create_argparser().parse_args()

    dist_util.setup_dist()
    logger.configure()

    if "consistency" in args.training_mode:
        distillation = True
    else:
        distillation = False

    logger.log("creating model and diffusion...")
    model, diffusion = create_model_and_diffusion(
        **args_to_dict(args, model_and_diffusion_defaults().keys()),
        distillation=distillation,
    )
    model.load_state_dict(
        dist_util.load_state_dict(args.model_path, map_location="cpu")
    )
    model.to(dist_util.dev())
    if args.use_fp16:
        model.convert_to_fp16()
    model.eval()

    fid_is = FIDAndIS()
    fid_is.set_ref_batch(args.ref_batch)
    (
        ref_fid_stats,
        ref_spatial_stats,
        ref_clip_stats,
    ) = fid_is.get_ref_batch(args.ref_batch)

    def sample_generator(ts):
        logger.log("sampling...")
        all_images = []
        all_labels = []
        all_preds = []

        generator = get_generator(args.generator, args.num_samples, args.seed)
        while len(all_images) * args.batch_size < args.num_samples:
            model_kwargs = {}
            if args.class_cond:
                classes = th.randint(
                    low=0,
                    high=NUM_CLASSES,
                    size=(args.batch_size,),
                    device=dist_util.dev(),
                )
                model_kwargs["y"] = classes

            def denoiser(x_t, sigma):
                _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
                if args.clip_denoised:
                    denoised = denoised.clamp(-1, 1)
                return denoised

            x_T = (
                generator.randn(
                    *(args.batch_size, 3, args.image_size, args.image_size),
                    device=dist_util.dev(),
                )
                * args.sigma_max
            )

            sample = stochastic_iterative_sampler(
                denoiser,
                x_T,
                ts,
                t_min=args.sigma_min,
                t_max=args.sigma_max,
                rho=diffusion.rho,
                steps=args.steps,
                generator=generator,
            )
            pred, spatial_pred, clip_pred, text_pred, _ = fid_is.get_preds(sample)

            sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
            sample = sample.permute(0, 2, 3, 1)
            sample = sample.contiguous()

            gathered_samples = [
                th.zeros_like(sample) for _ in range(dist.get_world_size())
            ]
            gathered_preds = [th.zeros_like(pred) for _ in range(dist.get_world_size())]
            dist.all_gather(gathered_samples, sample)  # gather not supported with NCCL
            dist.all_gather(gathered_preds, pred)
            all_images.extend([sample.cpu().numpy() for sample in gathered_samples])
            all_preds.extend([pred.cpu().numpy() for pred in gathered_preds])
            if args.class_cond:
                gathered_labels = [
                    th.zeros_like(classes) for _ in range(dist.get_world_size())
                ]
                dist.all_gather(gathered_labels, classes)
                all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])

            logger.log(f"created {len(all_images) * args.batch_size} samples")

        arr = np.concatenate(all_images, axis=0)
        arr = arr[: args.num_samples]
        preds = np.concatenate(all_preds, axis=0)
        preds = preds[: args.num_samples]
        if args.class_cond:
            label_arr = np.concatenate(all_labels, axis=0)
            label_arr = label_arr[: args.num_samples]

        dist.barrier()
        logger.log("sampling complete")

        return arr, preds

    @cache
    def get_fid(p, begin=(0,), end=(args.steps - 1,)):

        samples, preds = sample_generator(begin + (p,) + end)
        is_root = dist.get_rank() == 0
        if is_root:
            fid_stats = fid_is.get_statistics(preds, -1)
            fid = ref_fid_stats.frechet_distance(fid_stats)
            fid = MPI.COMM_WORLD.bcast(fid)
            # spatial_stats = fid_is.get_statistics(spatial_preds, -1)
            # sfid = ref_spatial_stats.frechet_distance(spatial_stats)
            # clip_stats = fid_is.get_statistics(clip_preds, -1)
            IS = fid_is.get_inception_score(preds)
            IS = MPI.COMM_WORLD.bcast(IS)
            # clip_fid = fid_is.get_clip_score(clip_preds, text_preds)
            # fcd = ref_clip_stats.frechet_distance(clip_stats)
        else:
            fid = MPI.COMM_WORLD.bcast(None)
            IS = MPI.COMM_WORLD.bcast(None)

        dist.barrier()
        return fid, IS

    def ternary_search(before=(0,), after=(17,)):
        left = before[-1]
        right = after[0]
        is_root = dist.get_rank() == 0
        while right - left >= 3:
            m1 = int(left + (right - left) / 3.0)
            m2 = int(right - (right - left) / 3.0)
            f1, is1 = get_fid(m1, before, after)
            if is_root:
                logger.log(f"fid at m1 = {m1} is {f1}, IS is {is1}")
            f2, is2 = get_fid(m2, before, after)
            if is_root:
                logger.log(f"fid at m2 = {m2} is {f2}, IS is {is2}")
            if f1 < f2:
                right = m2
            else:
                left = m1
            if is_root:
                logger.log(f"new interval is [{left}, {right}]")

        if right == left:
            p = right
        elif right - left == 1:
            f1, _ = get_fid(left, before, after)
            f2, _ = get_fid(right, before, after)
            p = m1 if f1 < f2 else m2
        elif right - left == 2:
            mid = left + 1
            f1, _ = get_fid(left, before, after)
            f2, _ = get_fid(right, before, after)
            fmid, ismid = get_fid(mid, before, after)
            if is_root:
                logger.log(f"fmid at mid = {mid} is {fmid}, ISmid is {ismid}")
            if fmid < f1 and fmid < f2:
                p = mid
            elif f1 < f2:
                p = m1
            else:
                p = m2

        return p

    # convert comma separated numbers to tuples
    begin = tuple(int(x) for x in args.begin.split(","))
    end = tuple(int(x) for x in args.end.split(","))

    optimal_p = ternary_search(begin, end)
    if dist.get_rank() == 0:
        logger.log(f"ternary_search_results: {optimal_p}")
        fid, IS = get_fid(optimal_p, begin, end)
        logger.log(f"fid at optimal p = {optimal_p} is {fid}, IS is {IS}")


def create_argparser():
    defaults = dict(
        begin="0",
        end="39",
        training_mode="consistency_distillation",
        generator="determ",
        clip_denoised=True,
        num_samples=10000,
        batch_size=16,
        sampler="heun",
        s_churn=0.0,
        s_tmin=0.0,
        s_tmax=float("inf"),
        s_noise=1.0,
        steps=40,
        model_path="",
        ref_batch="",
        seed=42,
    )
    defaults.update(model_and_diffusion_defaults())
    parser = argparse.ArgumentParser()
    add_dict_to_argparser(parser, defaults)
    return parser


if __name__ == "__main__":
    main()