import argparse import copy from copy import deepcopy import logging import os from pathlib import Path from collections import OrderedDict import json import numpy as np import torch import torch.nn.functional as F import torch.utils.checkpoint from tqdm.auto import tqdm from torch.utils.data import DataLoader from accelerate import Accelerator, DistributedDataParallelKwargs from accelerate.logging import get_logger from accelerate.utils import ProjectConfiguration, set_seed from models.sit import SiT_models from loss import SILoss from utils import load_encoders from dataset import CustomDataset from diffusers.models import AutoencoderKL # import wandb_utils import wandb import math from torchvision.utils import make_grid from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from torchvision.transforms import Normalize from PIL import Image logger = get_logger(__name__) def semantic_dim_from_enc_type(enc_type): """DINOv2 等 enc_type 字符串推断 class token 维度（与预处理特征一致）。""" 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 CLIP_DEFAULT_MEAN = (0.48145466, 0.4578275, 0.40821073) CLIP_DEFAULT_STD = (0.26862954, 0.26130258, 0.27577711) def preprocess_raw_image(x, enc_type): resolution = x.shape[-1] if 'clip' in enc_type: x = x / 255. x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic') x = Normalize(CLIP_DEFAULT_MEAN, CLIP_DEFAULT_STD)(x) elif 'mocov3' in enc_type or 'mae' in enc_type: x = x / 255. x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x) elif 'dinov2' in enc_type: x = x / 255. x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x) x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic') elif 'dinov1' in enc_type: x = x / 255. x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x) elif 'jepa' in enc_type: x = x / 255. x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x) x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic') return x def array2grid(x): nrow = round(math.sqrt(x.size(0))) x = make_grid(x.clamp(0, 1), nrow=nrow, value_range=(0, 1)) x = x.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy() return x @torch.no_grad() def sample_posterior(moments, latents_scale=1., latents_bias=0.): device = moments.device mean, std = torch.chunk(moments, 2, dim=1) z = mean + std * torch.randn_like(mean) z = (z * latents_scale + latents_bias) return z @torch.no_grad() def update_ema(ema_model, model, decay=0.9999): """ Step the EMA model towards the current model. """ ema_params = OrderedDict(ema_model.named_parameters()) model_params = OrderedDict(model.named_parameters()) for name, param in model_params.items(): name = name.replace("module.", "") # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) def create_logger(logging_dir): """ Create a logger that writes to a log file and stdout. """ logging.basicConfig( level=logging.INFO, format='[\033[34m%(asctime)s\033[0m] %(message)s', datefmt='%Y-%m-%d %H:%M:%S', handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")] ) logger = logging.getLogger(__name__) return logger def requires_grad(model, flag=True): """ Set requires_grad flag for all parameters in a model. """ for p in model.parameters(): p.requires_grad = flag ################################################################################# # Training Loop # ################################################################################# def main(args): # set accelerator logging_dir = Path(args.output_dir, args.logging_dir) accelerator_project_config = ProjectConfiguration( project_dir=args.output_dir, logging_dir=logging_dir ) accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with=args.report_to, project_config=accelerator_project_config, kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)] ) if accelerator.is_main_process: os.makedirs(args.output_dir, exist_ok=True) # Make results folder (holds all experiment subfolders) save_dir = os.path.join(args.output_dir, args.exp_name) os.makedirs(save_dir, exist_ok=True) args_dict = vars(args) # Save to a JSON file json_dir = os.path.join(save_dir, "args.json") with open(json_dir, 'w') as f: json.dump(args_dict, f, indent=4) checkpoint_dir = f"{save_dir}/checkpoints" # Stores saved model checkpoints os.makedirs(checkpoint_dir, exist_ok=True) logger = create_logger(save_dir) logger.info(f"Experiment directory created at {save_dir}") device = accelerator.device if torch.backends.mps.is_available(): accelerator.native_amp = False if args.seed is not None: set_seed(args.seed + accelerator.process_index) # Create model: assert args.resolution % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)." latent_size = args.resolution // 8 train_dataset = CustomDataset( args.data_dir, semantic_features_dir=args.semantic_features_dir ) use_preprocessed_semantic = train_dataset.use_preprocessed_semantic if use_preprocessed_semantic: encoders, encoder_types, architectures = [], [], [] z_dims = [semantic_dim_from_enc_type(args.enc_type)] if accelerator.is_main_process: logger.info( f"Preprocessed semantic features: skip loading online encoder, z_dims={z_dims}" ) elif args.enc_type is not None: encoders, encoder_types, architectures = load_encoders( args.enc_type, device, args.resolution ) z_dims = [encoder.embed_dim for encoder in encoders] else: raise NotImplementedError() block_kwargs = {"fused_attn": args.fused_attn, "qk_norm": args.qk_norm} model = SiT_models[args.model]( input_size=latent_size, num_classes=args.num_classes, use_cfg = (args.cfg_prob > 0), z_dims = z_dims, encoder_depth=args.encoder_depth, **block_kwargs ) model = model.to(device) ema = deepcopy(model).to(device) # Create an EMA of the model for use after training requires_grad(ema, False) latents_scale = torch.tensor( [0.18215, 0.18215, 0.18215, 0.18215] ).view(1, 4, 1, 1).to(device) latents_bias = torch.tensor( [0., 0., 0., 0.] ).view(1, 4, 1, 1).to(device) # VAE decoder：采样阶段将 latent 解码为图像（与根目录 train.py / 预处理一致：sd-vae-ft-mse） try: from preprocessing import dnnlib cache_dir = dnnlib.make_cache_dir_path("diffusers") os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1" os.environ["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() if accelerator.is_main_process: logger.info( "Loaded VAE 'stabilityai/sd-vae-ft-mse' from local diffusers cache " f"at '{cache_dir}' for intermediate sampling." ) except Exception as e_main: vae = None candidate_dir = None possible_roots = [ 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"), ] checked_roots = [] for root_dir in possible_roots: if not os.path.isdir(root_dir): continue checked_roots.append(root_dir) for root, dirs, 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: try: vae = AutoencoderKL.from_pretrained( candidate_dir, local_files_only=True, ).to(device) vae.eval() if accelerator.is_main_process: logger.info( "Loaded VAE 'stabilityai/sd-vae-ft-mse' from discovered local path " f"'{candidate_dir}'. Searched roots: {checked_roots}" ) except Exception as e_fallback: if accelerator.is_main_process: logger.warning( "Tried to load VAE from discovered local path " f"'{candidate_dir}' but failed: {e_fallback}" ) if vae is None and accelerator.is_main_process: logger.warning( "Could not load VAE 'stabilityai/sd-vae-ft-mse' via repo name or local search. " f"Last repo-level error: {e_main}" ) except Exception as e: vae = None if accelerator.is_main_process: logger.warning( f"Failed to initialize VAE loading logic (will skip image decoding): {e}" ) # create loss function loss_fn = SILoss( prediction=args.prediction, path_type=args.path_type, encoders=encoders, accelerator=accelerator, latents_scale=latents_scale, latents_bias=latents_bias, weighting=args.weighting, t_c=args.t_c, ot_cls=args.ot_cls, ) if accelerator.is_main_process: logger.info(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}") # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper): if args.allow_tf32: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True optimizer = torch.optim.AdamW( model.parameters(), lr=args.learning_rate, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, eps=args.adam_epsilon, ) # Setup data（train_dataset 已在上方创建） local_batch_size = int(args.batch_size // accelerator.num_processes) train_dataloader = DataLoader( train_dataset, batch_size=local_batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, drop_last=True ) if accelerator.is_main_process: logger.info(f"Dataset contains {len(train_dataset):,} images ({args.data_dir})") # Prepare models for training: update_ema(ema, model, decay=0) # Ensure EMA is initialized with synced weights model.train() # important! This enables embedding dropout for classifier-free guidance ema.eval() # EMA model should always be in eval mode # resume: global_step = 0 if args.resume_from_ckpt is not None: ckpt = torch.load(args.resume_from_ckpt, map_location="cpu") model.load_state_dict(ckpt["model"]) ema.load_state_dict(ckpt["ema"]) if "opt" in ckpt: optimizer.load_state_dict(ckpt["opt"]) global_step = int(ckpt.get("steps", 0)) if accelerator.is_main_process: logger.info( f"Resumed from ckpt: {args.resume_from_ckpt} (global_step={global_step})" ) elif args.resume_step > 0: ckpt_name = str(args.resume_step).zfill(7) +'.pt' ckpt = torch.load( f'{os.path.join(args.output_dir, args.exp_name)}/checkpoints/{ckpt_name}', map_location='cpu', ) model.load_state_dict(ckpt['model']) ema.load_state_dict(ckpt['ema']) optimizer.load_state_dict(ckpt['opt']) global_step = ckpt['steps'] model, optimizer, train_dataloader = accelerator.prepare( model, optimizer, train_dataloader ) if accelerator.is_main_process: tracker_config = vars(copy.deepcopy(args)) accelerator.init_trackers( project_name="REG", config=tracker_config, init_kwargs={ "wandb": {"name": f"{args.exp_name}"} }, ) progress_bar = tqdm( range(0, args.max_train_steps), initial=global_step, desc="Steps", # Only show the progress bar once on each machine. disable=not accelerator.is_local_main_process, ) # Labels to condition the model with (feel free to change): sample_batch_size = 64 // accelerator.num_processes first_batch = next(iter(train_dataloader)) if len(first_batch) == 4: gt_raw_images, gt_xs, _, _ = first_batch else: gt_raw_images, gt_xs, _ = first_batch assert gt_raw_images.shape[-1] == args.resolution gt_xs = gt_xs[:sample_batch_size] gt_xs = sample_posterior( gt_xs.to(device), latents_scale=latents_scale, latents_bias=latents_bias ) ys = torch.randint(1000, size=(sample_batch_size,), device=device) ys = ys.to(device) # Create sampling noise: n = ys.size(0) xT = torch.randn((n, 4, latent_size, latent_size), device=device) for epoch in range(args.epochs): model.train() for batch in train_dataloader: if len(batch) == 4: raw_image, x, r_preprocessed, y = batch use_sem_file = True else: raw_image, x, y = batch r_preprocessed = None use_sem_file = False raw_image = raw_image.to(device) x = x.squeeze(dim=1).to(device).float() y = y.to(device) if args.legacy: # In our early experiments, we accidentally apply label dropping twice: # once in train.py and once in sit.py. # We keep this option for exact reproducibility with previous runs. drop_ids = torch.rand(y.shape[0], device=y.device) < args.cfg_prob labels = torch.where(drop_ids, args.num_classes, y) else: labels = y with torch.no_grad(): x = sample_posterior(x, latents_scale=latents_scale, latents_bias=latents_bias) zs = [] if use_sem_file and r_preprocessed is not None: cls_token = r_preprocessed.to(device).float() if cls_token.dim() == 1: cls_token = cls_token.unsqueeze(0) while cls_token.dim() > 2: cls_token = cls_token.squeeze(1) base_m = model.module if hasattr(model, "module") else model n_pad = base_m.x_embedder.num_patches zs = [ torch.cat( [ cls_token.unsqueeze(1), cls_token.unsqueeze(1).expand(-1, n_pad, -1), ], dim=1, ) ] else: with accelerator.autocast(): for encoder, encoder_type, arch in zip( encoders, encoder_types, architectures ): raw_image_ = preprocess_raw_image(raw_image, encoder_type) z = encoder.forward_features(raw_image_) if 'dinov2' in encoder_type: dense_z = z['x_norm_patchtokens'] cls_token = z['x_norm_clstoken'] dense_z = torch.cat([cls_token.unsqueeze(1), dense_z], dim=1) else: exit() zs.append(dense_z) with accelerator.accumulate(model): model_kwargs = dict(y=labels) loss1, proj_loss1, time_input, noises, loss2 = loss_fn(model, x, model_kwargs, zs=zs, cls_token=cls_token, time_input=None, noises=None) loss_mean = loss1.mean() loss_mean_cls = loss2.mean() * args.cls proj_loss_mean = proj_loss1.mean() * args.proj_coeff tc_vel_loss = torch.tensor(0.0, device=device) if args.tc_velocity_loss_coeff > 0: tc_vel_loss = loss_fn.tc_velocity_loss( model, x, model_kwargs=model_kwargs, cls_token=cls_token, noises=noises, ).mean() loss = ( loss_mean + proj_loss_mean + loss_mean_cls + args.tc_velocity_loss_coeff * tc_vel_loss ) ## optimization accelerator.backward(loss) if accelerator.sync_gradients: params_to_clip = model.parameters() grad_norm = accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) optimizer.step() optimizer.zero_grad(set_to_none=True) if accelerator.sync_gradients: update_ema(ema, model) # change ema function ### enter if accelerator.sync_gradients: progress_bar.update(1) global_step += 1 if global_step % args.checkpointing_steps == 0 and global_step > 0: if accelerator.is_main_process: checkpoint = { "model": model.module.state_dict(), "ema": ema.state_dict(), "opt": optimizer.state_dict(), "args": args, "steps": global_step, } checkpoint_path = f"{checkpoint_dir}/{global_step:07d}.pt" torch.save(checkpoint, checkpoint_path) logger.info(f"Saved checkpoint to {checkpoint_path}") if (global_step == 1 or (global_step % args.sampling_steps == 0 and global_step > 0)): t_mid_vis = float(args.t_c) tc_tag = f"{t_mid_vis:.4f}".rstrip("0").rstrip(".").replace(".", "_") logging.info( f"Generating EMA samples (Euler-Maruyama; t≈{t_mid_vis:g} → t=0)..." ) ema.eval() with torch.no_grad(): latent_size = args.resolution // 8 n_samples = min(16, args.batch_size) base_model = model.module if hasattr(model, "module") else model cls_dim = base_model.z_dims[0] shared_seed = torch.randint(0, 2**32, (1,), device=device).item() torch.manual_seed(shared_seed) z_init = torch.randn(n_samples, base_model.in_channels, latent_size, latent_size, device=device) torch.manual_seed(shared_seed) cls_init = torch.randn(n_samples, cls_dim, device=device) y_samples = torch.randint(0, args.num_classes, (n_samples,), device=device) from samplers import euler_maruyama_sampler z_0, z_mid, _ = euler_maruyama_sampler( ema, z_init, y_samples, num_steps=50, cfg_scale=1.0, guidance_low=0.0, guidance_high=1.0, path_type=args.path_type, cls_latents=cls_init, args=args, return_mid_state=True, t_mid=t_mid_vis, ) samples_root = os.path.join(args.output_dir, args.exp_name, "samples") t0_dir = os.path.join(samples_root, "t0") t_mid_dir = os.path.join(samples_root, f"t0_{tc_tag}") os.makedirs(t0_dir, exist_ok=True) os.makedirs(t_mid_dir, exist_ok=True) if vae is not None: z_f = z_0.to(dtype=torch.float32) samples_final = vae.decode((z_f - latents_bias) / latents_scale).sample samples_final = (samples_final + 1) / 2.0 samples_final = samples_final.clamp(0, 1) grid_final = array2grid(samples_final) Image.fromarray(grid_final).save( os.path.join(t0_dir, f"step_{global_step:07d}_t0.png") ) if z_mid is not None: z_m = z_mid.to(dtype=torch.float32) samples_mid = vae.decode((z_m - latents_bias) / latents_scale).sample samples_mid = (samples_mid + 1) / 2.0 samples_mid = samples_mid.clamp(0, 1) grid_mid = array2grid(samples_mid) Image.fromarray(grid_mid).save( os.path.join(t_mid_dir, f"step_{global_step:07d}_t0_{tc_tag}.png") ) else: logging.warning( f"Sampling time grid did not bracket t_mid={t_mid_vis:g}; " f"skip t0_{tc_tag} image this step." ) del z_init, cls_init, y_samples, z_0 if z_mid is not None: del z_mid if vae is not None: del samples_final, grid_final if "samples_mid" in locals(): del samples_mid, grid_mid torch.cuda.empty_cache() logs = { "loss_final": accelerator.gather(loss).mean().detach().item(), "loss_mean": accelerator.gather(loss_mean).mean().detach().item(), "proj_loss": accelerator.gather(proj_loss_mean).mean().detach().item(), "loss_mean_cls": accelerator.gather(loss_mean_cls).mean().detach().item(), "loss_tc_vel": accelerator.gather(tc_vel_loss).mean().detach().item(), "grad_norm": accelerator.gather(grad_norm).mean().detach().item() } log_message = ", ".join(f"{key}: {value:.6f}" for key, value in logs.items()) logging.info(f"Step: {global_step}, Training Logs: {log_message}") progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step) if global_step >= args.max_train_steps: break if global_step >= args.max_train_steps: break model.eval() # important! This disables randomized embedding dropout # do any sampling/FID calculation/etc. with ema (or model) in eval mode ... accelerator.wait_for_everyone() if accelerator.is_main_process: logger.info("Done!") accelerator.end_training() def parse_args(input_args=None): parser = argparse.ArgumentParser(description="Training") # logging: parser.add_argument("--output-dir", type=str, default="exps") parser.add_argument("--exp-name", type=str, required=True) parser.add_argument("--logging-dir", type=str, default="logs") parser.add_argument("--report-to", type=str, default="wandb") parser.add_argument("--sampling-steps", type=int, default=2000) parser.add_argument("--resume-step", type=int, default=0) parser.add_argument( "--resume-from-ckpt", type=str, default=None, help="直接从指定 checkpoint 路径续训（优先于 --resume-step）。", ) # model parser.add_argument("--model", type=str) parser.add_argument("--num-classes", type=int, default=1000) parser.add_argument("--encoder-depth", type=int, default=8) parser.add_argument("--fused-attn", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--qk-norm", action=argparse.BooleanOptionalAction, default=False) parser.add_argument("--ops-head", type=int, default=16) # dataset parser.add_argument("--data-dir", type=str, default="../data/imagenet256") parser.add_argument( "--semantic-features-dir", type=str, default=None, help="预处理 DINOv2 class token 等特征目录（含 dataset.json）。" "默认 None 时若存在 data-dir/imagenet_256_features/dinov2-vit-b_tmp/gpu0 则自动使用。", ) parser.add_argument("--resolution", type=int, choices=[256, 512], default=256) parser.add_argument("--batch-size", type=int, default=256)#256 # precision parser.add_argument("--allow-tf32", action="store_true") parser.add_argument("--mixed-precision", type=str, default="fp16", choices=["no", "fp16", "bf16"]) # optimization parser.add_argument("--epochs", type=int, default=14000) parser.add_argument("--max-train-steps", type=int, default=10000000) parser.add_argument("--checkpointing-steps", type=int, default=10000) parser.add_argument("--gradient-accumulation-steps", type=int, default=1) parser.add_argument("--learning-rate", type=float, default=1e-4) parser.add_argument("--adam-beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") parser.add_argument("--adam-beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") parser.add_argument("--adam-weight-decay", type=float, default=0., help="Weight decay to use.") parser.add_argument("--adam-epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") parser.add_argument("--max-grad-norm", default=1.0, type=float, help="Max gradient norm.") # seed parser.add_argument("--seed", type=int, default=0) # cpu parser.add_argument("--num-workers", type=int, default=4) # loss parser.add_argument("--path-type", type=str, default="linear", choices=["linear", "cosine"]) parser.add_argument("--prediction", type=str, default="v", choices=["v"]) # currently we only support v-prediction parser.add_argument("--cfg-prob", type=float, default=0.1) parser.add_argument("--enc-type", type=str, default='dinov2-vit-b') parser.add_argument("--proj-coeff", type=float, default=0.5) parser.add_argument("--weighting", default="uniform", type=str, help="Max gradient norm.") parser.add_argument("--legacy", action=argparse.BooleanOptionalAction, default=False) parser.add_argument("--cls", type=float, default=0.03) parser.add_argument( "--t-c", type=float, default=0.5, help="语义分界时刻（与脚本内 t 约定一致：t=1 噪声→t=0 数据）。" "t∈(t_c,1]：cls 沿 OT 配对后的路径插值（CFM/OT-CFM 式 minibatch OT）；" "t∈[0,t_c]：cls 固定为真实 encoder cls，目标 cls 速度为 0。", ) parser.add_argument( "--ot-cls", action=argparse.BooleanOptionalAction, default=True, help="在 t>t_c 段对 cls 噪声与 batch 内 cls_gt 做 minibatch 最优传输配对（需 scipy）；关闭则退化为独立高斯噪声配对。", ) parser.add_argument( "--tc-velocity-loss-coeff", type=float, default=0.0, help="额外 t=t_c 图像速度场监督项权重（>0 启用，用于增强单步性）。", ) if input_args is not None: args = parser.parse_args(input_args) else: args = parser.parse_args() return args if __name__ == "__main__": args = parse_args() main(args)