sample_acc_vq.py

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
For a simple single-GPU/CPU sampling script, see sample.py.
"""
import random
from einops import rearrange, repeat
from omegaconf import OmegaConf
import torch
from tqdm import tqdm
import os
import numpy as np
import math
import hydra
import shutil
import uuid
from utils.train_utils import get_model, requires_grad
import accelerate
import wandb
from utils_vq import print_rank_0
from utils_vq import (
    vq_get_encoder_decoder,
    vq_get_generator,
    vq_get_dynamic,
    vq_get_vae,
    vq_get_sample_size,
)
from utils.openai_eval import _eval_by_npz

from utils.my_metrics_offline import MyMetric_Offline as MyMetric

from utils_vq import get_dataset_id2label
from torchvision.utils import save_image
from utils_vq import wandb_visual_dict,get_dataloader
from fvd_external import calculate_fvd_github

from torchvision.io import write_video


@hydra.main(config_path="config", config_name="default", version_base=None)
def main(args):
    return _main(args)


def _main(cfg):
    if cfg.debug:
        cfg.data.batch_size = 4
        cfg.ckpt_every = 10
        cfg.data.sample_fid_n = 1_00
        cfg.data.sample_fid_bs = 4
        cfg.data.sample_fid_every = 5
        cfg.data.sample_vis_every = 3
        cfg.data.sample_vis_n = 2
        #####
        cfg.num_fid_samples = 100
        cfg.offline.lbs = 4

    """
    Run sampling.
    """
    torch.backends.cuda.matmul.allow_tf32 = True # True: fast but may lead to some small numerical differences
    assert (
        torch.cuda.is_available()
    ), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage"
    torch.set_grad_enabled(False)
    if cfg.data.num_classes > 0:
        imagenet_id2label = get_dataset_id2label(cfg.data.name)

    slurm_job_id = os.environ.get("SLURM_JOB_ID")
    if slurm_job_id is None:
        slurm_job_id = "local"
    print_rank_0(f"slurm_job_id: {slurm_job_id}")
    skip_data_loader = True  # we didn't save real image in webdataset, so we need to load it from disk offline by preprocess
    if "coco14" in cfg.data.name:
        skip_data_loader = False
        print_rank_0(f"coco dataset,skip_data_loader force to be False: {skip_data_loader}")

    accelerator = accelerate.Accelerator()
    device = accelerator.device
    accelerate.utils.set_seed(cfg.global_seed, device_specific=True)
    rank = accelerator.state.process_index
    print_rank_0(
        f"Starting rank={rank}, world_size={accelerator.state.num_processes}, device={device}."
    )
    if accelerator.is_main_process:
        metric_fid_with_npz = MyMetric(npz_real=cfg.data.npz_real)

    assert cfg.ckpt is not None, "Must specify a checkpoint to sample from"
    model = get_model(cfg)
    model = model.to(device)

    if True:
        state_dict = torch.load(cfg.ckpt, map_location=lambda storage, loc: storage)
        _model_dict = state_dict["ema"]
        _model_dict = {k.replace("module.", ""): v for k, v in _model_dict.items()}
        model.load_state_dict(_model_dict)
        model = model.to(device)
        requires_grad(model, False)
        print_rank_0(f"Loaded checkpoint from {cfg.ckpt}")

    model.eval()  

    local_bs = cfg.offline.lbs
    cfg.data.batch_size = local_bs  # used for generating captions,etc.
    print_rank_0(f"local_bs: {local_bs}")

    model = accelerator.prepare(model)

    # Figure out how many samples we need to generate on each GPU and how many iterations we need to run:
    global_bs = local_bs * accelerator.state.num_processes
    total_samples = int(math.ceil(cfg.num_fid_samples / global_bs) * global_bs)
    assert (
        total_samples % accelerator.state.num_processes == 0
    ), "total_samples must be divisible by world_size"
    samples_needed_this_gpu = int(total_samples // accelerator.state.num_processes)
    assert (
        samples_needed_this_gpu % local_bs == 0
    ), "samples_needed_this_gpu must be divisible by the per-GPU batch size"
    iterations = int(samples_needed_this_gpu // local_bs)

    print_rank_0(
        f"Total number of images that will be sampled: {total_samples} with global_batch_size={global_bs}"
    )

    tokenizer_encode_fn, tokenizer_decode_fn = vq_get_encoder_decoder(
        cfg=cfg, device=device
    )
    training_losses_fn, sample_fn = vq_get_dynamic(
        cfg=cfg, device=device
    )
    vae = vq_get_vae(cfg, device)

    if not skip_data_loader:
        loader = get_dataloader(cfg)
        loader = accelerator.prepare(loader)
        data_generator, _generator, caption_generator = vq_get_generator(
            cfg, device, loader, accelerator.state.process_index, 0, vae
        )

    assert cfg.cfg_scale >= 0.0 # "In almost all cases, cfg_scale be >= 1.0"
    if cfg.cfg_scale > 0.0:
        model_fn = model.forward_with_cfg
    elif cfg.cfg_scale == 0.0:
        model_fn = model.forward_without_cfg
    else:
        raise ValueError(f"cfg_scale={cfg.cfg_scale} is not supported")



    wandb_name = "_".join(
        [
            "v1", 
            cfg.data.name,
            cfg.model.name,
            f"sampler={cfg.dynamic.disint.sampler}",
            f"scheduler={cfg.dynamic.disint.scheduler}",
            f"bs{cfg.offline.lbs}fid{cfg.num_fid_samples}cfg{cfg.cfg_scale}softmaxtem{cfg.sm_t}topk{cfg.top_k}topp{cfg.top_p}"
            f"dstep{cfg.dstep_num}",
            f"maxlast{int(cfg.offline.max_last)}",
            f"{slurm_job_id}",
        ]
    )
    sample_sample_dict_4wandb = dict(
        cfg_scale=cfg.cfg_scale,
        sm_t=cfg.sm_t,
        top_p=cfg.top_p,
        top_k=cfg.top_k,
        dstep_num=cfg.dstep_num,
        max_last=int(cfg.offline.max_last),
        fid_samples=cfg.num_fid_samples,
        offline_sample_local_bs=cfg.offline.lbs,
    )

    sample_folder_dir = f"{cfg.sample_dir}/{wandb_name}"
    if rank == 0:
        os.makedirs(sample_folder_dir, exist_ok=True)
        if cfg.use_wandb:
            entity = cfg.wandb.entity
            project = cfg.wandb.project + "_vis"
            print_rank_0(
                f"Logging to wandb entity={entity}, project={project},rank={rank}"
            )
            config_dict = OmegaConf.to_container(cfg, resolve=True)
            wandb.init(
                project=project,
                name=wandb_name,
                config=config_dict,
                dir=sample_folder_dir,
                resume="allow",
                mode="online",
            )

            wandb_project_url = (
                f"https://wandb.ai/dpose-team/{wandb.run.project}/runs/{wandb.run.id}"
            )
            wandb_sync_command = (
                f"wandb sync {sample_folder_dir}/wandb/latest-run --append"
            )
            wandb_desc = "\n".join(
                [
                    "*" * 24,
                    # str(config_dict),
                    wandb_name,
                    wandb_project_url,
                    wandb_sync_command,
                    "*" * 24,
                ]
            )
        else:
            wandb_project_url = "wandb_project_url_null"
            wandb_sync_command = "wandb_sync_command_null"
            wandb_desc = "wandb_desc_null"

        print_rank_0(f"Saving .png samples at {sample_folder_dir}")

        sample_img_dir = f"{sample_folder_dir}/samples"
        gt_img_dir = f"{sample_folder_dir}/gts"
        shutil.rmtree(sample_img_dir, ignore_errors=True)
        shutil.rmtree(gt_img_dir, ignore_errors=True)
        os.makedirs(sample_img_dir, exist_ok=True)
        os.makedirs(gt_img_dir, exist_ok=True)
        

    is_video = cfg.data.video_frames > 0

    accelerator.wait_for_everyone()
    pbar = range(iterations)
    pbar = tqdm(pbar, total=iterations, desc="sampling") if rank == 0 else pbar

    if rank == 0:
        print(wandb_desc)

    for bs_index in pbar:
        if rank == 0:
            print_rank_0(
                f"dataset.subset: {cfg.data.subset},skip_data_loader: {skip_data_loader},sample_img_dir: {sample_img_dir},gt_img_dir: {gt_img_dir}", 
        )
            print(wandb_desc)
        if cfg.data.num_classes > 0:
            y = torch.randint(0, cfg.data.num_classes, (local_bs,), device=device)
        elif cfg.data.num_classes == -666:  # a special value for caption generation
            cap_feat, cap_str = next(caption_generator)
            y = cap_feat[:local_bs]
            assert len(cap_str) == local_bs
        else:
            y = None

        model_kwargs = dict(
            y=y,
            temperature=cfg.sm_t,
            max_last=cfg.offline.max_last,
        )
        if cfg.cfg_scale != 0.0:
            model_kwargs.update(dict(cfg_scale=cfg.cfg_scale))

        if not skip_data_loader:
            gts = next(_generator)
            _sample_size = vq_get_sample_size(len(gts), cfg)

            with torch.no_grad():
                indices_chains = sample_fn(_sample_size, model_fn, **model_kwargs)
                _indices = indices_chains[-1]
                samples = tokenizer_decode_fn(_indices)

            # gts = gts[: len(samples)]
            gts = samples.clone()
        else:
            _sample_size = vq_get_sample_size(local_bs, cfg)

            with torch.no_grad():
                indices_chains = sample_fn(_sample_size, model_fn, **model_kwargs)
                _indices = indices_chains[-1]
                samples = tokenizer_decode_fn(_indices)

            gts = samples.clone()

        sam_4fid, gts_4fid = samples, gts
        gts_4fid = accelerator.gather(gts_4fid.to(device))
        sam_4fid = accelerator.gather(sam_4fid.to(device))
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            metric_fid_with_npz.update_fake(sam_4fid)

        if rank == 0:
            if cfg.use_wandb and bs_index <= 1:
                if cfg.data.num_classes > 0:
                    captions_sample = [
                        imagenet_id2label[y[_].item()]
                        for _ in range(min(16, len(sam_4fid)))
                    ]
                elif cfg.data.num_classes == -666:
                    captions_sample = [
                        cap_str[_] for _ in range(min(16, len(sam_4fid)))
                    ]
                else:
                    captions_sample = [
                        "null caption" for _ in range(min(16, len(sam_4fid)))
                    ]
                wandb_dict = {}
                wandb_dict.update(
                    wandb_visual_dict(
                        "vis/samples_single",
                        sam_4fid,
                        is_video=is_video,
                        num=16,
                        captions=captions_sample,
                    )
                )
                wandb_dict.update(
                    wandb_visual_dict(
                        "vis/gts_single",
                        gts_4fid,
                        is_video=is_video,
                        num=16,
                        captions=None,
                    )
                )
                wandb_dict.update(sample_sample_dict_4wandb)
                wandb.log(
                    wandb_dict,
                    step=bs_index,
                )
                print_rank_0("log_image into wandb")
                if not is_video:
                    wandb.log(
                        {
                            f"vis/samples": wandb.Image(sam_4fid[:16]),
                            f"vis/gts": wandb.Image(gts_4fid[:16]),
                        },
                        step=bs_index,
                    )
                print_rank_0("log_image into wandb")
        accelerator.wait_for_everyone()

        if rank == 0:
            print_rank_0(f"saving sample images, in {sample_img_dir}")
            if not is_video:
                # Save samples to disk as individual .png files
                for _iii, sample in enumerate(sam_4fid):
                    unique_id = uuid.uuid4().hex[:6]
                    save_image(sample / 255.0, f"{sample_img_dir}/{unique_id}.png")

                for _iii, sample in enumerate(gts_4fid):
                    unique_id = uuid.uuid4().hex[:6]
                    save_image(sample / 255.0, f"{gt_img_dir}/{unique_id}.png")

            else:
                for _video_id, (samples, gts) in enumerate(zip(sam_4fid, gts_4fid)):
                    samples = rearrange(samples, "t c h w -> t h w c")  # [0,255]
                    gts = rearrange(gts, "t c h w -> t h w c")  # [0,255]
                    samples = samples.cpu().numpy()
                    gts = gts.cpu().numpy()
                    unique_id = uuid.uuid4().hex[:6]
                    write_video(
                        os.path.join(sample_img_dir, f"{unique_id}.mp4"),
                        samples,
                        fps=cfg.data.fps,  
                    )
                    write_video(
                        os.path.join(gt_img_dir, f"{unique_id}.mp4"),
                        gts,
                        fps=cfg.data.fps,  
                    )


    if rank == 0:
        if not is_video:
            mymetric_dict = {}
            

            _metric_dict = metric_fid_with_npz.compute()
            fid_with_npz = _metric_dict["fid"]
            print_rank_0(f"fid_with_npz: {fid_with_npz}")
            mymetric_dict.update(fid_with_npz=fid_with_npz)
            
            wandb.log(mymetric_dict)

        else:
            fake_root = os.path.expanduser(sample_img_dir)
            real_root = os.path.expanduser(cfg.data.fvd_real_video_dir)
            # real_root = os.path.expanduser(gt_img_dir)
            print_rank_0(real_root)
            print_rank_0(fake_root)

            mymetric_dict = calculate_fvd_github(
                gen_dir=fake_root,
                gt_dir=real_root,
                frames=16,
                resolution=128,  # 128,64
            )

            mymetric_dict = {f"video/{k}": v for k, v in mymetric_dict.items()}
            wandb.log(mymetric_dict)
        try:
            shutil.rmtree(sample_img_dir)
            shutil.rmtree(gt_img_dir)
            print_rank_0(
                f"removed sample_img_dir and gt_img_dir\n,sample_img_dir: {sample_img_dir}\n,gt_img_dir: {gt_img_dir}"
            )
        except Exception as e:
            print_rank_0(f"Error removing directory {sample_img_dir},{gt_img_dir}: {e}")

    print_rank_0("done sampling")
    # accelerator.wait_for_everyone(), important! remove this.


if __name__ == "__main__":

    main()