Spaces:
Running on Zero
Running on Zero
| # 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) | |