code/sampler.py

from curses import update_lines_cols
from math import comb, ceil
import os
import time

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoImageProcessor, AutoModel

from LPNet import LPNet
from helpers.utils import (
    configure_inductor_for_low_memory_compile,
    is_dist_avail_and_initialized,
    is_main_process,
    get_world_size,
    get_rank,
    safe_barrier,
)
from models import parse_layer_string
from helpers.angle_sampler import Angle_Generator
from torch import autocast
import faiss
from tqdm import tqdm

class Sampler:
    def __init__(self, H, sz, preprocess_fn):
        
        self.device = torch.device("cuda", torch.cuda.current_device())
        self.world_size = get_world_size()
        self.rank = get_rank()

        self.pool_size = ceil(int(H.force_factor * sz) / H.imle_db_size) * H.imle_db_size
        self.preprocess_fn = preprocess_fn
        self.l2_loss = torch.nn.MSELoss(reduce=False).to(self.device)
        self.H = H
        self.latent_lr = H.latent_lr
        self.sz = sz
        self.entire_ds = torch.arange(sz)
        self.selected_latents = torch.empty([sz, H.latent_dim], dtype=torch.float32)
        self.last_selected_latents = torch.empty([sz, H.latent_dim], dtype=torch.float32)
        self.selected_latents_tmp = torch.empty([sz, H.latent_dim], dtype=torch.float32)

        blocks = parse_layer_string(H.dec_blocks)
        self.block_res = [s[0] for s in blocks]
        self.res = sorted(set([s[0] for s in blocks if s[0] <= H.max_hierarchy]))

        self.selected_dists = torch.empty([sz], dtype=torch.float32)
        self.selected_dists[:] = np.inf
        self.selected_dists_tmp = torch.empty([sz], dtype=torch.float32)

        self.selected_dists_lpips = torch.empty([sz], dtype=torch.float32)
        self.selected_dists_lpips[:] = np.inf

        self.selected_dists_l2 = torch.empty([sz], dtype=torch.float32)
        self.selected_dists_l2[:] = np.inf

        self.temp_latent_rnds = torch.empty([self.H.imle_db_size, self.H.latent_dim], dtype=torch.float32)
        self.temp_samples = torch.empty([self.H.imle_db_size, H.image_channels, self.H.image_size, self.H.image_size],
                                        dtype=torch.float32)

        self.pool_latents = None

        self.projections = []
        self.lpips_net = LPNet(pnet_type=H.lpips_net, path=H.lpips_path).to(self.device)
        self.lpips_net.eval()
        self.lpips_net.requires_grad_(False)

        if self.H.compile:
            configure_inductor_for_low_memory_compile()
            self.lpips_net = torch.compile(self.lpips_net)

        self._needs_dino = (H.image_size > 32) or (H.search_type == 'combined')

        self.dino_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073], device=self.device).view(1, 3, 1, 1)
        self.dino_std = torch.tensor([0.26862954, 0.26130258, 0.27577711], device=self.device).view(1, 3, 1, 1)

        if self._needs_dino:
            dino_cache_dir = getattr(H, "dino_cache_dir", None) \
                or os.environ.get("DINO_CACHE_DIR") \
                or "./dinov2_cache"
            snapshots_dir = os.path.join(
                dino_cache_dir,
                "models--facebook--dinov2-base",
                "snapshots",
            )
            local_snapshot = None
            if os.path.isdir(snapshots_dir):
                main_path = os.path.join(snapshots_dir, "main")
                if os.path.isdir(main_path):
                    local_snapshot = main_path
                else:
                    candidates = sorted(
                        d for d in os.listdir(snapshots_dir)
                        if os.path.isdir(os.path.join(snapshots_dir, d))
                    )
                    if candidates:
                        local_snapshot = os.path.join(snapshots_dir, candidates[0])
            if local_snapshot is not None:
                dino_model = AutoModel.from_pretrained(local_snapshot, local_files_only=True)
            else:
                dino_model = AutoModel.from_pretrained(
                    "facebook/dinov2-base",
                    cache_dir=dino_cache_dir,
                    local_files_only=True,
                )
            self.dino_encoder = dino_model.eval().to(self.device)
            
            if self.H.compile:
                configure_inductor_for_low_memory_compile()
                self.dino_encoder = torch.compile(self.dino_encoder)
        else:
            self.dino_encoder = None
        
        self.nn_search_batch = H.nn_search_batch



        self.l2_projection = None

        fake = torch.zeros(1, 3, H.image_size, H.image_size, device=self.device)

        safe_barrier()
        if(H.search_type == 'lpips'):
            interpolated = F.interpolate(fake,scale_factor = H.l2_search_downsample, antialias=True, mode='bicubic')
            out, shapes = self.lpips_net(interpolated)
            sum_dims = 0
            dims = [int(H.proj_dim * 1. / len(out)) for _ in range(len(out))]
            if H.proj_proportion:
                sm = sum([dim.shape[1] for dim in out])
                dims = [int(out[feat_ind].shape[1] * (H.proj_dim / sm)) for feat_ind in range(1,len(out))]
                dims.insert(0,H.proj_dim - sum(dims))
            for ind, feat in enumerate(out):
                self.projections.append(F.normalize(torch.randn(feat.shape[1], dims[ind], device=self.device), p=2, dim=1))
            sum_dims = sum(dims)

        elif(H.search_type == 'l2'):
            interpolated = F.interpolate(fake,scale_factor = H.l2_search_downsample, antialias=True, mode='bicubic')
            interpolated = interpolated.reshape(interpolated.shape[0],-1)
            self.l2_projection = F.normalize(torch.randn(interpolated.shape[1], H.proj_dim, device=self.device), p=2, dim=1)
            sum_dims = H.proj_dim
        
        elif(H.search_type == 'combined'):
            interpolated = F.interpolate(fake,scale_factor = H.l2_search_downsample, antialias=True, mode='bicubic')
            out, shapes = self.lpips_net(interpolated)
            sum_dims = 0
            dims = [int(H.proj_dim * 1. / len(out)) for _ in range(len(out))]
            if H.proj_proportion:
                sm = sum([dim.shape[1] for dim in out])
                dims = [int(out[feat_ind].shape[1] * (H.proj_dim / sm)) for feat_ind in range(1,len(out))]
                dims.insert(0,H.proj_dim - sum(dims))
            for ind, feat in enumerate(out):
                self.projections.append(F.normalize(torch.randn(feat.shape[1], dims[ind], device=self.device), p=2, dim=1))
            sum_dims = sum(dims)

            interpolated = self.preprocess_dino_tensor(fake)
            with torch.no_grad():
                out = self.dino_encoder(pixel_values=interpolated)
                out = out.last_hidden_state.mean(dim=1)            
            sum_dims += out.shape[-1]

        else:
            exit()

        self.dci_dim = sum_dims

        self.dataset_proj = torch.empty([sz, sum_dims], dtype=torch.float32, device='cpu')
        self.pool_samples_proj = None

        self.knn_ignore = H.knn_ignore
        self.ignore_radius = H.ignore_radius
        self.resample_angle = H.resample_angle

        self.total_excluded = 0
        self.total_excluded_percentage = 0

        self.dataset_size = sz
        self.db_iter = 0
        self.generator_seed = torch.Generator(device=self.device)         
        self.generator_seed.manual_seed(H.seed + self.rank)

        self.faiss_res = faiss.StandardGpuResources()  # one per process
        index_flat = faiss.IndexFlatL2(self.dci_dim)  # identical API to IndexFlatL2
        dev_id = torch.cuda.current_device()
        self.gpu_index_flat = faiss.index_cpu_to_gpu(self.faiss_res, dev_id, index_flat)


    def preprocess_dino_tensor(self, inp):
        # x: [B, C, H, W], range [0, 1]

        x = (inp + 1.0) / 2.0
        x = torch.clamp(x, 0.0, 1.0)

        x = F.interpolate(x, size=(224, 224), mode='bicubic', align_corners=False)
        return (x - self.dino_mean) / self.dino_std


    def get_projected(self, inp, permute=True):
        if(permute):
            inp = inp.permute(0, 3, 1, 2)
        
        interpolated = F.interpolate(inp,scale_factor = self.H.l2_search_downsample, antialias=True, mode='bicubic')
        out, _ = self.lpips_net(interpolated.to(self.device))
        gen_feat = []
        for i in range(len(out)):
            gen_feat.append(torch.mm(out[i], self.projections[i]))
            
        lpips_feat = torch.cat(gen_feat, dim=1)
        # lpips_feat = F.normalize(lpips_feat, p=2, dim=1)
        return lpips_feat
    
    def get_l2_feature(self, inp, permute=True):
        if(permute):
            inp = inp.permute(0, 3, 1, 2)
        interpolated = F.interpolate(inp,scale_factor = self.H.l2_search_downsample, antialias=True, mode='bicubic')
        interpolated = interpolated.reshape(interpolated.shape[0],-1)
        interpolated = torch.mm(interpolated, self.l2_projection)
        # interpolated = F.normalize(interpolated, p=2, dim=1)
        return interpolated
    
    def get_dino_features(self, inp, permute=True, scale_factor=10):
        if(permute):
            inp = inp.permute(0, 3, 1, 2)
        interpolated = self.preprocess_dino_tensor(inp)
        with torch.no_grad():
            out = self.dino_encoder(pixel_values=interpolated)
            out = out.last_hidden_state.mean(dim=1)   
            out = F.normalize(out, p=2, dim=1)
            out = out * scale_factor
        return out
    
    def get_combined_feature(self, inp, permute=True):
        lpisps_feat = self.get_projected(inp, permute)
        dino_feat = self.get_dino_features(inp, permute)
        # print(f'LPIPS is {torch.norm(lpisps_feat, p=2, dim=1).mean()} \n')
        # print(f'DINO is {torch.norm(dino_feat, p=2, dim=1).mean()} \n')
        combined_feat = torch.cat((lpisps_feat, dino_feat), dim=1)
        return combined_feat

    def init_projection(self, dataset):

        dataloader = DataLoader(
            dataset,
            batch_size=self.H.imle_batch,      # Get 32 samples per batch
        )

        if(is_main_process()):
            print("Starting Initialization")

        for ind, x in tqdm(enumerate(dataloader), total=len(dataloader), desc="Initializing"):
            batch_slice = slice(ind * self.H.imle_batch, ind * self.H.imle_batch + x[0].shape[0])
            if(self.H.search_type == 'lpips'):
                self.dataset_proj[batch_slice] = self.get_projected(self.preprocess_fn(x)[1]).cpu()
            elif(self.H.search_type == 'l2'):
                self.dataset_proj[batch_slice] = self.get_l2_feature(self.preprocess_fn(x)[1]).cpu()
            elif(self.H.search_type == 'combined'):
                self.dataset_proj[batch_slice] = self.get_combined_feature(self.preprocess_fn(x)[1]).cpu()
            else:
                exit()

        self.dataset_proj = self.dataset_proj.cpu().numpy().astype(np.float32)

    def sample(self, latents, gen, snoise=None):
        with torch.no_grad():
            with autocast(device_type='cuda'):
                latents = latents.to(self.device)
                px_z = gen(latents, None).permute(0, 2, 3, 1)
                xhat = (px_z + 1.0) * 127.5
                xhat = xhat.detach().cpu().numpy()
                
                xhat = np.nan_to_num(xhat, nan=0.0, posinf=255.0, neginf=0.0)
                xhat = np.minimum(np.maximum(0.0, xhat), 255.0).astype(np.uint8)
                return xhat

    def get_lpips_loss(self, inp, tar, use_mean=True):
        res = 0
        if(inp.shape[2] < 32):
            inp_interpolated = F.interpolate(inp, size=(32,32), mode='bicubic')
            tar_interpolated = F.interpolate(tar, size=(32,32), mode='bicubic')
        else:
            inp_interpolated = inp
            tar_interpolated = tar
        inp_feat, inp_shape = self.lpips_net(inp_interpolated)
        tar_feat, _ = self.lpips_net(tar_interpolated)
        for i, g_feat in enumerate(inp_feat):
            lpips_feature_residual = (g_feat - tar_feat[i])

            if self.H.loss_type == 'huber':
                lpips_feature_loss = self.pseudo_huber(lpips_feature_residual)
            elif self.H.loss_type == 'mclure':
                lpips_feature_loss = self.mclure_loss(lpips_feature_residual)
            elif self.H.loss_type == 'welsch':
                lpips_feature_loss = self.welsch_loss(lpips_feature_residual)
            else:
                lpips_feature_loss = lpips_feature_residual.pow(2)

            res += torch.sum(lpips_feature_loss, dim=1) / (inp_shape[i] ** 2)
        
        return res.mean()

    
    def get_dino_loss(self, inp, tar, use_mean=True):
        dino_feat = self.get_dino_features(inp, scale_factor=1, permute=False)
        tar_feat = self.get_dino_features(tar, scale_factor=1, permute=False)
        dino_residual = (dino_feat - tar_feat) 

        if self.H.loss_type == 'huber':
            dino_loss = self.pseudo_huber(dino_residual)
        elif self.H.loss_type == 'mclure':
            dino_loss = self.mclure_loss(dino_residual)
        elif self.H.loss_type == 'welsch':
            dino_loss = self.welsch_loss(dino_residual)
        else:
            dino_loss = dino_residual.pow(2)

        return dino_loss.mean()
    
    def pseudo_huber(self, residual):
        return self.H.huber_delta**2 * (torch.sqrt(1.0 + (residual / self.H.huber_delta) ** 2) - 1.0)

    def mclure_loss(self, residual):
        return (residual ** 2) / (residual ** 2 + self.H.loss_scale ** 2)

    def welsch_loss(self, residual):
        return 1 - torch.exp(-(residual / self.H.loss_scale)**2)

    def calc_loss(self, inp, tar, use_mean=True, logging=False):

        pixel_residual = (inp - tar)
        
        if self.H.loss_type == 'huber':
            pixel_loss = self.pseudo_huber(pixel_residual).mean()
        elif self.H.loss_type == 'mclure':
            pixel_loss = self.mclure_loss(pixel_residual).mean()
        elif self.H.loss_type == 'welsch':
            pixel_loss = self.welsch_loss(pixel_residual).mean()
        else:
            pixel_loss = pixel_residual.pow(2).mean()

        lpips_loss = self.get_lpips_loss(inp, tar, use_mean=True)

        loss = self.H.lpips_coef * lpips_loss + self.H.pixel_coef * pixel_loss

        if inp.shape[2] > 32:
            dino_loss = self.get_dino_loss(inp, tar, use_mean=True)
            loss = loss + self.H.dino_coef * dino_loss

        return loss.mean()

    def calc_dists_existing(self, dataset_tensor, gen, dists=None, dists_lpips = None, dists_l2 = None, latents=None, to_update=None, snoise=None, logging=False):
        if dists is None:
            dists = self.selected_dists
        if dists_lpips is None:
            dists_lpips = self.selected_dists_lpips
        if dists_l2 is None:
            dists_l2 = self.selected_dists_l2
        if latents is None:
            latents = self.selected_latents

        if to_update is not None:
            latents = latents[to_update]
            dists = dists[to_update]
            dataset_tensor = dataset_tensor[to_update]

        for ind, x in enumerate(DataLoader(TensorDataset(dataset_tensor), batch_size=self.H.n_batch)):
            _, target = self.preprocess_fn(x)
            batch_slice = slice(ind * self.H.n_batch, ind * self.H.n_batch + target.shape[0])
            cur_latents = latents[batch_slice]
            with torch.no_grad():
                with autocast(device_type='cuda'):
                    out = gen(cur_latents, None)
                    if(logging):
                        dist, dist_lpips, dist_l2 = self.calc_loss(target.permute(0, 3, 1, 2), out, use_mean=False, logging=True)
                        dists[batch_slice] = torch.squeeze(dist)
                        dists_lpips[batch_slice] = torch.squeeze(dist_lpips)
                        dists_l2[batch_slice] = torch.squeeze(dist_l2)
                    else:
                        dist = self.calc_loss(target.permute(0, 3, 1, 2), out, use_mean=False)
                        dists[batch_slice] = torch.squeeze(dist)
        
        if(logging):
            return dists, dists_lpips, dists_l2
        else:
            return dists

    def resample_pool(self, gen):

        gen.eval()   

        # Determine local pool size
        local_pool_size = ceil(self.pool_size / self.world_size)


        # Generate local pool latents and prepare container for projected features
        local_pool_latents = torch.randn((local_pool_size, self.H.latent_dim), 
                                         device=self.device, 
                                         generator=self.generator_seed)
        # Assuming pool_samples_proj is preallocated with shape (self.pool_size, projection_dim)

        local_pool_proj = torch.empty((local_pool_size, self.dci_dim), device=self.device)

        # Process local chunk in batches
        for j in range(local_pool_size // self.H.imle_batch):
            batch_slice = slice(j * self.H.imle_batch, (j + 1) * self.H.imle_batch)
            cur_latents = local_pool_latents[batch_slice]
            with torch.no_grad():
                with autocast(device_type='cuda'):
                    outputs = gen(cur_latents, None)
                    if self.H.search_type == 'lpips':
                        proj = self.get_projected(outputs, False)
                    elif self.H.search_type == 'l2':
                        proj = self.get_l2_feature(outputs, False)
                    elif self.H.search_type == 'combined':
                        proj = self.get_combined_feature(outputs, False)
                    else:
                        proj = self.get_combined_feature(outputs, False)
                    local_pool_proj[batch_slice] = proj

        safe_barrier()
        gathered_latents = [torch.empty_like(local_pool_latents) for _ in range(self.world_size)]
        gathered_proj = [torch.empty_like(local_pool_proj) for _ in range(self.world_size)]

        torch.distributed.all_gather(gathered_latents, local_pool_latents)
        torch.distributed.all_gather(gathered_proj, local_pool_proj)

        gen.train()

        safe_barrier()
        # Aggregate the full pool latents and projected features
        self.pool_latents = torch.cat(gathered_latents, dim=0).to('cpu')
        self.pool_samples_proj = torch.cat(gathered_proj, dim=0).to('cpu')
    

    def nn_search_batched(self, queries, dataset):

        topk = self.H.imle_db_topk
        tie_shuffle = True  # avoid ordering bias for equal/near-equal margins

        Nq = queries.shape[0]
        Nd = dataset.shape[0]
        if Nq == 0:
            return torch.empty(0, dtype=torch.float32), torch.empty(0, dtype=torch.long)

        topk = int(min(max(1, topk), Nd))

        # ---- Build index once on the full dataset ----
        self.gpu_index_flat.reset()
        self.gpu_index_flat.add(dataset)

        # ---- 1) Hardness (margin = d2 - d1) ----
        # Need k=2 even if topk==1, to get a margin; if Nd==1 margin is 0.
        if Nd >= 2:
            D2, _ = self.gpu_index_flat.search(queries, 2)  # (Nq,2)
            margin = D2[:, 0]
        else:
            margin = np.zeros(Nq, dtype=np.float32)

        if tie_shuffle:
            perm = np.random.permutation(Nq)
            order = perm[np.argsort(margin[perm], kind="stable")]
        else:
            order = np.argsort(margin, kind="stable")

        # ---- 2) Get Top-K candidate lists for all queries ----
        D, I = self.gpu_index_flat.search(queries, topk)  # (Nq,K), squared L2 + indices

        # ---- 3) Greedy unique assignment in hard-first order ----
        used = np.zeros(Nd, dtype=bool)

        out_idx = np.empty(Nq, dtype=np.int64)
        out_dst = np.empty(Nq, dtype=np.float32)

        I_ordered = I[order]
        D_ordered = D[order]

        for i in range(len(order)):
            qi = order[i]
            cand = I_ordered[i]
            cd   = D_ordered[i]

            mask = ~used[cand]
            valid = np.flatnonzero(mask)
            if len(valid) > 0:
                k = valid[0]
                chosen = cand[k]
                used[chosen] = True
                out_idx[qi] = chosen
                out_dst[qi] = cd[k]
            else:
                out_idx[qi] = cand[0]
                out_dst[qi] = cd[0]

        # ---- Cleanup ----
        self.gpu_index_flat.reset()

        return torch.from_numpy(out_dst), torch.from_numpy(out_idx)


    def imle_sample_force(self, gen, to_update=None):
        """
        Optimized force resampling routine using FAISS for batched nearest-neighbor search.
        In a DDP setting, each process handles a different subset of the dataset features,
        performs NN search locally, and then the results are merged and broadcast. 
        """
        if is_main_process():
            t1 = time.time()
            print("Starting pool resampling...")

        # Resample pool first (each process contributes its part);
        # this updates self.pool_samples_proj and self.pool_latents.
        self.resample_pool(gen)
        safe_barrier()  # Ensure all processes complete the pool resample

        if(is_main_process()):
            print(f"Resampling pool took {time.time() - t1:.2f} seconds")
        
        torch.cuda.empty_cache()

        self.selected_dists_tmp[:] = np.inf

        with torch.no_grad():

            if(is_main_process()):

                local_ds_feats = np.ascontiguousarray(self.dataset_proj, dtype=np.float32)

                # Pool features (as computed from resample_pool).
                pool_feats = np.ascontiguousarray(self.pool_samples_proj.cpu().numpy().astype(np.float32), dtype=np.float32)

                # Perform NN search for the local chunk. Returns arrays of shape (local_size, 1).
                local_distances, local_indices = self.nn_search_batched(local_ds_feats, pool_feats)

                new_latents = self.pool_latents[local_indices].clone()
            
            safe_barrier()  # Ensure all processes complete the gather

            if is_main_process():
                full_updated_latents = new_latents.to(self.device)
                perturbation = self.H.imle_perturb_coef * torch.randn(
                    (self.sz, self.H.latent_dim), 
                    device=self.device,
                    generator=self.generator_seed)
                full_updated_latents += perturbation
            else:
                full_updated_latents = torch.empty(self.sz, self.H.latent_dim, dtype=torch.float32, device=self.device)

            safe_barrier()

            torch.distributed.broadcast(full_updated_latents, src=0)

            safe_barrier()

            # Move the broadcasted results to CPU if desired.
            self.selected_latents_tmp = full_updated_latents.cpu()

            # Update last and current selected latents on all processes.
            self.last_selected_latents = self.selected_latents.clone()
            self.selected_latents = self.selected_latents_tmp.clone()

            if is_main_process():
                print(f"Force resampling took {time.time() - t1:.2f} seconds")

        safe_barrier()  # Ensure synchronization before leaving the function
        self.gpu_index_flat.reset()