patch-reparameterization/src/sample_ddp.py

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

"""
Samples a large number of images from a pre-trained stage-2 model using DDP and
stores results for downstream metrics. For single-device sampling, use sample.py.

When --reference-npz is provided, FID is computed after sampling completes.
Reference NPZ can be either:
  - Precomputed stats (keys: mu, sigma) from Inception features
  - Raw images (key: arr_0) - FID will be computed via torch_fidelity
"""
import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import argparse
import json
import logging
import math
from typing import Callable, Optional

logger = logging.getLogger(__name__)

import numpy as np
import torch
import torch.distributed as dist
from PIL import Image
from torch.cuda.amp import autocast
from tqdm import tqdm
from pathlib import Path
from omegaconf import OmegaConf
from utils.config_utils import load_config  # noqa: F401  (registers oc.env resolver)
from utils.model_utils import instantiate_from_config
from stage1 import PatchReparam
from stage2.models import Stage2ModelProtocol
from stage2.transport import create_transport, Sampler
from utils.train_utils import parse_configs, center_crop_arr
from eval.fid import calculate_gfid, calculate_rfid
from torchvision.datasets import ImageFolder
from torchvision import transforms


def _safe_load_checkpoint(path: str, map_location: str = "cpu"):
    try:
        return torch.load(path, map_location=map_location, weights_only=False)
    except TypeError:
        # For older PyTorch versions without weights_only arg
        return torch.load(path, map_location=map_location)


def _load_patch_reparam_checkpoint(pr: PatchReparam, ckpt_path: str, use_ema: bool = True) -> None:
    """
    Load checkpoint weights into PatchReparam, choosing between EMA and model weights.
    Mirrors the logic in eval_inversion_error._load_patch_reparam_checkpoint.
    """
    logger.info("Loading PatchReparam checkpoint from %s (use_ema=%s)", ckpt_path, use_ema)
    ckpt = _safe_load_checkpoint(ckpt_path, map_location="cpu")
    state_dict = ckpt.get("ema" if use_ema else "model")
    if state_dict is None:
        raise KeyError(f"Checkpoint must contain 'ema' or 'model'. Got keys: {list(ckpt.keys())}")

    clean_state = {}
    for k, v in state_dict.items():
        clean_k = k
        for prefix in ["module._orig_mod.", "_orig_mod.", "module."]:
            if clean_k.startswith(prefix):
                clean_k = clean_k[len(prefix):]
                break
        clean_state[clean_k] = v

    missing, unexpected = pr.load_state_dict(clean_state, strict=False)
    if missing:
        logger.warning("PatchReparam checkpoint missing %d keys (first 5): %s...", len(missing), missing[:5])
    if unexpected:
        logger.warning("PatchReparam checkpoint unexpected %d keys (first 5): %s...", len(unexpected), unexpected[:5])
    logger.info("PatchReparam checkpoint loaded successfully.")


def create_npz_from_sample_folder(sample_dir, num=50_000):
    """
    Builds a single .npz file from a folder of .png samples.
    """
    samples = []
    for i in tqdm(range(num), desc="Building .npz file from samples"):
        sample_pil = Image.open(f"{sample_dir}/{i:06d}.png")
        sample_np = np.asarray(sample_pil).astype(np.uint8)
        samples.append(sample_np)
    samples = np.stack(samples)
    assert samples.shape == (num, samples.shape[1], samples.shape[2], 3)
    npz_path = f"{sample_dir}.npz"
    np.savez(npz_path, arr_0=samples)
    logger.info("Saved .npz file to %s [shape=%s].", npz_path, samples.shape)
    return npz_path


def compute_fid_from_samples(
    gen_npz_path: str,
    reference_npz_path: str,
    device: torch.device,
    batch_size: int = 128,
) -> float:
    """
    Compute FID between generated samples and reference.
    Reference NPZ can be:
      - Precomputed stats: keys 'mu', 'sigma' (Inception moments)
      - Raw images: key 'arr_0'
    """
    gen_data = np.load(gen_npz_path)
    gen_arr = gen_data["arr_0"]

    ref_data = np.load(reference_npz_path)
    ref_files = list(ref_data.files)

    device_str = "cuda" if device.type == "cuda" else "cpu"

    if "mu" in ref_files and "sigma" in ref_files:
        fid = calculate_gfid(gen_arr, ref_data, batch_size=batch_size, device=device_str)
    elif "arr_0" in ref_files:
        ref_arr = ref_data["arr_0"]
        fid = calculate_rfid(gen_arr, ref_arr, bs=batch_size, device=device_str)
    else:
        raise KeyError(
            f"Reference NPZ must have 'mu'/'sigma' (precomputed stats) or 'arr_0' (raw images). "
            f"Got keys: {ref_files}"
        )
    return float(fid)

def build_label_sampler(
    sampling_mode: str,
    num_classes: int,
    num_fid_samples: int,
    total_samples: int,
    samples_needed_this_device: int,
    batch_size: int,
    device: torch.device,
    rank: int,
    iterations: int,
    seed: int,
    eval_data: Optional[str] = None,
    image_size: int = 256,
) -> Callable[[int], torch.Tensor]:
    """Create a callable that returns a batch of labels for the given step index.
    When mode='dataset', labels follow val dataset order (same as training eval)."""

    if sampling_mode == "dataset":
        if not eval_data or not os.path.isdir(eval_data):
            raise ValueError(
                "dataset mode requires --eval-data (or config eval.data_path) pointing to ImageFolder val dir."
            )
        dataset = ImageFolder(
            eval_data,
            transform=transforms.Compose([
                transforms.Lambda(lambda x: center_crop_arr(x, image_size)),
                transforms.ToTensor(),
            ]),
        )
        N = min(len(dataset), num_fid_samples)
        world_size = total_samples // samples_needed_this_device
        chunk = total_samples // world_size  # == samples_needed_this_device
        start = rank * chunk
        end = min(start + chunk, total_samples)
        rank_indices = list(range(start, end))
        device_pool = torch.tensor([dataset.targets[i] for i in rank_indices], dtype=torch.long)
        needed = iterations * batch_size
        if len(device_pool) < needed:
            gen = torch.Generator().manual_seed(seed + rank)
            device_pool = torch.cat([
                device_pool,
                torch.randint(0, num_classes, (needed - len(device_pool),), generator=gen),
            ])
        device_pool = device_pool[:needed].view(iterations, batch_size)

        def dataset_sampler(step_idx: int) -> torch.Tensor:
            return device_pool[step_idx].to(device)

        return dataset_sampler

    if sampling_mode == "random":
        def random_sampler(_step_idx: int) -> torch.Tensor:
            return torch.randint(0, num_classes, (batch_size,), device=device)

        return random_sampler

    if sampling_mode == "equal":
        if num_fid_samples % num_classes != 0:
            raise ValueError(
                f"Equal label sampling requires num_fid_samples ({num_fid_samples}) to be divisible by num_classes ({num_classes})."
            )

        labels_per_class = num_fid_samples // num_classes
        base_pool = torch.arange(num_classes, dtype=torch.long).repeat_interleave(labels_per_class)

        generator = torch.Generator()
        generator.manual_seed(seed)
        permutation = torch.randperm(base_pool.numel(), generator=generator)
        base_pool = base_pool[permutation]

        if total_samples > num_fid_samples:
            tail = torch.randint(0, num_classes, (total_samples - num_fid_samples,), generator=generator)
            global_pool = torch.cat([base_pool, tail], dim=0)
        else:
            global_pool = base_pool

        start = rank * samples_needed_this_device
        end = start + samples_needed_this_device
        device_pool = global_pool[start:end]
        device_pool = device_pool.view(iterations, batch_size)

        def equal_sampler(step_idx: int) -> torch.Tensor:
            labels = device_pool[step_idx]
            return labels.to(device)

        return equal_sampler
    raise ValueError(f"Unknown label sampling mode: {sampling_mode}")

def main(args):
    """Run sampling with distributed execution."""
    if not torch.cuda.is_available():
        raise RuntimeError("Sampling with DDP requires at least one GPU. Use sample.py for single-device usage.")

    torch.backends.cuda.matmul.allow_tf32 = args.tf32
    torch.backends.cudnn.allow_tf32 = args.tf32
    torch.set_grad_enabled(False)

    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device_idx = rank % torch.cuda.device_count()
    torch.cuda.set_device(device_idx)
    device = torch.device("cuda", device_idx)

    if rank == 0:
        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    else:
        logger.addHandler(logging.NullHandler())
        logger.propagate = False

    seed = args.global_seed * world_size + rank
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    if rank == 0:
        logger.info("Starting rank=%s, seed=%s, world_size=%s.", rank, seed, world_size)

    use_bf16 = args.precision == "bf16"
    if use_bf16 and not torch.cuda.is_bf16_supported():
        raise ValueError("Requested bf16 precision, but the current CUDA device does not support bfloat16.")
    autocast_kwargs = dict(dtype=torch.bfloat16, enabled=use_bf16)
    cfg = OmegaConf.load(args.config)
    if args.ckpt is not None:
        OmegaConf.update(cfg, "stage_2.ckpt", args.ckpt, force_add=True)
    if args.cfg_scale is not None:
        OmegaConf.update(cfg, "guidance.scale", args.cfg_scale, force_add=True)
    if args.precomputed_latents_dir is not None:
        val = None if args.precomputed_latents_dir.lower() == "none" else args.precomputed_latents_dir
        OmegaConf.update(cfg, "misc.precomputed_latents_dir", val, force_add=True)
    pr_config, model_config, transport_config, sampler_config, guidance_config, misc, _, _ = parse_configs(cfg)
    # ── VAE mode: if vae_decoder is specified, use external VAE instead of PatchReparam ──
    vae_decoder_cfg = cfg.get("vae_decoder", None)
    use_external_vae = vae_decoder_cfg is not None

    if not use_external_vae and (pr_config is None or model_config is None):
        raise ValueError("Config must provide both stage_1 and stage_2 entries (or use vae_decoder).")
    if use_external_vae and model_config is None:
        raise ValueError("Config must provide stage_2 section.")

    misc = {} if misc is None else dict(misc)

    latent_size = tuple(int(dim) for dim in misc.get("latent_size", (768, 16, 16)))
    shift_dim = misc.get("time_dist_shift_dim", math.prod(latent_size))
    shift_base = misc.get("time_dist_shift_base", 4096)
    time_dist_shift = math.sqrt(shift_dim / shift_base)
    if rank == 0:
        logger.info("Using time_dist_shift=%.4f.", time_dist_shift)

    # ------------------------------------------------------------------
    # Optional DiT latent caching / reuse
    # ------------------------------------------------------------------
    precomputed_latents_dir = None
    save_latents = False
    latents_dir = None
    if misc is not None:
        precomputed_latents_dir = misc.get("precomputed_latents_dir", None)
        save_latents = bool(misc.get("save_latents", False))
        latents_dir = misc.get("latents_dir", None)

    external_vae = None
    vae_scaling_factor = None
    pr = None
    proj_down_weight = None
    proj_norm_mode = None
    proj_norm_mean = None
    proj_norm_std = None

    if use_external_vae:
        from diffusers import AutoencoderKL
        vae_path = vae_decoder_cfg.get("path") if hasattr(vae_decoder_cfg, "get") else str(vae_decoder_cfg)
        logger.info(f"Loading external VAE from {vae_path}")
        external_vae = AutoencoderKL.from_pretrained(vae_path).to(device).eval()
        external_vae.requires_grad_(False)
        vae_scaling_factor = external_vae.config.scaling_factor
        logger.info(f"External VAE loaded: scaling_factor={vae_scaling_factor}")

        # Optional projection (for high-dim experiment decoding)
        proj_cfg = vae_decoder_cfg.get("projection", None) if hasattr(vae_decoder_cfg, "get") else None
        if proj_cfg is not None:
            proj_cfg = OmegaConf.to_container(proj_cfg, resolve=True) if hasattr(proj_cfg, '_dict') or type(proj_cfg).__name__ == 'DictConfig' else dict(proj_cfg)
            weight_path = proj_cfg.get("weight_path", None)
            if not weight_path or not os.path.exists(weight_path):
                raise FileNotFoundError(f"Projection weight_path required for sampling: {weight_path}")
            proj_up_weight = torch.load(weight_path, map_location="cpu").to(device)
            proj_down_weight = proj_up_weight.T.contiguous()
            logger.info(f"Loaded projection matrix from {weight_path}: {proj_up_weight.shape}")

            proj_norm_mode = proj_cfg.get("norm_mode", "global")
            if proj_norm_mode == "per_dim":
                norm_stat_path = proj_cfg.get("norm_stat_path", None)
                if norm_stat_path and os.path.exists(norm_stat_path):
                    stats = torch.load(norm_stat_path, map_location="cpu")
                    proj_norm_mean = stats["mean"].to(device)
                    proj_norm_std = torch.sqrt(stats["var"].to(device) + 1e-5)
                    logger.info(f"Loaded per-dim norm stats from {norm_stat_path}")
                else:
                    raise FileNotFoundError(f"per_dim norm requires norm_stat_path: {norm_stat_path}")
    else:
        pr: PatchReparam = instantiate_from_config(pr_config).to(device)
        pr_ckpt_cfg = pr_config.get("checkpoint") if hasattr(pr_config, "get") else None
        if pr_ckpt_cfg:
            if isinstance(pr_ckpt_cfg, (str, Path)):
                ckpt_path = str(pr_ckpt_cfg)
                use_ema = True
            else:
                ckpt_path = pr_ckpt_cfg.get("path") or pr_ckpt_cfg.get("ckpt_path")
                use_ema = pr_ckpt_cfg.get("use_ema", True)
            if ckpt_path and str(ckpt_path).strip():
                _load_patch_reparam_checkpoint(pr, str(ckpt_path), use_ema=use_ema)
            else:
                logger.warning("PatchReparam checkpoint config present but path is empty.")
        else:
            logger.info("No PatchReparam checkpoint specified.")
        pr.eval()

    model: Stage2ModelProtocol = instantiate_from_config(model_config).to(device)
    model.eval()

    transport_params = {}
    if transport_config is not None:
        transport_params = dict(transport_config.get("params", {}))
    transport = create_transport(
        **transport_params,
        time_dist_shift=time_dist_shift,
    )
    sampler = Sampler(transport)

    sampler_config = {} if sampler_config is None else dict(sampler_config)
    sampler_mode = sampler_config.get("mode", "ODE")
    sampler_params = dict(sampler_config.get("params", {}))
    mode = sampler_mode.upper()
    if mode == "ODE":
        sample_fn = sampler.sample_ode(**sampler_params)
    elif mode == "SDE":
        sample_fn = sampler.sample_sde(**sampler_params)
    else:
        raise NotImplementedError(f"Invalid sampling mode {sampler_mode}.")

    guidance_config = {} if guidance_config is None else dict(guidance_config)

    def guidance_value(key: str, default: float):
        if key in guidance_config:
            return guidance_config[key]
        dashed_key = key.replace("_", "-")
        return guidance_config.get(dashed_key, default)

    guidance_scale = guidance_config.get("scale", 1.0)
    guidance_method = guidance_config.get("method", "cfg")
    t_min = guidance_value("t_min", 0.0)
    t_max = guidance_value("t_max", 1.0)

    guid_model_forward = None
    if guidance_scale > 1.0 and guidance_method == "autoguidance":
        guid_model_config = guidance_config.get("guidance_model")
        if guid_model_config is None:
            raise ValueError("Please provide a guidance model config when using autoguidance.")
        guid_model: Stage2ModelProtocol = instantiate_from_config(guid_model_config).to(device)
        guid_model.eval()
        guid_model_forward = guid_model.forward

    num_classes = int(misc.get("num_classes", 1000))
    null_label = int(misc.get("null_label", num_classes))

    model_target = model_config.get("target", "stage2")
    model_string_name = str(model_target).split(".")[-1]
    ckpt_path = model_config.get("ckpt")
    ckpt_string_name = "pretrained" if not ckpt_path else os.path.splitext(os.path.basename(str(ckpt_path)))[0]
    sampling_method = sampler_params.get("sampling_method", "na")
    num_steps = sampler_params.get("num_steps", sampler_params.get("steps", "na"))
    guidance_tag = f"cfg-{guidance_scale:.2f}"
    base_components = [model_string_name, ckpt_string_name, guidance_tag, f"bs{args.per_proc_batch_size}"]
    logger.info(f"args.per_proc_batch_size{args.per_proc_batch_size}")
    if mode == "ODE":
        detail_components = [mode, str(num_steps), str(sampling_method), args.precision]
    else:
        diffusion_form = sampler_params.get("diffusion_form", "na")
        last_step = sampler_params.get("last_step", "na")
        last_step_size = sampler_params.get("last_step_size", "na")
        detail_components = [mode, str(num_steps), str(sampling_method), str(diffusion_form), str(last_step), str(last_step_size), args.precision]
    folder_name = "-".join(component.replace(os.sep, "-") for component in base_components + detail_components)
    possible_folder_name = os.environ.get('SAVE_FOLDER', None)
    if possible_folder_name:
        sample_folder_dir = os.path.join(args.sample_dir, possible_folder_name)
    else:
        sample_folder_dir = os.path.join(args.sample_dir, folder_name)
    # Parent dir for result files (one level above image folder)
    result_parent_dir = Path(sample_folder_dir).parent
    folder_base = Path(sample_folder_dir).name

    if rank == 0:
        os.makedirs(sample_folder_dir, exist_ok=True)
        logger.info("Saving .png samples at %s", sample_folder_dir)

        # Determine directory to read/write DiT latents.
        if latents_dir is None:
            latents_dir = os.path.join(sample_folder_dir, "latents")

        # Save config and startup args for reproducibility (in parent dir, not inside image folder)
        config_save_path = result_parent_dir / f"{folder_base}_sample_config.yaml"
        OmegaConf.save(cfg, config_save_path)
        logger.info("Config saved to %s", config_save_path)
        args_dict = vars(args)
        args_save_path = result_parent_dir / f"{folder_base}_sample_args.json"
        with open(args_save_path, "w", encoding="utf-8") as f:
            json.dump(args_dict, f, indent=2, ensure_ascii=False)
        logger.info("Args saved to %s", args_save_path)
    # Broadcast latents_dir from rank 0 so all ranks agree.
    latents_dir_obj = {"latents_dir": latents_dir}
    obj_list = [latents_dir_obj]
    dist.broadcast_object_list(obj_list, src=0)
    latents_dir = obj_list[0]["latents_dir"]
    if save_latents and latents_dir is not None and rank == 0:
        os.makedirs(latents_dir, exist_ok=True)
    dist.barrier()

    # If npz already exists with enough samples, skip sampling to avoid redundant computation
    npz_path = f"{sample_folder_dir}.npz"
    skip_sampling = False
    if rank == 0:
        logger.info("NPZ path: %s, isfile: %s", npz_path, os.path.isfile(npz_path))
        if os.path.isfile(npz_path):
            try:
                data = np.load(npz_path)
                arr = data["arr_0"]
                logger.info("NPZ already exists at %s with %d samples.", npz_path, arr.shape[0])
                if arr.shape[0] >= args.num_fid_samples:
                    skip_sampling = True
                    logger.info("NPZ already exists at %s with %d samples, skipping sampling.", npz_path, arr.shape[0])
            except Exception as e:
                logger.exception("Error loading NPZ: %s", e)
                # Fall back to existing PNG samples if they are complete enough.
                try:
                    existing_files = [
                        name
                        for name in os.listdir(sample_folder_dir)
                        if os.path.isfile(os.path.join(sample_folder_dir, name)) and name.endswith(".png")
                    ]
                    existing_set = set(existing_files)
                    expected_files = {f"{i:06d}.png" for i in range(args.num_fid_samples)}
                    if expected_files.issubset(existing_set):
                        logger.info(
                            "Found %d PNG samples and a broken NPZ. Rebuilding NPZ from existing images and skipping sampling.",
                            len(existing_files),
                        )
                        create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)
                        skip_sampling = True
                    else:
                        logger.info(
                            "Existing PNG samples are incomplete (%d found, need at least %d). Proceeding to resample.",
                            len(existing_files),
                            args.num_fid_samples,
                        )
                except Exception as e2:
                    logger.exception("Error while checking/rebuilding NPZ from PNGs: %s", e2)
    skip_tensor = torch.tensor(1 if skip_sampling else 0, device=device)
    dist.broadcast(skip_tensor, 0)
    skip_sampling = skip_tensor.item() == 1
    logger.info("Skip sampling (existing images/NPZ): %s", skip_sampling)
    n = args.per_proc_batch_size
    global_batch_size = n * world_size
    existing = [name for name in os.listdir(sample_folder_dir) if (os.path.isfile(os.path.join(sample_folder_dir, name)) and name.endswith(".png"))]
    num_samples = len(existing)
    total_samples = int(math.ceil(args.num_fid_samples / global_batch_size) * global_batch_size)
    if rank == 0:
        logger.info("Total number of images that will be sampled: %d", total_samples)
    if total_samples % world_size != 0:
        raise ValueError("Total samples must be divisible by world size.")
    samples_needed_this_gpu = total_samples // world_size
    if samples_needed_this_gpu % n != 0:
        raise ValueError("Per-rank sample count must be divisible by the per-GPU batch size.")
    iterations = samples_needed_this_gpu // n
    # Resume: align to global_batch_size boundary to preserve label consistency
    total = (num_samples // global_batch_size) * global_batch_size
    start_step = total // global_batch_size
    remaining_iterations = iterations - start_step
    if rank == 0 and start_step > 0:
        logger.info("Resuming from step %d (%d images exist, aligned to %d), %d steps remaining.",
                     start_step, num_samples, total, remaining_iterations)
    pbar = tqdm(range(remaining_iterations)) if rank == 0 else range(remaining_iterations)

    eval_data = args.eval_data
    if eval_data is None and cfg is not None:
        eval_section = cfg.get("eval", None)
        if eval_section is not None:
            eval_data = eval_section.get("data_path", None)
    label_sampler = build_label_sampler(
        args.label_sampling,
        num_classes,
        args.num_fid_samples,
        total_samples,
        samples_needed_this_gpu,
        n,
        device,
        rank,
        iterations,
        args.global_seed,
        eval_data=eval_data if args.label_sampling == "dataset" else None,
        image_size=args.image_size,
    )

    # Decide whether to reuse precomputed DiT latents (skip DiT forward) or run DiT.
    use_precomputed_latents = False
    # if precomputed_latents_dir is None:
    #     # If a dedicated precomputed directory is not provided, reuse latents_dir if present.
    #     precomputed_latents_dir = latents_dir
    if (
        not skip_sampling
        and precomputed_latents_dir is not None
        and os.path.isdir(precomputed_latents_dir)
    ):
        if rank == 0:
            logger.info("Using precomputed DiT latents from %s.", precomputed_latents_dir)
        use_precomputed_latents = True
    use_precomputed_tensor = torch.tensor(1 if use_precomputed_latents else 0, device=device)
    dist.broadcast(use_precomputed_tensor, 0)
    use_precomputed_latents = use_precomputed_tensor.item() == 1
    if rank == 0:
        logger.info("Use precomputed DiT latents: %s", use_precomputed_latents)

    using_cfg = guidance_scale > 1.0
    if not skip_sampling:
        for step_idx in pbar:
            actual_step = start_step + step_idx
            with autocast(**autocast_kwargs):
                if use_precomputed_latents:
                    # Load precomputed DiT latents and only run the decoder.
                    latents = []
                    for local_idx in range(n):
                        index = local_idx * world_size + rank + total
                        latent_path = os.path.join(precomputed_latents_dir, f"{index:06d}.pt")
                        if not os.path.isfile(latent_path):
                            raise FileNotFoundError(
                                f"Expected latent file {latent_path} but it does not exist. "
                                "Please ensure precomputed latents cover all required indices."
                            )
                        latent = torch.load(latent_path, map_location=device)
                        latents.append(latent)
                    latents = torch.stack(latents, dim=0).to(device=device)
                else:
                    # Run DiT to generate latents.
                    z = torch.randn(n, *latent_size, device=device)
                    y = label_sampler(actual_step)

                    model_kwargs = dict(y=y)
                    model_fn = model.forward

                    if using_cfg:
                        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,
                            cfg_scale=guidance_scale,
                            cfg_interval=(t_min, t_max),
                        )
                        if guidance_method == "autoguidance":
                            if guid_model_forward is None:
                                raise RuntimeError("Guidance model forward is not initialized.")
                            model_kwargs["additional_model_forward"] = guid_model_forward
                            model_fn = model.forward_with_autoguidance
                        else:
                            model_fn = model.forward_with_cfg
                    latents = sample_fn(z, model_fn, **model_kwargs)[-1]
                    if using_cfg:
                        latents, _ = latents.chunk(2, dim=0)

                    # Optionally cache DiT latents for future runs.
                    if save_latents and latents_dir is not None:
                        for local_idx, latent in enumerate(latents):
                            index = local_idx * world_size + rank + total
                            latent_path = os.path.join(latents_dir, f"{index:06d}.pt")
                            torch.save(latent.to("cpu"), latent_path)

                # Decode latents into images.
                if use_external_vae:
                    # Reverse projection if used
                    if proj_down_weight is not None:
                        if proj_norm_mode == "per_dim" and proj_norm_mean is not None:
                            latents = latents * proj_norm_std + proj_norm_mean
                        latents = torch.nn.functional.conv2d(latents, proj_down_weight[:, :, None, None])
                    decoded = external_vae.decode(latents / vae_scaling_factor).sample
                    decoded = ((decoded + 1) / 2).clamp(0, 1)
                else:
                    decoded = pr.decode(latents)
                    mid_x_rec = getattr(decoded, 'mid_x_rec', None)
                    if hasattr(decoded, 'x_rec'):
                        decoded = decoded.x_rec
                    decoded = decoded.clamp(0, 1)
                if mid_x_rec is not None:
                    mid_decoded = mid_x_rec.clamp(0, 1)
                    mid_decoded = (
                        mid_decoded.mul(255)
                        .permute(0, 2, 3, 1)
                        .to("cpu", dtype=torch.uint8)
                        .numpy()
                    )
                else:
                    mid_decoded = None
                decoded = (
                    decoded.mul(255)
                    .permute(0, 2, 3, 1)
                    .to("cpu", dtype=torch.uint8)
                    .numpy()
                )

            for local_idx, sample in enumerate(decoded):
                index = local_idx * world_size + rank + total
                Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png")
                if mid_decoded is not None:
                    mid_dir = sample_folder_dir + "_mid"
                    os.makedirs(mid_dir, exist_ok=True)
                    Image.fromarray(mid_decoded[local_idx]).save(f"{mid_dir}/{index:06d}.png")

            total += global_batch_size
            dist.barrier()

    dist.barrier()
    if rank == 0:
        if not skip_sampling:
            npz_path = create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)
        # else: npz_path already set and file exists

        reference_npz = args.reference_npz
        if reference_npz is None and cfg is not None:
            eval_section = cfg.get("eval", None)
            if eval_section is not None:
                reference_npz = eval_section.get("reference_npz_path", None)

        if reference_npz is not None and not args.skip_fid:
            logger.info("Computing FID against reference...")
            fid_value = compute_fid_from_samples(
                npz_path,
                reference_npz,
                device,
                batch_size=args.fid_batch_size,
            )
            logger.info("FID: %.6f", fid_value)

            result_path = result_parent_dir / f"{folder_base}_fid_result.json"
            with open(result_path, "w", encoding="utf-8") as f:
                json.dump(
                    {
                        "fid": fid_value,
                        "reference_npz": reference_npz,
                        "num_samples": args.num_fid_samples,
                        "sample_folder": sample_folder_dir,
                    },
                    f,
                    indent=2,
                )
            logger.info("FID result saved to %s", result_path)

        logger.info("Done.")
    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, required=True, help="Path to the config file.")
    parser.add_argument("--sample-dir", type=str, default="samples")
    parser.add_argument("--per-proc-batch-size", type=int, default=125)
    parser.add_argument("--num-fid-samples", type=int, default=50_000)
    parser.add_argument("--global-seed", type=int, default=0)
    parser.add_argument("--precision", type=str, choices=["fp32", "bf16"], default="fp32")
    parser.add_argument("--tf32", action=argparse.BooleanOptionalAction, default=True,
                        help="Enable TF32 matmuls (Ampere+). Disable if deterministic results are required.")
    parser.add_argument(
        "--label-sampling",
        type=str,
        choices=["random", "equal", "dataset"],
        default="equal",
        help="Choose how to sample class labels. Use 'dataset' to match training eval (val order).",
    )
    parser.add_argument(
        "--eval-data",
        type=str,
        default=None,
        help="Path to eval ImageFolder (val dir). Required for --label-sampling dataset. Overrides config eval.data_path.",
    )
    parser.add_argument(
        "--image-size",
        type=int,
        default=256,
        help="Image size for center crop when using dataset mode.",
    )
    parser.add_argument(
        "--reference-npz",
        type=str,
        default=None,
        help="Path to reference NPZ for FID computation. Can be precomputed stats (mu/sigma) or raw images (arr_0). "
        "Overrides config eval.reference_npz_path if set.",
    )
    parser.add_argument(
        "--fid-batch-size",
        type=int,
        default=128,
        help="Batch size for FID computation.",
    )
    parser.add_argument(
        "--skip-fid",
        action="store_true",
        help="Generate PNG/NPZ artifacts without computing FID in this process.",
    )
    parser.add_argument(
        "--ckpt",
        type=str,
        default=None,
        help="Override stage_2.ckpt in config with this checkpoint path.",
    )
    parser.add_argument(
        "--cfg-scale",
        type=float,
        default=None,
        help="Override guidance.scale in config.",
    )
    parser.add_argument(
        "--precomputed-latents-dir",
        type=str,
        default=None,
        help="Override misc.precomputed_latents_dir in config. Pass 'none' to disable precomputed latents.",
    )

    args = parser.parse_args()
    main(args)