#!/usr/bin/env python3 """ DDP 多卡采样脚本（单路径，不做 dual-compare，不保存 t_c 中间态图）。 用法（4 卡示例）： torchrun --nproc_per_node=4 sample_from_checkpoint_ddp.py \ --ckpt exps/jsflow-experiment/checkpoints/0290000.pt \ --out-dir ./my_samples_ddp \ --num-images 50000 \ --batch-size 16 \ --t-c 0.75 --steps-before-tc 100 --steps-after-tc 5 \ --sampler em_image_noise_before_tc """ from __future__ import annotations import argparse import math import os import sys import types import numpy as np import torch import torch.distributed as dist from diffusers.models import AutoencoderKL from PIL import Image from tqdm import tqdm from models.sit import SiT_models from samplers import ( euler_maruyama_image_noise_before_tc_sampler, euler_maruyama_image_noise_sampler, euler_maruyama_sampler, euler_ode_sampler, ) def create_npz_from_sample_folder(sample_dir: str, num: int): """ 将 sample_dir 下 000000.png... 组装为单个 .npz（arr_0）。 """ samples = [] for i in tqdm(range(num), desc="Building .npz file from samples"): sample_pil = Image.open(os.path.join(sample_dir, f"{i:06d}.png")) sample_np = np.asarray(sample_pil).astype(np.uint8) samples.append(sample_np) samples = np.stack(samples) npz_path = f"{sample_dir}.npz" np.savez(npz_path, arr_0=samples) print(f"Saved .npz file to {npz_path} [shape={samples.shape}].") return npz_path def semantic_dim_from_enc_type(enc_type): if enc_type is None: return 768 s = str(enc_type).lower() if "vit-g" in s or "vitg" in s: return 1536 if "vit-l" in s or "vitl" in s: return 1024 if "vit-s" in s or "vits" in s: return 384 return 768 def load_train_args_from_ckpt(ckpt: dict) -> argparse.Namespace | None: a = ckpt.get("args") if a is None: return None if isinstance(a, argparse.Namespace): return a if isinstance(a, dict): return argparse.Namespace(**a) if isinstance(a, types.SimpleNamespace): return argparse.Namespace(**vars(a)) return None def load_vae(device: torch.device): try: from preprocessing import dnnlib cache_dir = dnnlib.make_cache_dir_path("diffusers") os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1") os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") os.environ["HF_HOME"] = cache_dir try: vae = AutoencoderKL.from_pretrained( "stabilityai/sd-vae-ft-mse", cache_dir=cache_dir, local_files_only=True, ).to(device) vae.eval() return vae except Exception: pass candidate_dir = None for root_dir in [ cache_dir, os.path.join(os.path.expanduser("~"), ".cache", "dnnlib", "diffusers"), os.path.join(os.path.expanduser("~"), ".cache", "diffusers"), os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub"), ]: if not os.path.isdir(root_dir): continue for root, _, files in os.walk(root_dir): if "config.json" in files and "sd-vae-ft-mse" in root.replace("\\", "/"): candidate_dir = root break if candidate_dir is not None: break if candidate_dir is not None: vae = AutoencoderKL.from_pretrained(candidate_dir, local_files_only=True).to(device) vae.eval() return vae except Exception: pass vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device) vae.eval() return vae def build_model_from_train_args(ta: argparse.Namespace, device: torch.device): res = int(getattr(ta, "resolution", 256)) latent_size = res // 8 enc_type = getattr(ta, "enc_type", "dinov2-vit-b") z_dims = [semantic_dim_from_enc_type(enc_type)] block_kwargs = { "fused_attn": getattr(ta, "fused_attn", True), "qk_norm": getattr(ta, "qk_norm", False), } cfg_prob = float(getattr(ta, "cfg_prob", 0.1)) if ta.model not in SiT_models: raise ValueError(f"未知 model={ta.model!r}，可选：{list(SiT_models.keys())}") model = SiT_models[ta.model]( input_size=latent_size, num_classes=int(getattr(ta, "num_classes", 1000)), use_cfg=(cfg_prob > 0), z_dims=z_dims, encoder_depth=int(getattr(ta, "encoder_depth", 8)), **block_kwargs, ).to(device) return model, z_dims[0] def resolve_tc_schedule(cli, ta): sb = cli.steps_before_tc sa = cli.steps_after_tc tc = cli.t_c if sb is None and sa is None: return None, None, None if sb is None or sa is None: print("使用分段步数时必须同时指定 --steps-before-tc 与 --steps-after-tc。", file=sys.stderr) sys.exit(1) if tc is None: tc = getattr(ta, "t_c", None) if ta is not None else None if tc is None: print("分段采样需要 --t-c，或检查点 args 中含 t_c。", file=sys.stderr) sys.exit(1) return float(tc), int(sb), int(sa) def parse_cli(): p = argparse.ArgumentParser(description="REG DDP 检查点采样（单路径，无 at_tc 图）") p.add_argument("--ckpt", type=str, required=True) p.add_argument("--out-dir", type=str, required=True) p.add_argument("--num-images", type=int, required=True) p.add_argument("--batch-size", type=int, default=16) p.add_argument("--seed", type=int, default=0) p.add_argument("--weights", type=str, choices=("ema", "model"), default="ema") p.add_argument("--device", type=str, default="cuda") p.add_argument("--num-steps", type=int, default=50) p.add_argument("--t-c", type=float, default=None) p.add_argument("--steps-before-tc", type=int, default=None) p.add_argument("--steps-after-tc", type=int, default=None) p.add_argument("--cfg-scale", type=float, default=1.0) p.add_argument("--cls-cfg-scale", type=float, default=0.0) p.add_argument("--guidance-low", type=float, default=0.0) p.add_argument("--guidance-high", type=float, default=1.0) p.add_argument("--path-type", type=str, default=None, choices=["linear", "cosine"]) p.add_argument("--legacy", action=argparse.BooleanOptionalAction, default=False) p.add_argument("--model", type=str, default=None) p.add_argument("--resolution", type=int, default=None, choices=[256, 512]) p.add_argument("--num-classes", type=int, default=1000) p.add_argument("--encoder-depth", type=int, default=None) p.add_argument("--enc-type", type=str, default=None) p.add_argument("--fused-attn", action=argparse.BooleanOptionalAction, default=None) p.add_argument("--qk-norm", action=argparse.BooleanOptionalAction, default=None) p.add_argument("--cfg-prob", type=float, default=None) p.add_argument( "--sampler", type=str, default="em_image_noise_before_tc", choices=["ode", "em", "em_image_noise", "em_image_noise_before_tc"], ) p.add_argument( "--save-fixed-trajectory", action="store_true", help="保存本 rank 轨迹（npy）到 out-dir/trajectory_rank{rank}", ) p.add_argument( "--save-npz", action=argparse.BooleanOptionalAction, default=True, help="采样结束后由 rank0 汇总 PNG 并保存 out-dir.npz", ) return p.parse_args() def _decode_to_uint8_hwc(latents, latents_bias, latents_scale, vae): imgs = vae.decode((latents - latents_bias) / latents_scale).sample imgs = (imgs + 1) / 2.0 imgs = torch.clamp(imgs, 0, 1) return ( (imgs * 255.0) .round() .to(torch.uint8) .permute(0, 2, 3, 1) .cpu() .numpy() ) def init_ddp(): if "RANK" in os.environ and "WORLD_SIZE" in os.environ: rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) local_rank = int(os.environ.get("LOCAL_RANK", 0)) dist.init_process_group(backend="nccl", init_method="env://") torch.cuda.set_device(local_rank) return True, rank, world_size, local_rank return False, 0, 1, 0 def main(): cli = parse_cli() use_ddp, rank, world_size, local_rank = init_ddp() if torch.cuda.is_available(): device = torch.device(f"cuda:{local_rank}" if use_ddp else cli.device) torch.backends.cuda.matmul.allow_tf32 = True else: device = torch.device("cpu") try: ckpt = torch.load(cli.ckpt, map_location="cpu", weights_only=False) except TypeError: ckpt = torch.load(cli.ckpt, map_location="cpu") ta = load_train_args_from_ckpt(ckpt) if ta is None: if cli.model is None or cli.resolution is None or cli.enc_type is None: print("检查点中无 args，请至少指定：--model --resolution --enc-type", file=sys.stderr) sys.exit(1) ta = argparse.Namespace( model=cli.model, resolution=cli.resolution, num_classes=cli.num_classes if cli.num_classes is not None else 1000, encoder_depth=cli.encoder_depth if cli.encoder_depth is not None else 8, enc_type=cli.enc_type, fused_attn=cli.fused_attn if cli.fused_attn is not None else True, qk_norm=cli.qk_norm if cli.qk_norm is not None else False, cfg_prob=cli.cfg_prob if cli.cfg_prob is not None else 0.1, ) else: if cli.model is not None: ta.model = cli.model if cli.resolution is not None: ta.resolution = cli.resolution if cli.num_classes is not None: ta.num_classes = cli.num_classes if cli.encoder_depth is not None: ta.encoder_depth = cli.encoder_depth if cli.enc_type is not None: ta.enc_type = cli.enc_type if cli.fused_attn is not None: ta.fused_attn = cli.fused_attn if cli.qk_norm is not None: ta.qk_norm = cli.qk_norm if cli.cfg_prob is not None: ta.cfg_prob = cli.cfg_prob path_type = cli.path_type if cli.path_type is not None else getattr(ta, "path_type", "linear") tc_split = resolve_tc_schedule(cli, ta) if rank == 0: if tc_split[0] is not None: print( f"时间网格：t_c={tc_split[0]}, 步数 (1→t_c)={tc_split[1]}, (t_c→0)={tc_split[2]}" ) else: print(f"时间网格：均匀 num_steps={cli.num_steps}") if cli.sampler == "ode": sampler_fn = euler_ode_sampler elif cli.sampler == "em": sampler_fn = euler_maruyama_sampler elif cli.sampler == "em_image_noise_before_tc": sampler_fn = euler_maruyama_image_noise_before_tc_sampler else: sampler_fn = euler_maruyama_image_noise_sampler model, cls_dim = build_model_from_train_args(ta, device) wkey = cli.weights if wkey not in ckpt: raise KeyError(f"检查点中无 '{wkey}' 键，现有键：{list(ckpt.keys())}") state = ckpt[wkey] if cli.legacy: from utils import load_legacy_checkpoints state = load_legacy_checkpoints( state_dict=state, encoder_depth=int(getattr(ta, "encoder_depth", 8)) ) model.load_state_dict(state, strict=True) model.eval() vae = load_vae(device) latents_scale = torch.tensor([0.18215] * 4, device=device).view(1, 4, 1, 1) latents_bias = torch.tensor([0.0] * 4, device=device).view(1, 4, 1, 1) sampler_args = argparse.Namespace(cls_cfg_scale=float(cli.cls_cfg_scale)) os.makedirs(cli.out_dir, exist_ok=True) traj_dir = None if cli.save_fixed_trajectory and cli.sampler != "em": traj_dir = os.path.join(cli.out_dir, f"trajectory_rank{rank}") os.makedirs(traj_dir, exist_ok=True) latent_size = int(getattr(ta, "resolution", 256)) // 8 n_total = int(cli.num_images) b = max(1, int(cli.batch_size)) global_batch_size = b * world_size total_samples = int(math.ceil(n_total / global_batch_size) * global_batch_size) samples_needed_this_gpu = int(total_samples // world_size) if samples_needed_this_gpu % b != 0: raise ValueError("samples_needed_this_gpu must be divisible by per-rank batch size") iterations = int(samples_needed_this_gpu // b) seed_rank = int(cli.seed) + int(rank) torch.manual_seed(seed_rank) if device.type == "cuda": torch.cuda.manual_seed_all(seed_rank) if rank == 0: print(f"Total number of images that will be sampled: {total_samples}") pbar = range(iterations) pbar = tqdm(pbar, desc="sampling") if rank == 0 else pbar total = 0 written_local = 0 for _ in pbar: cur = b z = torch.randn(cur, model.in_channels, latent_size, latent_size, device=device) y = torch.randint(0, int(ta.num_classes), (cur,), device=device) cls_z = torch.randn(cur, cls_dim, device=device) with torch.no_grad(): em_kw = dict( num_steps=cli.num_steps, cfg_scale=cli.cfg_scale, guidance_low=cli.guidance_low, guidance_high=cli.guidance_high, path_type=path_type, cls_latents=cls_z, args=sampler_args, ) if tc_split[0] is not None: em_kw["t_c"] = tc_split[0] em_kw["num_steps_before_tc"] = tc_split[1] em_kw["num_steps_after_tc"] = tc_split[2] if cli.save_fixed_trajectory and cli.sampler != "em": if cli.sampler == "em_image_noise_before_tc": latents, traj = sampler_fn( model, z, y, **em_kw, return_trajectory=True ) else: latents, traj = sampler_fn( model, z, y, **em_kw, return_trajectory=True ) else: latents = sampler_fn(model, z, y, **em_kw) traj = None latents = latents.to(torch.float32) imgs = _decode_to_uint8_hwc(latents, latents_bias, latents_scale, vae) for i, img in enumerate(imgs): gidx = i * world_size + rank + total if gidx < n_total: Image.fromarray(img).save(os.path.join(cli.out_dir, f"{gidx:06d}.png")) written_local += 1 if traj is not None and traj_dir is not None: traj_np = torch.stack(traj, dim=0).to(torch.float32).cpu().numpy() first_idx = rank + total if first_idx < n_total: np.save(os.path.join(traj_dir, f"{first_idx:06d}_traj.npy"), traj_np) total += global_batch_size if use_ddp: dist.barrier() if rank == 0 and hasattr(pbar, "close"): pbar.close() if use_ddp: dist.barrier() if rank == 0: if cli.save_npz: create_npz_from_sample_folder(cli.out_dir, n_total) print(f"Done. Saved {n_total} images under {cli.out_dir} (world_size={world_size}).") if use_ddp: dist.destroy_process_group() if __name__ == "__main__": main()