code/fewshot/train.py

import os as _os, sys as _sys
_sys.path.append(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__))))

from helpers.improved_precision_recall import compute_prec_recall
from visual.utils import (generate_and_save, generate_for_NN,
                          generate_images_initial,
                          get_sample_for_visualization)
from visual.spatial_visual import spatial_vissual
from visual.nn_interplate import nn_interp
from visual.interpolate import random_interp
from visual.generate_sample_nn import generate_sample_nn
from visual.generate_rnd_nn import generate_rnd_nn
from visual.generate_rnd import generate_rnd
from sampler import Sampler
from metrics.ppl_uniform import calc_ppl_uniform
from metrics.ppl import calc_ppl
from helpers.utils import ZippedDataset, get_cpu_stats_over_ranks
from helpers.train_helpers_fewshot import (load_imle, load_opt, save_latents,
                                           save_latents_latest, save_model,
                                           save_snoise, set_up_hyperparams, update_ema)
from helpers.imle_helpers import backtrack, reconstruct
from data import set_up_data
from torch.utils.data import DataLoader, TensorDataset
import csv
import json
import os
import signal
import shutil
import time
from pathlib import Path

try:
    from comet_ml import Experiment, ExistingExperiment
except ImportError:
    Experiment = None
    ExistingExperiment = None
import imageio
import torch
import torch.nn.functional as F
import torchvision
import wandb
import torch.nn as nn
from cleanfid import fid
import cleanfid.features as _cleanfid_feat
import cleanfid.inception_torchscript as _cleanfid_incept

_inception_cache = os.path.join(os.path.expanduser("~"), ".cache", "cleanfid")
os.makedirs(_inception_cache, exist_ok=True)
_orig_feature_extractor = _cleanfid_feat.feature_extractor


def _patched_feature_extractor(name="torchscript_inception", device=torch.device("cuda"),
                               resize_inside=False, use_dataparallel=True):
    if name == "torchscript_inception":
        model = _cleanfid_incept.InceptionV3W(
            _inception_cache, download=True, resize_inside=resize_inside).to(device)
        model.eval()
        if use_dataparallel:
            model = torch.nn.DataParallel(model)
        return lambda x: model(x)
    return _orig_feature_extractor(name, device, resize_inside, use_dataparallel)


_cleanfid_feat.feature_extractor = _patched_feature_extractor


# ---------- metrics CSV / best_metrics helpers ----------

def append_metrics_csv(save_dir, row: dict):
    csv_path = os.path.join(save_dir, "metrics.csv")
    file_exists = os.path.isfile(csv_path)
    with open(csv_path, "a", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(row.keys()))
        if not file_exists:
            writer.writeheader()
        writer.writerow(row)


def update_best_metrics(save_dir, row: dict):
    best_path = os.path.join(save_dir, "best_metrics.json")
    if os.path.isfile(best_path):
        with open(best_path, "r") as f:
            best = json.load(f)
    else:
        best = {}
    updated = False
    fid_val = row.get("fid")
    if fid_val is not None and (best.get("best_fid") is None or fid_val < best["best_fid"]):
        best["best_fid"] = fid_val
        best["best_fid_epoch"] = row.get("epoch")
        updated = True
    prec_val = row.get("precision")
    if prec_val is not None and (best.get("best_precision") is None or prec_val > best["best_precision"]):
        best["best_precision"] = prec_val
        best["best_precision_epoch"] = row.get("epoch")
    rec_val = row.get("recall")
    if rec_val is not None and (best.get("best_recall") is None or rec_val > best["best_recall"]):
        best["best_recall"] = rec_val
        best["best_recall_epoch"] = row.get("epoch")
    with open(best_path, "w") as f:
        json.dump(best, f, indent=2)
    return updated


def save_images_to_dir(images_nchw: torch.Tensor, out_dir: Path):
    out_dir.mkdir(parents=True, exist_ok=True)
    images_nchw = images_nchw.clamp(0.0, 1.0)
    for i in range(images_nchw.shape[0]):
        img = (images_nchw[i].permute(
            1, 2, 0).cpu().numpy() * 255).astype("uint8")
        imageio.imwrite(str(out_dir / f"{i}.png"), img)


def save_grid_image(grid_chw: torch.Tensor, out_path: str):
    # Save visualization grids directly under checkpoint save_dir.
    grid_img = torchvision.transforms.functional.to_pil_image(
        grid_chw.detach().cpu().clamp(0.0, 1.0))
    grid_img.save(out_path)


def slerp(a: torch.Tensor, b: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    a = F.normalize(a, dim=-1)
    b = F.normalize(b, dim=-1)
    dot = torch.sum(a * b, dim=-1, keepdim=True).clamp(-1.0, 1.0)
    omega = torch.acos(dot)
    sin_omega = torch.sin(omega)
    t = t.view(-1, 1)
    factor1 = torch.sin((1.0 - t) * omega) / sin_omega
    factor2 = torch.sin(t * omega) / sin_omega
    return factor1 * a + factor2 * b


def training_step_imle(H, n, targets, latents, snoise, imle, ema_imle, optimizer, loss_fn):
    t0 = time.time()
    imle.zero_grad()

    cur_batch_latents = latents

    px_z = imle(cur_batch_latents, snoise)
    loss = loss_fn(px_z, targets.permute(0, 3, 1, 2))
    loss.backward()
    optimizer.step()
    if ema_imle is not None:
        update_ema(imle, ema_imle, H.ema_rate)

    stats = get_cpu_stats_over_ranks(dict(loss_nans=0, loss=loss))
    stats.update(skipped_updates=0, iter_time=time.time() - t0, grad_norm=0)
    return stats


def train_loop_imle(H, data_train, data_valid, preprocess_fn, imle, ema_imle, logprint, experiment=None):
    subset_len = len(data_train)
    if H.subset_len != -1:
        subset_len = H.subset_len
    for data_train in DataLoader(data_train, batch_size=subset_len):
        data_train = TensorDataset(data_train[0])
        break

    optimizer, scheduler, _, iterate, starting_epoch = load_opt(
        H, imle, logprint)

    print("Starting epoch: ", starting_epoch)
    print("Starting iteration: ", iterate)

    stats = []
    H.ema_rate = torch.as_tensor(H.ema_rate)

    subset_len = H.subset_len
    if subset_len == -1:
        subset_len = len(data_train)

    sampler = Sampler(H, subset_len, preprocess_fn)

    last_updated = torch.zeros(subset_len, dtype=torch.int16).cuda()
    times_updated = torch.zeros(subset_len, dtype=torch.int8).cuda()
    change_thresholds = torch.empty(subset_len).cuda()
    change_thresholds[:] = H.change_threshold
    best_fid = 100000
    epoch = starting_epoch - 1

    _sigterm_received = [False]

    def _save_full_checkpoint(split_ind, sampler, imle, ema_imle, optimizer,
                              scheduler, change_thresholds, last_updated,
                              times_updated):
        fp = os.path.join(H.save_dir, 'latest')
        logprint(f'SIGTERM checkpoint save @ epoch {epoch} iter {iterate} to {fp}')
        save_model(fp, imle, ema_imle, optimizer, scheduler, H)
        save_latents_latest(H, split_ind, sampler.selected_latents)
        save_latents_latest(H, split_ind, change_thresholds, name='threshold_latest')
        save_latents_latest(H, split_ind, last_updated, name='last_updated')
        save_latents_latest(H, split_ind, times_updated, name='times_updated')

    def _sigterm_handler(signum, frame):
        print(f'SIGTERM received at epoch {epoch}, saving checkpoint...')
        _sigterm_received[0] = True

    def _safe_restore_tensor(path, label):
        if not path or not os.path.isfile(str(path)):
            return None
        try:
            return torch.load(path, map_location='cpu')
        except Exception as e:
            logprint(f'WARNING: Failed to restore {label} from {path}: {type(e).__name__}: {e}. Ignoring this restore artifact.')
            return None

    prev_handler = signal.signal(signal.SIGTERM, _sigterm_handler)

    for split_ind, split_x_tensor in enumerate(DataLoader(data_train, batch_size=subset_len, pin_memory=True)):
        split_x_tensor = split_x_tensor[0].contiguous()
        split_x = TensorDataset(split_x_tensor)
        sampler.init_projection(split_x_tensor)
        viz_batch_original, _ = get_sample_for_visualization(
            split_x, preprocess_fn, H.num_images_visualize, H.dataset)

        print('Outer batch - {}'.format(split_ind, len(split_x)))

        while (epoch < H.num_epochs):

            if _sigterm_received[0]:
                _save_full_checkpoint(
                    split_ind, sampler, imle, ema_imle, optimizer,
                    scheduler, change_thresholds, last_updated, times_updated)
                signal.signal(signal.SIGTERM, prev_handler)
                os.kill(os.getpid(), signal.SIGTERM)
                return

            epoch += 1
            last_updated[:] = last_updated + 1

            restored_latents = False

            if (epoch == starting_epoch):
                latents = _safe_restore_tensor(H.restore_latent_path, 'latents')
                if latents is not None:
                    sampler.selected_latents[:] = latents.to(sampler.selected_latents.device)[:]
                    restored_latents = True
                    print(f'loaded latest latents (shape={latents.shape})')

                threshold = _safe_restore_tensor(H.restore_threshold_path, 'thresholds')
                if threshold is not None:
                    change_thresholds[:] = threshold.to(change_thresholds.device)[:]
                    print('loaded thresholds', torch.mean(change_thresholds))

                lu_path = getattr(H, 'restore_last_updated_path', None)
                restored_last_updated = _safe_restore_tensor(lu_path, 'last_updated')
                if restored_last_updated is not None:
                    last_updated[:] = restored_last_updated.to(last_updated.device)
                    last_updated[:] = last_updated + 1
                    print('loaded last_updated', torch.mean(last_updated.float()))

                tu_path = getattr(H, 'restore_times_updated_path', None)
                restored_times_updated = _safe_restore_tensor(tu_path, 'times_updated')
                if restored_times_updated is not None:
                    times_updated[:] = restored_times_updated.to(times_updated.device)
                    print('loaded times_updated', torch.mean(times_updated.float()))

            sampler.selected_dists[:] = sampler.calc_dists_existing(
                split_x_tensor, imle, dists=sampler.selected_dists)
            dists_in_threshold = sampler.selected_dists < change_thresholds
            updated_enough = last_updated >= H.imle_staleness
            updated_too_much = last_updated >= H.imle_force_resample
            in_threshold = torch.logical_and(
                dists_in_threshold, updated_enough)

            if (H.use_adaptive):
                all_conditions = torch.logical_or(
                    in_threshold, updated_too_much)
            else:
                all_conditions = updated_too_much

            to_update = torch.nonzero(
                all_conditions, as_tuple=False).squeeze(1)

            if (epoch == starting_epoch):
                if not restored_latents:
                    to_update = sampler.entire_ds
                else:
                    for x in DataLoader(split_x, batch_size=H.num_images_visualize, pin_memory=True):
                        break
                    batch_slice = slice(0, x[0].size()[0])
                    latents_viz = sampler.selected_latents[batch_slice]
                    with torch.no_grad():
                        snoise = [s[batch_slice]
                                  for s in sampler.selected_snoise]
                        generate_for_NN(sampler, x[0], latents_viz, snoise, viz_batch_original.shape, imle,
                                        f'{H.save_dir}/NN-samples_{epoch}-{split_ind}-imle.png', logprint)

            change_thresholds[to_update] = sampler.selected_dists[to_update].clone(
            ) * (1 - H.change_coef)

            sampler.imle_sample_force(split_x_tensor, imle, to_update)

            to_update = to_update.cpu()
            last_updated[to_update] = 0
            times_updated[to_update] = times_updated[to_update] + 1

            save_latents_latest(H, split_ind, sampler.selected_latents)
            save_latents_latest(
                H, split_ind, change_thresholds, name='threshold_latest')
            save_latents_latest(H, split_ind, last_updated, name='last_updated')
            save_latents_latest(H, split_ind, times_updated, name='times_updated')

            if to_update.shape[0] >= H.num_images_visualize + 8 and epoch % H.fid_freq == 0:
                latents = sampler.selected_latents[to_update[:H.num_images_visualize]]
                with torch.no_grad():
                    generate_for_NN(sampler, split_x_tensor[to_update[:H.num_images_visualize]], latents,
                                    [s[to_update[:H.num_images_visualize]]
                                        for s in sampler.selected_snoise],
                                    viz_batch_original.shape, imle,
                                    f'{H.save_dir}/NN-samples_{epoch}-imle.png', logprint)

            comb_dataset = ZippedDataset(
                split_x, TensorDataset(sampler.selected_latents))
            data_loader = DataLoader(comb_dataset, batch_size=H.n_batch, pin_memory=True,
                                     shuffle=False, num_workers=4, persistent_workers=False)

            start_time = time.time()

            for cur, indices in data_loader:
                x = cur[0]
                latents = cur[1][0]
                _, target = preprocess_fn(x)

                # if(H.use_snoise):
                cur_snoise = [s[indices] for s in sampler.selected_snoise]

                for i in range(len(H.res)):
                    cur_snoise[i].zero_()
                # else:
                #     cur_snoise = [s[indices] for s in sampler.selected_snoise]

                stat = training_step_imle(
                    H, target.shape[0], target, latents, cur_snoise, imle, ema_imle, optimizer, sampler.calc_loss)
                stats.append(stat)

                if (iterate <= H.warmup_iters):
                    # print("Warmup iteration: ", iterate)
                    scheduler.step()

                iterate += 1
                if iterate % H.iters_per_save == 0:
                    fp = os.path.join(H.save_dir, 'latest')
                    logprint(f'Saving model@ {iterate} to {fp}')
                    save_model(fp, imle, ema_imle, optimizer, scheduler, H)
                    save_latents_latest(H, split_ind, sampler.selected_latents)
                    save_latents_latest(
                        H, split_ind, change_thresholds, name='threshold_latest')
                    save_latents_latest(H, split_ind, last_updated, name='last_updated')
                    save_latents_latest(H, split_ind, times_updated, name='times_updated')

                if iterate % H.iters_per_ckpt == 0:
                    save_model(os.path.join(
                        H.save_dir, f'iter-{iterate}'), imle, ema_imle, optimizer, scheduler, H)
                    save_latents(H, iterate, split_ind,
                                 sampler.selected_latents)
                    save_latents(H, iterate, split_ind,
                                 change_thresholds, name='threshold')
                    save_snoise(H, iterate, sampler.selected_snoise)

                if _sigterm_received[0]:
                    _save_full_checkpoint(
                        split_ind, sampler, imle, ema_imle, optimizer,
                        scheduler, change_thresholds, last_updated,
                        times_updated)
                    signal.signal(signal.SIGTERM, prev_handler)
                    os.kill(os.getpid(), signal.SIGTERM)
                    return

            print(f'Epoch {epoch} took {time.time() - start_time} seconds')

            if (iterate > H.warmup_iters):
                scheduler.step()

            cur_dists = torch.empty([subset_len], dtype=torch.float32).cuda()
            cur_dists_lpips = torch.empty(
                [subset_len], dtype=torch.float32).cuda()
            cur_dists_l2 = torch.empty(
                [subset_len], dtype=torch.float32).cuda()

            cur_dists[:], cur_dists_lpips[:], cur_dists_l2[:] = sampler.calc_dists_existing(split_x_tensor, imle,
                                                                                            dists=cur_dists,
                                                                                            dists_lpips=cur_dists_lpips,
                                                                                            dists_l2=cur_dists_l2,
                                                                                            logging=True)

            # torch.save(cur_dists, f'{H.save_dir}/latent/dists-{epoch}.npy')

            metrics = {
                'mean_loss': torch.mean(cur_dists).item(),
                'std_loss': torch.std(cur_dists).item(),
                'max_loss': torch.max(cur_dists).item(),
                'min_loss': torch.min(cur_dists).item(),
                'mean_loss_lpips': torch.mean(cur_dists_lpips).item(),
                'std_loss_lpips': torch.std(cur_dists_lpips).item(),
                'max_loss_lpips': torch.max(cur_dists_lpips).item(),
                'min_loss_lpips': torch.min(cur_dists_lpips).item(),
                'mean_loss_l2': torch.mean(cur_dists_l2).item(),
                'std_loss_l2': torch.std(cur_dists_l2).item(),
                'max_loss_l2': torch.max(cur_dists_l2).item(),
                'min_loss_l2': torch.min(cur_dists_l2).item(),
                'total_excluded': sampler.total_excluded,
                'total_excluded_percentage': sampler.total_excluded_percentage,
            }

            if (epoch > 0 and epoch % H.fid_freq == 0):
                print("Learning rate: ", optimizer.param_groups[0]['lr'])

                # --- FID: 5000 generated samples (temp dir, cleaned after) ---
                generate_and_save(H, imle, sampler, 5000)
                print(f'{H.data_root}/img', f'{H.save_dir}/fid/')
                cur_fid = fid.compute_fid(f'{H.data_root}/img', f'{H.save_dir}/fid/',
                                          verbose=False, num_workers=0)

                # --- Save a 4x4 sample grid from the generated images ---
                with torch.no_grad():
                    grid_z = torch.randn(16, H.latent_dim).cuda()
                    grid_sn = [s[:16].normal_() for s in sampler.snoise_tmp]
                    grid_np = sampler.sample(grid_z, imle, grid_sn)
                    grid_t = torch.from_numpy(grid_np).float() / 255.0
                    grid_t = grid_t.permute(0, 3, 1, 2)
                    grid_img = torchvision.utils.make_grid(
                        grid_t, nrow=4, padding=2)
                    grid_pil = torchvision.transforms.functional.to_pil_image(
                        grid_img)
                    grid_pil.save(f'{H.save_dir}/samples_epoch_{epoch}.png')

                shutil.rmtree(f'{H.save_dir}/fid', ignore_errors=True)
                os.makedirs(f'{H.save_dir}/fid', exist_ok=True)

                is_new_best = False
                if cur_fid < best_fid:
                    best_fid = cur_fid
                    is_new_best = True
                    fp = os.path.join(H.save_dir, 'best')
                    logprint(
                        f'Saving best model (fid={best_fid:.4f}) @ {iterate} to {fp}')
                    save_model(fp, imle, ema_imle, optimizer, scheduler, H)

                # --- P/R: 1000 generated samples (temp dir, cleaned after) ---
                os.makedirs(f'{H.save_dir}/prec_rec', exist_ok=True)
                generate_and_save(H, imle, sampler, 1000, subdir='prec_rec')
                precision, recall = compute_prec_recall(
                    f'{H.data_root}/img', f'{H.save_dir}/prec_rec/')
                shutil.rmtree(f'{H.save_dir}/prec_rec', ignore_errors=True)

                metrics['fid'] = cur_fid
                metrics['best_fid'] = best_fid
                metrics['precision'] = precision
                metrics['recall'] = recall

                # --- Save metrics CSV and best_metrics.json ---
                csv_row = dict(epoch=epoch, fid=cur_fid,
                               precision=precision, recall=recall)
                append_metrics_csv(H.save_dir, csv_row)
                update_best_metrics(H.save_dir, csv_row)

                # --- Periodic checkpoint ---
                ckpt_fp = os.path.join(H.save_dir, f'epoch_{epoch}')
                logprint(f'Saving periodic ckpt @ epoch {epoch} to {ckpt_fp}')
                save_model(ckpt_fp, imle, ema_imle, optimizer, scheduler, H)

                # --- Slerp interpolation ---
                interp_num_pairs = 5
                interp_steps = 10
                with torch.no_grad():
                    z1 = torch.randn(interp_num_pairs, H.latent_dim).cuda()
                    z2 = torch.randn(interp_num_pairs, H.latent_dim).cuda()
                    t_vals = torch.linspace(0.0, 1.0, interp_steps).cuda()
                    all_rows = []
                    for pi in range(interp_num_pairs):
                        z_interp = slerp(
                            z1[pi:pi + 1].repeat(interp_steps, 1),
                            z2[pi:pi + 1].repeat(interp_steps, 1),
                            t_vals,
                        )
                        snoise_tmp = [s[:interp_steps].normal_()
                                      for s in sampler.snoise_tmp]
                        preds = sampler.sample(z_interp, imle, snoise_tmp)
                        preds_t = torch.from_numpy(preds).float() / 255.0
                        preds_t = preds_t.permute(0, 3, 1, 2)
                        all_rows.append(preds_t)
                    interp_grid = torchvision.utils.make_grid(
                        torch.cat(all_rows, dim=0), nrow=interp_steps)
                    save_grid_image(
                        interp_grid, f'{H.save_dir}/interp_slerp_{epoch}.png')

                # --- Latent walk ---
                n_perturb_seeds = 5
                n_perturb_steps = 10
                perturb_max_sigma = 1.0
                with torch.no_grad():
                    perturb_rows = []
                    for _ in range(n_perturb_seeds):
                        z_seed = torch.randn(1, H.latent_dim).cuda()
                        direction = torch.randn(1, H.latent_dim).cuda()
                        direction = direction / \
                            direction.norm(dim=1, keepdim=True)
                        magnitudes = torch.linspace(
                            0.0, perturb_max_sigma, n_perturb_steps).cuda()
                        z_walk = z_seed + magnitudes.unsqueeze(1) * direction
                        snoise_tmp = [s[:n_perturb_steps].normal_()
                                      for s in sampler.snoise_tmp]
                        preds = sampler.sample(z_walk, imle, snoise_tmp)
                        preds_t = torch.from_numpy(preds).float() / 255.0
                        preds_t = preds_t.permute(0, 3, 1, 2)
                        perturb_rows.append(preds_t)
                    perturb_all = torch.cat(perturb_rows, dim=0)
                    perturb_grid = torchvision.utils.make_grid(
                        perturb_all, nrow=n_perturb_steps, padding=2)
                    save_grid_image(
                        perturb_grid, f'{H.save_dir}/latent_walk_{epoch}.png')

                # --- Visual recall (LPIPS nearest) ---
                try:
                    import lpips as lpips_module
                    vr_num_real = 10
                    vr_num_fake = 200
                    vr_topk = 5
                    with torch.no_grad():
                        vr_real_imgs = split_x_tensor[:vr_num_real]
                        _, vr_real_processed = preprocess_fn([vr_real_imgs])
                        if vr_real_processed.dim() == 4 and vr_real_processed.shape[-1] in (1, 3):
                            vr_real_processed = vr_real_processed.permute(
                                0, 3, 1, 2)

                        z_vr = torch.randn(vr_num_fake, H.latent_dim).cuda()
                        snoise_tmp = [s[:min(H.imle_batch, vr_num_fake)].normal_()
                                      for s in sampler.snoise_tmp]
                        vr_fake_np = sampler.sample(z_vr[:min(H.imle_batch, vr_num_fake)],
                                                    imle, snoise_tmp)
                        all_fake = [torch.from_numpy(
                            vr_fake_np).float() / 255.0]
                        for fi in range(H.imle_batch, vr_num_fake, H.imle_batch):
                            bsz = min(H.imle_batch, vr_num_fake - fi)
                            sn = [s[:bsz].normal_()
                                  for s in sampler.snoise_tmp]
                            f_np = sampler.sample(z_vr[fi:fi + bsz], imle, sn)
                            all_fake.append(
                                torch.from_numpy(f_np).float() / 255.0)
                        vr_fake_all = torch.cat(all_fake, dim=0)
                        vr_fake_all = vr_fake_all.permute(0, 3, 1, 2)

                        lpips_fn = lpips_module.LPIPS(net='vgg').cuda()
                        vr_real_for_grid = (vr_real_processed + 1.0) / 2.0
                        for qi in range(vr_num_real):
                            dists = []
                            real_qi = vr_real_processed[qi:qi + 1].cuda()
                            for ci in range(0, vr_fake_all.shape[0], 8):
                                chunk = vr_fake_all[ci:ci + 8].cuda()
                                d = lpips_fn(real_qi.expand(chunk.shape[0], -1, -1, -1),
                                             chunk * 2 - 1)
                                dists.append(d.view(-1).cpu())
                            dists = torch.cat(dists)
                            topk_idx = torch.topk(-dists,
                                                  k=min(vr_topk, len(dists))).indices
                            row = torch.cat([
                                vr_real_for_grid[qi:qi + 1].cpu(),
                                vr_fake_all[topk_idx].cpu()
                            ], dim=0)
                            row_grid = torchvision.utils.make_grid(
                                row, nrow=vr_topk + 1)
                            save_grid_image(
                                row_grid, f'{H.save_dir}/visual_recall_{qi}_{epoch}.png')
                        del lpips_fn
                        torch.cuda.empty_cache()
                except ImportError:
                    print("lpips not installed, skipping visual recall")

            if (to_update.shape[0] != 0):
                metrics['mean_loss_resample'] = torch.mean(cur_dists).item()
                metrics['std_loss_resample'] = torch.std(cur_dists).item()
                metrics['max_loss_resample'] = torch.max(cur_dists).item()
                metrics['min_loss_resample'] = torch.min(cur_dists).item()

            log_metrics = {k: v for k, v in metrics.items()
                          if not k.startswith("viz/")}
            logprint(model=H.desc, type='train_loss',
                     epoch=epoch, step=iterate, **log_metrics)

            if epoch % H.fid_freq == 0:
                with torch.no_grad():
                    generate_images_initial(H, sampler, viz_batch_original,
                                            sampler.selected_latents[0: H.num_images_visualize],
                                            [s[0: H.num_images_visualize]
                                                for s in sampler.selected_snoise],
                                            viz_batch_original.shape, imle, ema_imle,
                                            f'{H.save_dir}/latest.png', logprint, experiment)

            if H.use_wandb:
                wandb.log(metrics, step=iterate)

            if experiment is not None:
                experiment.log_metrics(metrics, epoch=epoch, step=iterate)


def main(H=None):
    H_cur, logprint = set_up_hyperparams()
    if not H:
        H = H_cur
    H, data_train, data_valid_or_test, preprocess_fn = set_up_data(H)
    imle, ema_imle = load_imle(H, logprint)

    if H.use_comet and H.comet_api_key:
        if (H.comet_experiment_key):
            print("Resuming experiment")
            experiment = ExistingExperiment(
                api_key=H.comet_api_key,
                previous_experiment=H.comet_experiment_key
            )
            experiment.log_parameters(H)

        else:
            experiment = Experiment(
                api_key=H.comet_api_key,
                project_name=getattr(H, 'comet_project', 'adaptiveimle'),
                workspace=getattr(H, 'comet_workspace', None),
            )
            experiment.set_name(H.comet_name)
            experiment.log_parameters(H)
    else:
        experiment = None

    if H.use_wandb:
        wandb.init(
            name=H.wandb_name,
            project=H.wandb_project,
            config=H,
            mode=H.wandb_mode,
            dir=H.save_dir,
        )

    os.makedirs(f'{H.save_dir}/fid', exist_ok=True)

    if H.mode == 'eval':

        os.makedirs(f'{H.save_dir}/eval', exist_ok=True)
        print(H)

        with torch.no_grad():
            # Generating
            sampler = Sampler(H, len(data_train), preprocess_fn)
            n_samp = H.n_batch
            temp_latent_rnds = torch.randn(
                [n_samp, H.latent_dim], dtype=torch.float32).cuda()
            for i in range(0, H.num_images_to_generate // n_samp):
                if (i % 10 == 0):
                    print(i * n_samp)
                temp_latent_rnds.normal_()
                tmp_snoise = [s[:n_samp].normal_() for s in sampler.snoise_tmp]
                torch.save(temp_latent_rnds,
                           f'{H.save_dir}/eval/temp_latent_rnds_{i}.pt')
                torch.save(tmp_snoise, f'{H.save_dir}/eval/tmp_snoise_{i}.pt')
                samp = sampler.sample(temp_latent_rnds, imle, tmp_snoise)
                for j in range(n_samp):
                    imageio.imwrite(
                        f'{H.save_dir}/eval/{i * n_samp + j}.png', samp[j])

    elif H.mode == 'eval_fid':
        subset_len = H.subset_len
        if subset_len == -1:
            subset_len = len(data_train)
        sampler = Sampler(H, len(data_train), preprocess_fn)

        n_samp = getattr(H, 'num_fid_samples', 5000)
        subdir = getattr(H, 'eval_fid_subdir', 'fid')
        eval_model = ema_imle if ema_imle is not None else imle
        which = 'ema' if ema_imle is not None else 'main'
        print(f'[eval_fid] dumping {n_samp} samples to {H.save_dir}/{subdir}/ (using {which} model)')
        generate_and_save(H, eval_model, sampler, n_samp, subdir=subdir)

        if getattr(H, 'skip_cleanfid', False):
            print('[eval_fid] skip_cleanfid=True, not running cleanfid.compute_fid')
        else:
            try:
                print(f'{H.data_root}/img', f'{H.save_dir}/{subdir}/')
                cur_fid = fid.compute_fid(
                    f'{H.data_root}/img', f'{H.save_dir}/{subdir}/', verbose=False)
                print("FID: ", cur_fid)
            except Exception as e:
                print(f'[eval_fid] cleanfid.compute_fid failed: {e!r} '
                      '(ignored; samples already dumped)')

    elif H.mode == 'reconstruct':

        subset_len = H.subset_len
        if subset_len == -1:
            subset_len = len(data_train)
        ind = 0
        for split_ind, split_x_tensor in enumerate(DataLoader(data_train, batch_size=H.subset_len, pin_memory=True)):
            if (ind == 14):
                break
            split_x = TensorDataset(split_x_tensor[0])
            ind += 1

        for param in imle.parameters():
            param.requires_grad = False
        viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                             H.num_images_visualize, H.dataset)
        if os.path.isfile(str(H.restore_latent_path)):
            latents = torch.tensor(torch.load(
                H.restore_latent_path), requires_grad=True)
        else:
            latents = torch.randn(
                [viz_batch_original.shape[0], H.latent_dim], requires_grad=True)
        sampler = Sampler(H, subset_len, preprocess_fn)
        reconstruct(H, sampler, imle, preprocess_fn, viz_batch_original,
                    latents, 'reconstruct', logprint, training_step_imle)

    elif H.mode == 'backtrack':
        for param in imle.parameters():
            param.requires_grad = False
        for split_x in DataLoader(data_train, batch_size=H.subset_len):
            split_x = split_x[0]
            pass
        print(f'split shape is {split_x.shape}')
        sampler = Sampler(H, H.subset_len, preprocess_fn)
        backtrack(H, sampler, imle, preprocess_fn,
                  split_x, logprint, training_step_imle)

    elif H.mode == 'train':
        print(H)
        train_loop_imle(H, data_train, data_valid_or_test,
                        preprocess_fn, imle, ema_imle, logprint, experiment)

    elif H.mode == 'ppl':
        sampler = Sampler(H, H.subset_len, preprocess_fn)
        calc_ppl(H, imle, sampler)

    elif H.mode == 'ppl_uniform':
        sampler = Sampler(H, H.subset_len, preprocess_fn)
        calc_ppl_uniform(H, imle, sampler)

    elif H.mode == 'interpolate':
        subset_len = H.subset_len
        if subset_len == -1:
            subset_len = len(data_train)
        with torch.no_grad():
            for split_x in DataLoader(data_train, batch_size=subset_len):
                split_x = split_x[0]
            viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                                 H.num_images_visualize, H.dataset)
            sampler = Sampler(H, subset_len, preprocess_fn)
            for i in range(H.num_images_to_generate):
                random_interp(H, sampler, (0, 256, 256, 3), imle,
                              f'{H.save_dir}/interp-{i}.png', logprint)

    elif H.mode == 'spatial_visual':
        with torch.no_grad():
            for split_x in DataLoader(data_train, batch_size=H.subset_len):
                split_x = split_x[0]
            viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                                 H.num_images_visualize, H.dataset)
            sampler = Sampler(H, H.subset_len, preprocess_fn)
            for i in range(H.num_images_to_generate):
                print(H.num_images_to_generate, i)
                spatial_vissual(H, sampler, (0, 256, 256, 3),
                                imle, f'{H.save_dir}/interp-{i}.png', logprint)

    elif H.mode == 'generate_rnd':
        with torch.no_grad():
            for split_x in DataLoader(data_train, batch_size=H.subset_len):
                split_x = split_x[0]
            viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                                 H.num_images_visualize, H.dataset)
            sampler = Sampler(H, H.subset_len, preprocess_fn)
            generate_rnd(H, sampler, (0, 256, 256, 3), imle,
                         f'{H.save_dir}/rnd.png', logprint)

    elif H.mode == 'generate_rnd_nn':
        with torch.no_grad():
            for split_x in DataLoader(data_train, batch_size=len(data_train)):
                split_x = split_x[0]
            viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                                 H.num_images_visualize, H.dataset)
            sampler = Sampler(H, H.subset_len, preprocess_fn)
            generate_rnd_nn(H, split_x,  sampler, (0, 256, 256, 3),
                            imle, f'{H.save_dir}', logprint, preprocess_fn)

    elif H.mode == 'nn_interp':
        with torch.no_grad():
            for split_x in DataLoader(data_train, batch_size=len(data_train)):
                split_x = split_x[0]
            viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                                 H.num_images_visualize, H.dataset)
            sampler = Sampler(H, H.subset_len, preprocess_fn)
            nn_interp(H, split_x,  sampler, (0, 256, 256, 3), imle,
                      f'{H.save_dir}', logprint, preprocess_fn)

    elif H.mode == 'generate_sample_nn':
        with torch.no_grad():
            for split_x in DataLoader(data_train, batch_size=len(data_train)):
                split_x = split_x[0]
            viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                                 H.num_images_visualize, H.dataset)
            sampler = Sampler(H, H.subset_len, preprocess_fn)
            generate_sample_nn(H, split_x,  sampler, (0, 256, 256, 3),
                               imle, f'{H.save_dir}/rnd2.png', logprint, preprocess_fn)

    elif H.mode == 'backtrack_interpolate':
        subset_len = H.subset_len
        if subset_len == -1:
            subset_len = len(data_train)
        with torch.no_grad():
            for split_x in DataLoader(data_train, batch_size=subset_len):
                split_x = split_x[0]
            viz_batch_original, _ = get_sample_for_visualization(split_x, preprocess_fn,
                                                                 H.num_images_visualize, H.dataset)
            sampler = Sampler(H, subset_len, preprocess_fn)
            latents = torch.tensor(torch.load(
                f'{H.restore_latent_path}'), requires_grad=True, dtype=torch.float32, device='cuda')
            for i in range(latents.shape[0] - 1):
                lat0 = latents[i:i+1]
                lat1 = latents[i+1:i+2]
                sn1 = None
                sn2 = None
                random_interp(H, sampler, (0, 256, 256, 3), imle,
                              f'{H.save_dir}/back-interp-{i}.png', logprint, lat0, lat1, sn1, sn2)

    elif H.mode == 'prec_rec':

        os.makedirs(f'{H.save_dir}/prec_rec', exist_ok=True)

        subset_len = H.subset_len
        if subset_len == -1:
            subset_len = len(data_train)
        sampler = Sampler(H, len(data_train), preprocess_fn)
        # generate_and_save(H, imle, sampler, 5000)

        print("Generating images")
        generate_and_save(H, imle, sampler, 1000, subdir='prec_rec')
        print(f'{H.data_root}/img', f'{H.save_dir}/prec_rec/')
        precision, recall = compute_prec_recall(
            f'{H.data_root}/img', f'{H.save_dir}/prec_rec/')
        print("Precision: ", precision)
        print("Recall: ", recall)


if __name__ == "__main__":
    main()