patch-reparameterization/src/eval/__init__.py

from .ref_iqa import calculate_psnr, calculate_lpips, calculate_ssim
from .fid import calculate_rfid, calculate_gfid
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch.cuda.amp import autocast
from typing import Union, Dict, Optional
try:
    from accelerate import Accelerator
except ImportError:
    Accelerator = None
from torch.utils.data import DataLoader, Subset
from tqdm import tqdm
import os
import sys
from functools import lru_cache
@lru_cache(maxsize=4)
def _load_encoder_cached(model_path: str, device_str: str):
    """
    懒加载并缓存 encoder（按 model_path + device 缓存，进程内只加载一次）。
    使用 transformers.AutoModel + AutoProcessor。
    """
    from transformers import AutoModel, AutoProcessor
    print(f"[CosSim] Loading encoder from {model_path} onto {device_str} (cached)...")
    processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
    model = model.to(device_str).eval()
    return model, processor


@torch.no_grad()
def calculate_encoder_cos_sim(
    ref_arr: np.ndarray,
    rec_arr: np.ndarray,
    encoder_model: str,
    device: torch.device,
    batch_size: int = 64,
) -> float:
    """
    用指定 encoder 对 GT 图像和重建图像编码，计算 image embedding 的 cos 相似度均值。

    Args:
        ref_arr: GT 图像 [N, H, W, C] uint8
        rec_arr: 重建图像 [N, H, W, C] uint8
        encoder_model: HuggingFace 模型路径或名称
        device: 计算设备
        batch_size: 每批编码的图像数

    Returns:
        所有样本的 cos 相似度均值 (float)
    """
    device_str = "cuda" if device.type == "cuda" else "cpu"
    model, processor = _load_encoder_cached(encoder_model, device_str)

    N = len(ref_arr)
    cos_sims = []

    for start in range(0, N, batch_size):
        end = min(start + batch_size, N)
        ref_batch = [Image.fromarray(img) for img in ref_arr[start:end]]
        rec_batch = [Image.fromarray(img) for img in rec_arr[start:end]]

        # 预处理：使用 processor 处理两组图像
        ref_inputs = processor(images=ref_batch, return_tensors="pt").to(device_str)
        rec_inputs = processor(images=rec_batch, return_tensors="pt").to(device_str)

        # 提取 image embedding
        ref_feats = model.get_image_features(**ref_inputs)  # [B, D]
        rec_feats = model.get_image_features(**rec_inputs)  # [B, D]

        # L2 归一化后计算 cos 相似度
        ref_feats = F.normalize(ref_feats.float(), dim=-1)
        rec_feats = F.normalize(rec_feats.float(), dim=-1)
        batch_cos = (ref_feats * rec_feats).sum(dim=-1)  # [B]
        cos_sims.append(batch_cos.cpu())

    cos_sim_mean = torch.cat(cos_sims).mean().item()
    return cos_sim_mean


def compute_reconstruction_metrics(
    ref_arr: np.ndarray,
    rec_arr: np.ndarray,
    device: torch.device,
    batch_size: int = 128,
    metrics_to_compute = ("psnr", "ssim", "rfid"),
    disable_bar: bool = True,
    encoder_model: Optional[str] = None,
    psnr_resize: Optional[int] = None,
    l1_resize: Optional[int] = None,
) -> Dict[str, float]:
    """
    Compute reconstruction metrics between reference and reconstructed images.

    Args:
        ref_arr: Reference images [N, H, W, C] uint8
        rec_arr: Reconstructed images [N, H, W, C] uint8
        device: Device for computation
        batch_size: Batch size for metric computation

    Returns:
        Dictionary with metrics: eval/psnr, eval/ssim, eval/rfid
        Note: LPIPS is not computed here since it's already tracked during training
    """
    device_str = "cuda" if device.type == "cuda" else "cpu"
    # For paired metrics (psnr, ssim, l1, cos_sim), align ref/rec count
    n_paired = min(len(ref_arr), len(rec_arr))
    ref_paired = ref_arr[:n_paired]
    rec_paired = rec_arr[:n_paired]
    results_dict = {}
    if 'psnr' in metrics_to_compute:
        psnr = calculate_psnr(ref_paired, rec_paired, batch_size, device_str, disable_bar=disable_bar, psnr_resize=psnr_resize)
        results_dict["psnr"] = psnr
    if 'ssim' in metrics_to_compute:
        ssim = calculate_ssim(ref_paired, rec_paired, batch_size, device_str, disable_bar=disable_bar)
        results_dict["ssim"] = ssim
    if 'rfid' in metrics_to_compute:
        rfid = calculate_rfid(ref_arr, rec_arr, batch_size, device_str)
        results_dict["rfid"] = rfid
    if 'l1' in metrics_to_compute:
        # L1 = mean(|rec - ref|), same as train: (recon - images).abs().mean()
        ref_f = ref_paired.astype(np.float32) / 255.0
        rec_f = rec_paired.astype(np.float32) / 255.0
        if l1_resize is not None:
            # resize via torch interpolate to keep consistent with psnr_resize style
            # Process in small batches to avoid CUDA OOM on large datasets
            device_str = "cuda" if device.type == "cuda" else "cpu"
            resize_batch = 64
            n = ref_f.shape[0]
            l1_sum = 0.0
            l1_count = 0
            for i in range(0, n, resize_batch):
                ref_b = torch.from_numpy(ref_f[i:i+resize_batch]).permute(0, 3, 1, 2).to(device_str)
                rec_b = torch.from_numpy(rec_f[i:i+resize_batch]).permute(0, 3, 1, 2).to(device_str)
                ref_b = F.interpolate(ref_b, size=(l1_resize, l1_resize), mode="bilinear", align_corners=False)
                rec_b = F.interpolate(rec_b, size=(l1_resize, l1_resize), mode="bilinear", align_corners=False)
                l1_sum += (ref_b - rec_b).abs().sum().item()
                l1_count += ref_b.numel()
            l1 = l1_sum / l1_count
        else:
            l1 = np.abs(rec_f - ref_f).mean()
        results_dict["l1"] = float(l1)
    if 'cos_sim' in metrics_to_compute:
        assert encoder_model is not None, \
            "'cos_sim' metric requires encoder_model (HuggingFace model path) to be specified."
        print(f"[CosSim] Computing encoder cos_sim with model: {encoder_model}")
        cos_sim = calculate_encoder_cos_sim(
            ref_paired, rec_paired,
            encoder_model=encoder_model,
            device=device,
            batch_size=batch_size,
        )
        results_dict["cos_sim"] = cos_sim
    assert len(results_dict) > 0, "No metrics were computed."
    return results_dict
def compute_generation_metrics(
    ref_arr: np.ndarray,
    rec_arr: np.ndarray,
    device: torch.device,
    batch_size: int = 128,
):
    device_str = "cuda" if device.type == "cuda" else "cpu"
    # only eval FID
    fid = calculate_gfid(rec_arr, ref_arr, batch_size, device_str)
    return {
        'fid': fid
    }
@torch.no_grad()
def evaluate_generation_distributed(
    model_fn,
    sample_fn,
    latent_size, # for noise
    additional_model_kwargs,
    use_guidance: bool,
    pr,
    val_dataset,
    num_samples: int,
    batch_size: int,
    rank: Union[int, None] = None,
    world_size: Union[int, None] = None,
    device: torch.device = None,
    experiment_dir: str = None,
    global_step: int = None,
    autocast_kwargs: dict = None,
    metric_batch_size: int = 128,
    reference_npz_path: Optional[str] = None,
    accelerator: Union[Accelerator, None] = None,
) -> Optional[Dict[str, float]]:
    """
    Evaluate generation metrics using all GPUs in a distributed manner.

    Args:
        val_dataset: Validation dataset
        batch_size: Batch size per GPU for generation
        rank: Current GPU rank (deprecated, use accelerator instead)
        world_size: Total number of GPUs (deprecated, use accelerator instead)
        device: Device to use (deprecated, use accelerator instead)
        experiment_dir: Experiment directory
        global_step: Current training step
        autocast_kwargs: Autocast configuration
        metric_batch_size: Batch size for metric computation (on rank 0)
        reference_npz_path: Optional path to existing reference NPZ file
        accelerator: Accelerator instance (preferred over rank/world_size/device)

    Returns:
        Dictionary of metrics (only on rank 0, None on other ranks)
    """
    # Support both Accelerator and legacy DDP interface
    if accelerator is not None:
        rank = accelerator.process_index
        world_size = accelerator.num_processes
        device = accelerator.device
        is_main_process = accelerator.is_main_process
        wait_for_everyone = accelerator.wait_for_everyone
    else:
        # Legacy DDP mode
        import torch.distributed as dist
        if rank is None or world_size is None or device is None:
            raise ValueError("Either accelerator must be provided, or rank/world_size/device must be provided")
        is_main_process = (rank == 0)
        wait_for_everyone = dist.barrier

    # model.eval()
    # Save shard NPZ
    temp_dir = os.path.join(experiment_dir, "eval_npzs")
    if is_main_process:
        print(f"\n[Eval] Starting distributed sampling evaluation at step {global_step}")
        os.makedirs(temp_dir, exist_ok=True)

    # Wait for rank 0 to create the directory before other ranks try to save
    wait_for_everyone()
    # print(f"[Rank {rank}] Starting sampling...")
    # Each rank processes its shard
    N = min(len(val_dataset), num_samples)
    chunk = N // world_size

    if rank < world_size - 1:
        start = rank * chunk
        end   = (rank + 1) * chunk
    else:
        # Last rank takes the remainder (and handles N < world_size gracefully)
        start = rank * chunk
        end   = N

    rank_indices = list(range(start, end))
    subset = Subset(val_dataset, rank_indices)
    loader = DataLoader(
        subset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True,
        drop_last=False,
    )

    # Reconstruct images on this rank
    generations = []
    iterator = tqdm(loader, desc=f"[Rank {rank}] Sampling", file=sys.stdout) if is_main_process else loader

    with torch.inference_mode():
        for _, label in iterator: # don't actually need images at sampling time
            n = label.size(0)
            z = torch.randn(n, *latent_size, device = device)
            y = label.to(device)
            if use_guidance:
                z = torch.cat([z, z], dim=0)
                y_null = torch.full((n,), null_label, device=device)
                y = torch.cat([y, y_null], dim=0)
            model_kwargs = dict(y=y, **additional_model_kwargs)
            with autocast(**autocast_kwargs):
                samples = sample_fn(z, model_fn, **model_kwargs)[-1]
                if use_guidance:
                    samples = samples.chunk(2, dim = 0)
                if callable(pr) and not hasattr(pr, 'decode'):
                    samples = pr(samples)
                else:
                    samples = pr.decode(samples)
                if hasattr(samples, 'x_rec'):
                    samples = samples.x_rec
                samples = samples.clamp(0, 1)
            gen_np = samples.mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
            for img in gen_np:
                generations.append(img)

    generations = np.stack(generations)
    shard_path = os.path.join(temp_dir, f"gen_{global_step:07d}_{rank:02d}.npz")
    np.savez(shard_path, arr_0=generations)

    if is_main_process:
        print(f"[Rank {rank}] Saved {len(generations)} generation to {shard_path}")

    # Wait for all ranks to finish generation
    wait_for_everyone()

    # Rank 0 computes metrics
    metrics = None
    if is_main_process:
        # Combine all generation shards
        all_gens = []
        for r in range(world_size):
            shard_file = os.path.join(temp_dir, f"gen_{global_step:07d}_{r:02d}.npz")
            shard_data = np.load(shard_file)["arr_0"]
            all_gens.append(shard_data)

        combined_recons = np.concatenate(all_gens, axis=0)[:num_samples]
        print(f"[Eval] Combined generation NPZ shape: {combined_recons.shape}")

        # Load reference NPZ
        ref_npz_path = reference_npz_path

        if not os.path.exists(ref_npz_path):
            raise FileNotFoundError(f"Reference NPZ not found at {ref_npz_path}")

        ref_stats = np.load(ref_npz_path)
        print(f"[Eval] Loaded reference NPZ from {ref_npz_path}")

        # Compute metrics
        print("[Eval] Computing metrics...")
        metrics = compute_generation_metrics(
            ref_stats,
            combined_recons,
            device,
            metric_batch_size,
        )

        # Print results
        print(f"[Eval] Step {global_step} Metrics:")
        for key, value in metrics.items():
            print(f"  {key}: {value:.6f}")

        # Cleanup reconstruction shards
        for r in range(world_size):
            shard_file = os.path.join(temp_dir, f"gen_{global_step:07d}_{r:02d}.npz")
            if os.path.exists(shard_file):
                os.remove(shard_file)

    wait_for_everyone()
    return metrics
@torch.no_grad()
def evaluate_reconstruction_distributed(
    model,
    val_dataset,
    num_samples: int,
    batch_size: int,
    rank: Union[int, None] = None,
    world_size: Union[int, None] = None,
    device: torch.device = None,
    experiment_dir: str = None,
    global_step: int = None,
    autocast_kwargs: dict = None,
    metric_batch_size: int = 128,
    reference_npz_path: Optional[str] = None,
    metrics_to_compute: Optional[list] = ("psnr", "ssim", "rfid"),
    accelerator: Union[Accelerator, None] = None,
    save_recon_dir: Optional[str] = None,
    image_features_mode: Optional[str] = None,
    encoder_cos_sim_model: Optional[str] = None,
    psnr_resize: Optional[int] = None,
    l1_resize: Optional[int] = None,
) -> Optional[Dict[str, float]]:
    """
    Evaluate reconstruction metrics using all GPUs in a distributed manner.

    Args:
        model: Model to evaluate (should be in eval mode)
        val_dataset: Validation dataset
        batch_size: Batch size per GPU for reconstruction
        rank: Current GPU rank (deprecated, use accelerator instead)
        world_size: Total number of GPUs (deprecated, use accelerator instead)
        device: Device to use (deprecated, use accelerator instead)
        experiment_dir: Experiment directory
        global_step: Current training step
        autocast_kwargs: Autocast configuration
        metric_batch_size: Batch size for metric computation (on rank 0)
        reference_npz_path: Optional path to existing reference NPZ file
        accelerator: Accelerator instance (preferred over rank/world_size/device)

    Returns:
        Dictionary of metrics (only on rank 0, None on other ranks)
    """
    # Support both Accelerator and legacy DDP interface
    if accelerator is not None:
        rank = accelerator.process_index
        world_size = accelerator.num_processes
        device = accelerator.device
        is_main_process = accelerator.is_main_process
        wait_for_everyone = accelerator.wait_for_everyone
    else:
        # Legacy DDP mode
        import torch.distributed as dist
        if rank is None or world_size is None or device is None:
            raise ValueError("Either accelerator must be provided, or rank/world_size/device must be provided")
        is_main_process = (rank == 0)
        wait_for_everyone = dist.barrier

    model.eval()
    # Save shard NPZ
    temp_dir = os.path.join(experiment_dir, "eval_npzs")
    if is_main_process:
        print(f"\n[Eval] Starting distributed reconstruction evaluation at step {global_step}")
        os.makedirs(temp_dir, exist_ok=True)
    # Wait for rank 0 to create the directory before other ranks try to save
    wait_for_everyone()
    # print(f"[Rank {rank}] Starting reconstruction...")
    # Each rank processes its shard
    N = min(len(val_dataset), num_samples)
    chunk = N // world_size

    if rank < world_size - 1:
        start = rank * chunk
        end   = (rank + 1) * chunk
    else:
        # Last rank takes the remainder (and handles N < world_size gracefully)
        start = rank * chunk
        end   = N

    rank_indices = list(range(start, end))
    subset = Subset(val_dataset, rank_indices)
    loader = DataLoader(
        subset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True,
        drop_last=False,
    )

    # Reconstruct images on this rank, keep originals for distributed metrics
    reconstructions = []
    originals = []
    iterator = tqdm(loader, desc=f"[Rank {rank}] Reconstructing", file=sys.stdout) if is_main_process else loader

    with torch.inference_mode():
        for images, _ in iterator:
            images = images.to(device, non_blocking=True)
            # Save originals before model forward
            orig_np = images.clamp(0, 1).mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
            with autocast(**autocast_kwargs):
                if image_features_mode is not None:
                    recon = model(images, image_features_mode=image_features_mode)
                else:
                    recon = model(images)

            # Handle DecodeOutput from split_adaln decoder
            if hasattr(recon, 'x_rec'):
                recon = recon.x_rec
            # Convert to numpy uint8 [H, W, C]
            recon = recon.clamp(0, 1)
            recon_np = recon.mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()

            for img in recon_np:
                reconstructions.append(img)
            for img in orig_np:
                originals.append(img)

    reconstructions = np.stack(reconstructions)
    originals = np.stack(originals)

    # Save shard for backward compat
    shard_path = os.path.join(temp_dir, f"recon_{global_step:07d}_{rank:02d}.npz")
    np.savez(shard_path, arr_0=reconstructions)
    if is_main_process:
        print(f"[Rank {rank}] Saved {len(reconstructions)} reconstructions to {shard_path}")

    # Wait for all ranks to finish reconstruction
    wait_for_everyone()

    # ── Distributed metrics: each rank computes on its own subset ──────────
    import torch.distributed as dist
    from .fid import compute_inception_features, calculate_rfid_from_features
    local_n = len(reconstructions)

    # PSNR / SSIM: compute locally then all_reduce weighted sum
    scalar_metrics_to_compute = [m for m in (metrics_to_compute or []) if m in ('psnr', 'ssim', 'lpips')]
    local_scalar = {}
    if scalar_metrics_to_compute:
        if is_main_process:
            print(f"[Eval] Rank {rank}: computing {scalar_metrics_to_compute} on {local_n} images...")
        local_scalar = compute_reconstruction_metrics(
            originals, reconstructions, device, metric_batch_size,
            metrics_to_compute=scalar_metrics_to_compute,
            disable_bar=is_main_process,
            encoder_model=encoder_cos_sim_model,
            psnr_resize=psnr_resize,
            l1_resize=l1_resize,
        )

    # Weighted all_reduce: sum(metric * count) / total_count
    global_scalar = {}
    count_t = torch.tensor(float(local_n), device=device)
    distributed = world_size > 1 and dist.is_available() and dist.is_initialized()
    if distributed:
        dist.all_reduce(count_t, op=dist.ReduceOp.SUM)
    for k, v in local_scalar.items():
        val_t = torch.tensor(v * local_n, dtype=torch.float64, device=device)
        if distributed:
            dist.all_reduce(val_t, op=dist.ReduceOp.SUM)
        global_scalar[k] = (val_t / count_t).item()

    # rFID: each rank computes inception features, all_gather, rank 0 computes FID
    global_rfid = {}
    if 'rfid' in (metrics_to_compute or []):
        if is_main_process:
            print(f"[Eval] Rank {rank}: computing inception features on {local_n} images...")
        local_recon_feats = compute_inception_features(reconstructions, metric_batch_size, device)
        local_ref_feats   = compute_inception_features(originals,       metric_batch_size, device)

        recon_t = torch.from_numpy(local_recon_feats).to(device)
        ref_t   = torch.from_numpy(local_ref_feats).to(device)

        if distributed:
            all_recon = [torch.zeros_like(recon_t) for _ in range(world_size)]
            all_ref   = [torch.zeros_like(ref_t)   for _ in range(world_size)]
            dist.all_gather(all_recon, recon_t)
            dist.all_gather(all_ref, ref_t)
        else:
            all_recon = [recon_t]
            all_ref = [ref_t]

        if is_main_process:
            combined_recon_feats = torch.cat(all_recon, dim=0).cpu().double().numpy()
            combined_ref_feats   = torch.cat(all_ref,   dim=0).cpu().double().numpy()
            rfid = calculate_rfid_from_features(combined_recon_feats, combined_ref_feats)
            global_rfid = {'rfid': rfid}

    metrics = None
    if is_main_process:
        metrics = {**global_scalar, **global_rfid}
        # Persist the all-reduced count so release validation proves how many
        # examples were actually processed rather than trusting CLI intent.
        metrics["_evaluated_samples"] = int(count_t.item())
        print(f"[Eval] Step {global_step} Metrics:")
        for key, value in metrics.items():
            print(f"  {key}: {value:.6f}")

        # Optionally save reconstructed images
        if save_recon_dir:
            import threading
            recons_copy = np.copy(reconstructions)
            save_dir = save_recon_dir
            def _save_recon_images():
                os.makedirs(save_dir, exist_ok=True)
                for i, img in enumerate(recons_copy):
                    Image.fromarray(img).save(os.path.join(save_dir, f"recon_{i:06d}.png"))
                print(f"[Eval] Saved {len(recons_copy)} reconstruction images to {save_dir}")
            threading.Thread(target=_save_recon_images, daemon=False).start()

        # Cleanup shards
        for r in range(world_size):
            shard_file = os.path.join(temp_dir, f"recon_{global_step:07d}_{r:02d}.npz")
            if os.path.exists(shard_file):
                os.remove(shard_file)

    wait_for_everyone()
    model.train()

    return metrics
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--ref-img", type=str, default="samples/imagenet-256-val.npz")
    parser.add_argument("--rec-img", type=str, default="samples/sdvae-ft-mse-f8d4.npz")
    parser.add_argument("--bs", type=int, default=128)
    args = parser.parse_args()

    # Load images
    device = "cuda"
    ref_img = np.load(args.ref_img)["arr_0"]
    rec_img = np.load(args.rec_img)["arr_0"]
    print(f"Loaded images: ref: {ref_img.shape}, rec: {rec_img.shape}")

    psnr = calculate_psnr(ref_img, rec_img, args.bs, device)
    print(f"PSNR: {psnr:.6f}")
    lpips = calculate_lpips(ref_img, rec_img, args.bs, device)
    print(f"LPIPS: {lpips:.6f}")
    ssim_val = calculate_ssim(ref_img, rec_img, args.bs, device)
    print(f"SSIM: {ssim_val:.6f}")
    rfid = calculate_rfid(ref_img, rec_img, args.bs, device)
    print(f"rFID: {rfid:.6f}")