Add files using upload-large-folder tool

REG copy/train.sh

@@ -0,0 +1,28 @@
+NUM_GPUS=8
+random_number=$((RANDOM % 100 + 1200))
+
+
+accelerate launch --multi_gpu --num_processes $NUM_GPUS train.py \
+    --report-to="wandb" \
+    --allow-tf32 \
+    --mixed-precision="fp16" \
+    --seed=0 \
+    --path-type="linear" \
+    --prediction="v" \
+    --weighting="uniform" \
+    --model="SiT-XL/2" \
+    --enc-type="dinov2-vit-b" \
+    --proj-coeff=0.5 \
+    --encoder-depth=8 \     #SiT-L/XL use 8, SiT-B use 4
+    --output-dir="your_path/reg_xlarge_dinov2_base_align_8_cls" \
+    --exp-name="linear-dinov2-b-enc8" \
+    --batch-size=256 \
+    --data-dir="data_path/imagenet_vae" \
+    --cls=0.03
+
+
+    #Dataset Path
+    #For example: your_path/imagenet-vae
+    #This folder contains two folders
+    #(1) The imagenet's RGB image: your_path/imagenet-vae/imagenet_256-vae/
+    #(2) The imagenet's VAE latent: your_path/imagenet-vae/vae-sd/

VIRTUAL_imagenet256_labeled/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

VIRTUAL_imagenet256_labeled/README.md

@@ -0,0 +1,11 @@
+---
+license: Apache License 2.0
+---
+数据集文件元信息以及数据文件，请浏览“数据集文件”页面获取。
+
+当前数据集卡片使用的是默认模版，数据集的贡献者未提供更加详细的数据集介绍，但是您可以通过如下GIT Clone命令，或者ModelScope SDK来下载数据集
+
+#### 下载方法 
+:modelscope-code[]{type="sdk"}
+:modelscope-code[]{type="git"}
+

back/dataset.py

@@ -0,0 +1,149 @@
+import os
+import json
+
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+
+from PIL import Image
+import PIL.Image
+try:
+    import pyspng
+except ImportError:
+    pyspng = None
+
+
+class CustomDataset(Dataset):
+    """
+    data_dir 下 VAE latent：imagenet_256_vae/
+    无预处理语义时：VAE 统计量/配对文件在 vae-sd/（与原 REG 一致）。
+    有 semantic_features_dir 时：与主仓库 dataset 一致，从该目录 dataset.json 索引，
+    按特征文件名推断 imagenet_256_vae 中对应 npy。
+    """
+
+    def __init__(self, data_dir, semantic_features_dir=None):
+        PIL.Image.init()
+        supported_ext = PIL.Image.EXTENSION.keys() | {'.npy'}
+
+        self.images_dir = os.path.join(data_dir, 'imagenet_256_vae')
+
+        if semantic_features_dir is None:
+            potential_semantic_dir = os.path.join(
+                data_dir, 'imagenet_256_features', 'dinov2-vit-b_tmp', 'gpu0'
+            )
+            if os.path.exists(potential_semantic_dir):
+                self.semantic_features_dir = potential_semantic_dir
+                self.use_preprocessed_semantic = True
+                print(f"Found preprocessed semantic features at: {self.semantic_features_dir}")
+            else:
+                self.semantic_features_dir = None
+                self.use_preprocessed_semantic = False
+        else:
+            self.semantic_features_dir = semantic_features_dir
+            self.use_preprocessed_semantic = True
+            print(f"Using preprocessed semantic features from: {self.semantic_features_dir}")
+
+        if self.use_preprocessed_semantic:
+            label_fname = os.path.join(self.semantic_features_dir, 'dataset.json')
+            if not os.path.exists(label_fname):
+                raise FileNotFoundError(f"Label file not found: {label_fname}")
+
+            print(f"Using {label_fname}.")
+            with open(label_fname, 'rb') as f:
+                data = json.load(f)
+            labels_list = data.get('labels', None)
+            if labels_list is None:
+                raise ValueError(f"'labels' field is missing in {label_fname}")
+
+            semantic_fnames = []
+            labels = []
+            for entry in labels_list:
+                if entry is None:
+                    continue
+                fname, lab = entry
+                semantic_fnames.append(fname)
+                labels.append(0 if lab is None else lab)
+
+            self.semantic_fnames = semantic_fnames
+            self.labels = np.array(labels, dtype=np.int64)
+            self.num_samples = len(self.semantic_fnames)
+            print(f"Loaded {self.num_samples} semantic entries from dataset.json")
+        else:
+            self.features_dir = os.path.join(data_dir, 'vae-sd')
+
+            self._image_fnames = {
+                os.path.relpath(os.path.join(root, fname), start=self.images_dir)
+                for root, _dirs, files in os.walk(self.images_dir) for fname in files
+            }
+            self.image_fnames = sorted(
+                fname for fname in self._image_fnames if self._file_ext(fname) in supported_ext
+            )
+            self._feature_fnames = {
+                os.path.relpath(os.path.join(root, fname), start=self.features_dir)
+                for root, _dirs, files in os.walk(self.features_dir) for fname in files
+            }
+            self.feature_fnames = sorted(
+                fname for fname in self._feature_fnames if self._file_ext(fname) in supported_ext
+            )
+
+            fname = os.path.join(self.features_dir, 'dataset.json')
+            if os.path.exists(fname):
+                print(f"Using {fname}.")
+            else:
+                raise FileNotFoundError("Neither of the specified files exists.")
+
+            with open(fname, 'rb') as f:
+                labels = json.load(f)['labels']
+            labels = dict(labels)
+            labels = [labels[fname.replace('\\', '/')] for fname in self.feature_fnames]
+            labels = np.array(labels)
+            self.labels = labels.astype({1: np.int64, 2: np.float32}[labels.ndim])
+
+    def _file_ext(self, fname):
+        return os.path.splitext(fname)[1].lower()
+
+    def __len__(self):
+        if self.use_preprocessed_semantic:
+            return self.num_samples
+        assert len(self.image_fnames) == len(self.feature_fnames), \
+            "Number of feature files and label files should be same"
+        return len(self.feature_fnames)
+
+    def __getitem__(self, idx):
+        if self.use_preprocessed_semantic:
+            semantic_fname = self.semantic_fnames[idx]
+            basename = os.path.basename(semantic_fname)
+            idx_str = basename.split('-')[-1].split('.')[0]
+            subdir = idx_str[:5]
+            vae_relpath = os.path.join(subdir, f"img-mean-std-{idx_str}.npy")
+            vae_path = os.path.join(self.images_dir, vae_relpath)
+
+            with open(vae_path, 'rb') as f:
+                image = np.load(f)
+
+            semantic_path = os.path.join(self.semantic_features_dir, semantic_fname)
+            semantic_features = np.load(semantic_path)
+
+            return (
+                torch.from_numpy(image).float(),
+                torch.from_numpy(image).float(),
+                torch.from_numpy(semantic_features).float(),
+                torch.tensor(self.labels[idx]),
+            )
+
+        image_fname = self.image_fnames[idx]
+        feature_fname = self.feature_fnames[idx]
+        image_ext = self._file_ext(image_fname)
+        with open(os.path.join(self.images_dir, image_fname), 'rb') as f:
+            if image_ext == '.npy':
+                image = np.load(f)
+                image = image.reshape(-1, *image.shape[-2:])
+            elif image_ext == '.png' and pyspng is not None:
+                image = pyspng.load(f.read())
+                image = image.reshape(*image.shape[:2], -1).transpose(2, 0, 1)
+            else:
+                image = np.array(PIL.Image.open(f))
+                image = image.reshape(*image.shape[:2], -1).transpose(2, 0, 1)
+
+        features = np.load(os.path.join(self.features_dir, feature_fname))
+        return torch.from_numpy(image), torch.from_numpy(features), torch.tensor(self.labels[idx])

back/generate.py

@@ -0,0 +1,227 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# 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 SiT model using DDP.
+Subsequently saves a .npz file that can be used to compute FID and other
+evaluation metrics via the ADM repo: https://github.com/openai/guided-diffusion/tree/main/evaluations
+
+For a simple single-GPU/CPU sampling script, see sample.py.
+"""
+import torch
+import torch.distributed as dist
+from models.sit import SiT_models
+from diffusers.models import AutoencoderKL
+from tqdm import tqdm
+import os
+from PIL import Image
+import numpy as np
+import math
+import argparse
+from samplers import euler_maruyama_sampler
+from utils import load_legacy_checkpoints, download_model
+
+
+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)
+    print(f"Saved .npz file to {npz_path} [shape={samples.shape}].")
+    return npz_path
+
+
+def main(args):
+    """
+    Run sampling.
+    """
+    torch.backends.cuda.matmul.allow_tf32 = args.tf32  # True: fast but may lead to some small numerical differences
+    assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage"
+    torch.set_grad_enabled(False)
+
+    # Setup DDP:cd
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    device = rank % torch.cuda.device_count()
+    seed = args.global_seed * dist.get_world_size() + rank
+    torch.manual_seed(seed)
+    torch.cuda.set_device(device)
+    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
+
+    # Load model:
+    block_kwargs = {"fused_attn": args.fused_attn, "qk_norm": args.qk_norm}
+    latent_size = args.resolution // 8
+    model = SiT_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes,
+        use_cfg = True,
+        z_dims = [int(z_dim) for z_dim in args.projector_embed_dims.split(',')],
+        encoder_depth=args.encoder_depth,
+        **block_kwargs,
+    ).to(device)
+    # Auto-download a pre-trained model or load a custom SiT checkpoint from train.py:
+    ckpt_path = args.ckpt
+
+
+    #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+    if ckpt_path is None:
+        args.ckpt = 'SiT-XL-2-256x256.pt'
+        assert args.model == 'SiT-XL/2'
+        assert len(args.projector_embed_dims.split(',')) == 1
+        assert int(args.projector_embed_dims.split(',')[0]) == 768
+        state_dict = download_model('last.pt')
+    else:
+        state_dict = torch.load(ckpt_path, map_location=f'cuda:{device}')['ema']
+
+    if args.legacy:
+        state_dict = load_legacy_checkpoints(
+            state_dict=state_dict, encoder_depth=args.encoder_depth
+            )
+    model.load_state_dict(state_dict)
+    #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+
+    model.eval()  # important!
+    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
+    #vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path="your_local_path/weight/").to(device)
+
+
+    # Create folder to save samples:
+    model_string_name = args.model.replace("/", "-")
+    ckpt_string_name = os.path.basename(args.ckpt).replace(".pt", "") if args.ckpt else "pretrained"
+    folder_name = f"{model_string_name}-{ckpt_string_name}-size-{args.resolution}-vae-{args.vae}-" \
+                  f"cfg-{args.cfg_scale}-seed-{args.global_seed}-{args.mode}-{args.guidance_high}-{args.cls_cfg_scale}"
+    sample_folder_dir = f"{args.sample_dir}/{folder_name}"
+    if rank == 0:
+        os.makedirs(sample_folder_dir, exist_ok=True)
+        print(f"Saving .png samples at {sample_folder_dir}")
+    dist.barrier()
+
+    # Figure out how many samples we need to generate on each GPU and how many iterations we need to run:
+    n = args.per_proc_batch_size
+    global_batch_size = n * dist.get_world_size()
+    # To make things evenly-divisible, we'll sample a bit more than we need and then discard the extra samples:
+    total_samples = int(math.ceil(args.num_fid_samples / global_batch_size) * global_batch_size)
+    if rank == 0:
+        print(f"Total number of images that will be sampled: {total_samples}")
+        print(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
+        print(f"projector Parameters: {sum(p.numel() for p in model.projectors.parameters()):,}")
+    assert total_samples % dist.get_world_size() == 0, "total_samples must be divisible by world_size"
+    samples_needed_this_gpu = int(total_samples // dist.get_world_size())
+    assert samples_needed_this_gpu % n == 0, "samples_needed_this_gpu must be divisible by the per-GPU batch size"
+    iterations = int(samples_needed_this_gpu // n)
+    pbar = range(iterations)
+    pbar = tqdm(pbar) if rank == 0 else pbar
+    total = 0
+    for _ in pbar:
+        # Sample inputs:
+        z = torch.randn(n, model.in_channels, latent_size, latent_size, device=device)
+        y = torch.randint(0, args.num_classes, (n,), device=device)
+        cls_z = torch.randn(n, args.cls, device=device)
+
+        # Sample images:
+        sampling_kwargs = dict(
+            model=model, 
+            latents=z,
+            y=y,
+            num_steps=args.num_steps, 
+            heun=args.heun,
+            cfg_scale=args.cfg_scale,
+            guidance_low=args.guidance_low,
+            guidance_high=args.guidance_high,
+            path_type=args.path_type,
+            cls_latents=cls_z,
+            args=args
+        )
+        with torch.no_grad():
+            if args.mode == "sde":
+                samples = euler_maruyama_sampler(**sampling_kwargs).to(torch.float32)
+            elif args.mode == "ode":# will support
+                exit()
+                #samples = euler_sampler(**sampling_kwargs).to(torch.float32)
+            else:
+                raise NotImplementedError()
+
+            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)
+            samples = vae.decode((samples -  latents_bias) / latents_scale).sample
+            samples = (samples + 1) / 2.
+            samples = torch.clamp(
+                255. * samples, 0, 255
+                ).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+
+            # Save samples to disk as individual .png files
+            for i, sample in enumerate(samples):
+                index = i * dist.get_world_size() + rank + total
+                Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png")
+        total += global_batch_size
+
+    # Make sure all processes have finished saving their samples before attempting to convert to .npz
+    dist.barrier()
+    if rank == 0:
+        create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)
+        print("Done.")
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # seed
+    parser.add_argument("--global-seed", type=int, default=0)
+
+    # precision
+    parser.add_argument("--tf32", action=argparse.BooleanOptionalAction, default=True,
+                        help="By default, use TF32 matmuls. This massively accelerates sampling on Ampere GPUs.")
+
+    # logging/saving:
+    parser.add_argument("--ckpt", type=str, default=None, help="Optional path to a SiT checkpoint.")
+    parser.add_argument("--sample-dir", type=str, default="samples")
+
+    # model
+    parser.add_argument("--model", type=str, choices=list(SiT_models.keys()), default="SiT-XL/2")
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--encoder-depth", type=int, default=8)
+    parser.add_argument("--resolution", type=int, choices=[256, 512], default=256)
+    parser.add_argument("--fused-attn", action=argparse.BooleanOptionalAction, default=False)
+    parser.add_argument("--qk-norm", action=argparse.BooleanOptionalAction, default=False)
+    # vae
+    parser.add_argument("--vae",  type=str, choices=["ema", "mse"], default="ema")
+
+    # number of samples
+    parser.add_argument("--per-proc-batch-size", type=int, default=32)
+    parser.add_argument("--num-fid-samples", type=int, default=50_000)
+
+    # sampling related hyperparameters
+    parser.add_argument("--mode", type=str, default="ode")
+    parser.add_argument("--cfg-scale",  type=float, default=1.5)
+    parser.add_argument("--cls-cfg-scale",  type=float, default=1.5)
+    parser.add_argument("--projector-embed-dims", type=str, default="768,1024")
+    parser.add_argument("--path-type", type=str, default="linear", choices=["linear", "cosine"])
+    parser.add_argument("--num-steps", type=int, default=50)
+    parser.add_argument("--heun", action=argparse.BooleanOptionalAction, default=False) # only for ode
+    parser.add_argument("--guidance-low", type=float, default=0.)
+    parser.add_argument("--guidance-high", type=float, default=1.)
+    parser.add_argument('--local-rank', default=-1, type=int)
+    parser.add_argument('--cls', default=768, type=int)
+    # will be deprecated
+    parser.add_argument("--legacy", action=argparse.BooleanOptionalAction, default=False) # only for ode
+
+
+    args = parser.parse_args()
+    main(args)

back/sample_from_checkpoint_ddp.py

@@ -0,0 +1,416 @@
+#!/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()
+

back/samplers.py

@@ -0,0 +1,840 @@
+import torch
+import numpy as np
+
+
+def expand_t_like_x(t, x_cur):
+    """Function to reshape time t to broadcastable dimension of x
+    Args:
+      t: [batch_dim,], time vector
+      x: [batch_dim,...], data point
+    """
+    dims = [1] * (len(x_cur.size()) - 1)
+    t = t.view(t.size(0), *dims)
+    return t
+
+def get_score_from_velocity(vt, xt, t, path_type="linear"):
+    """Wrapper function: transfrom velocity prediction model to score
+    Args:
+        velocity: [batch_dim, ...] shaped tensor; velocity model output
+        x: [batch_dim, ...] shaped tensor; x_t data point
+        t: [batch_dim,] time tensor
+    """
+    t = expand_t_like_x(t, xt)
+    if path_type == "linear":
+        alpha_t, d_alpha_t = 1 - t, torch.ones_like(xt, device=xt.device) * -1
+        sigma_t, d_sigma_t = t, torch.ones_like(xt, device=xt.device)
+    elif path_type == "cosine":
+        alpha_t = torch.cos(t * np.pi / 2)
+        sigma_t = torch.sin(t * np.pi / 2)
+        d_alpha_t = -np.pi / 2 * torch.sin(t * np.pi / 2)
+        d_sigma_t =  np.pi / 2 * torch.cos(t * np.pi / 2)
+    else:
+        raise NotImplementedError
+
+    mean = xt
+    reverse_alpha_ratio = alpha_t / d_alpha_t
+    var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+    score = (reverse_alpha_ratio * vt - mean) / var
+
+    return score
+
+
+def compute_diffusion(t_cur):
+    return 2 * t_cur
+
+
+def build_sampling_time_steps(
+    num_steps=50,
+    t_c=None,
+    num_steps_before_tc=None,
+    num_steps_after_tc=None,
+    t_floor=0.04,
+):
+    """
+    构造从 t=1 → t=0 的时间网格（与原先一致：最后一段到 0 前保留 t_floor，再接到 0）。
+
+    - 默认：均匀 linspace(1, t_floor, num_steps)，再 append 0。
+    - 分段：t∈(t_c,1] 用 num_steps_before_tc 步（从 1 线性到 t_c）；
+            t∈[0,t_c] 用 num_steps_after_tc 步（从 t_c 线性到 t_floor），再 append 0。
+    """
+    t_floor = float(t_floor)
+    if t_c is None or num_steps_before_tc is None or num_steps_after_tc is None:
+        ns = int(num_steps)
+        if ns < 1:
+            raise ValueError("num_steps must be >= 1")
+        t_steps = torch.linspace(1.0, t_floor, ns, dtype=torch.float64)
+        return torch.cat([t_steps, torch.tensor([0.0], dtype=torch.float64)])
+
+    t_c = float(t_c)
+    nb = int(num_steps_before_tc)
+    na = int(num_steps_after_tc)
+    if nb < 1 or na < 1:
+        raise ValueError("num_steps_before_tc and num_steps_after_tc must be >= 1 when using t_c")
+    if not (0.0 < t_c < 1.0):
+        raise ValueError("t_c must be in (0, 1)")
+    if t_c <= t_floor:
+        raise ValueError(f"t_c ({t_c}) must be > t_floor ({t_floor})")
+
+    p1 = torch.linspace(1.0, t_c, nb + 1, dtype=torch.float64)
+    p2 = torch.linspace(t_c, t_floor, na + 1, dtype=torch.float64)
+    t_steps = torch.cat([p1, p2[1:]])
+    return torch.cat([t_steps, torch.tensor([0.0], dtype=torch.float64)])
+
+
+def _tc_segmented_freeze_cls(t_c, num_steps_before_tc, num_steps_after_tc):
+    """仅在 1→t_c→0 分段网格下启用：t∈[0,t_c] 段固定使用到达 t_c 时的 cls。"""
+    return (
+        t_c is not None
+        and num_steps_before_tc is not None
+        and num_steps_after_tc is not None
+    )
+
+
+def _cls_effective_and_freeze(
+    cls_x_cur, cls_frozen, t_cur, t_c_v, freeze_after_tc
+):
+    """
+    时间从 1 减到 0：当 t_cur <= t_c 时冻结 cls（取首次进入该段时的 cls_x_cur）。
+    返回 (用于前向的 cls, 更新后的 cls_frozen)。
+    """
+    if not freeze_after_tc or t_c_v is None:
+        return cls_x_cur, cls_frozen
+    if float(t_cur) <= float(t_c_v) + 1e-9:
+        if cls_frozen is None:
+            cls_frozen = cls_x_cur.clone()
+        return cls_frozen, cls_frozen
+    return cls_x_cur, cls_frozen
+
+
+def _build_euler_sampler_time_steps(
+    num_steps, t_c, num_steps_before_tc, num_steps_after_tc, device
+):
+    """
+    euler_sampler / REG ODE 用时间网格：默认 linspace(1,0)；分段时为 1→t_c→0 直连，无 t_floor。
+    """
+    if t_c is None or num_steps_before_tc is None or num_steps_after_tc is None:
+        ns = int(num_steps)
+        if ns < 1:
+            raise ValueError("num_steps must be >= 1")
+        return torch.linspace(1.0, 0.0, ns + 1, dtype=torch.float64, device=device)
+    t_c = float(t_c)
+    nb = int(num_steps_before_tc)
+    na = int(num_steps_after_tc)
+    if nb < 1 or na < 1:
+        raise ValueError(
+            "num_steps_before_tc and num_steps_after_tc must be >= 1 when using t_c"
+        )
+    if not (0.0 < t_c < 1.0):
+        raise ValueError("t_c must be in (0, 1)")
+    p1 = torch.linspace(1.0, t_c, nb + 1, dtype=torch.float64, device=device)
+    p2 = torch.linspace(t_c, 0.0, na + 1, dtype=torch.float64, device=device)
+    return torch.cat([p1, p2[1:]])
+
+
+def euler_maruyama_sampler(
+        model,
+        latents,
+        y,
+        num_steps=20,
+        heun=False,  # not used, just for compatability
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear",
+        cls_latents=None,
+        args=None,
+        return_mid_state=False,
+        t_mid=0.5,
+        t_c=None,
+        num_steps_before_tc=None,
+        num_steps_after_tc=None,
+        deterministic=False,
+        return_trajectory=False,
+        ):
+    """
+    Euler–Maruyama：漂移项与 score/velocity 变换与 euler_ode_sampler（euler_sampler）一致；
+    deterministic=True 时关闭扩散噪声项。ODE 使用 euler_sampler 的 linspace(1→0) / t_c 分段网格（无 t_floor），
+    本函数仍用 build_sampling_time_steps（含 t_floor），与 EM/SDE 对齐。
+    """
+    # setup conditioning
+    if cfg_scale > 1.0:
+        y_null = torch.tensor([1000] * y.size(0), device=y.device)
+        #[1000, 1000]
+    _dtype = latents.dtype
+    cls_cfg = getattr(args, "cls_cfg_scale", 0) if args is not None else 0
+
+    t_steps = build_sampling_time_steps(
+        num_steps=num_steps,
+        t_c=t_c,
+        num_steps_before_tc=num_steps_before_tc,
+        num_steps_after_tc=num_steps_after_tc,
+    )
+    freeze_after_tc = _tc_segmented_freeze_cls(t_c, num_steps_before_tc, num_steps_after_tc)
+    t_c_v = float(t_c) if freeze_after_tc else None
+    x_next = latents.to(torch.float64)
+    cls_x_next = cls_latents.to(torch.float64)
+    device = x_next.device
+    z_mid = cls_mid = None
+    t_mid = float(t_mid)
+    cls_frozen = None
+    traj = [x_next.clone()] if return_trajectory else None
+
+    with torch.no_grad():
+        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-2], t_steps[1:-1])):
+            dt = t_next - t_cur
+            x_cur = x_next
+            cls_x_cur = cls_x_next
+
+            cls_model_input, cls_frozen = _cls_effective_and_freeze(
+                cls_x_cur, cls_frozen, t_cur, t_c_v, freeze_after_tc
+            )
+
+            tc, tn = float(t_cur), float(t_next)
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                if abs(tc - t_mid) < abs(tn - t_mid):
+                    z_mid = x_cur.clone()
+                    cls_mid = cls_model_input.clone()
+
+            if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                model_input = torch.cat([x_cur] * 2, dim=0)
+                cls_model_input = torch.cat([cls_model_input] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+            else:
+                model_input = x_cur
+                y_cur = y
+
+            kwargs = dict(y=y_cur)
+            time_input = torch.ones(model_input.size(0)).to(device=device, dtype=torch.float64) * t_cur
+            diffusion = compute_diffusion(t_cur)
+
+            if deterministic:
+                deps = torch.zeros_like(x_cur)
+                cls_deps = torch.zeros_like(cls_model_input[: x_cur.size(0)])
+            else:
+                eps_i = torch.randn_like(x_cur).to(device)
+                cls_eps_i = torch.randn_like(cls_model_input[: x_cur.size(0)]).to(device)
+                deps = eps_i * torch.sqrt(torch.abs(dt))
+                cls_deps = cls_eps_i * torch.sqrt(torch.abs(dt))
+
+            # compute drift
+            v_cur, _, cls_v_cur = model(
+                model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs, cls_token=cls_model_input.to(dtype=_dtype)
+                )
+            v_cur = v_cur.to(torch.float64)
+            cls_v_cur = cls_v_cur.to(torch.float64)
+
+            s_cur = get_score_from_velocity(v_cur, model_input, time_input, path_type=path_type)
+            d_cur = v_cur - 0.5 * diffusion * s_cur
+
+            cls_s_cur = get_score_from_velocity(cls_v_cur, cls_model_input, time_input, path_type=path_type)
+            cls_d_cur = cls_v_cur - 0.5 * diffusion * cls_s_cur
+
+            if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+                d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+                d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+                cls_d_cur_cond, cls_d_cur_uncond = cls_d_cur.chunk(2)
+                if cls_cfg > 0:
+                    cls_d_cur = cls_d_cur_uncond + cls_cfg * (cls_d_cur_cond - cls_d_cur_uncond)
+                else:
+                    cls_d_cur = cls_d_cur_cond
+            x_next = x_cur + d_cur * dt + torch.sqrt(diffusion) * deps
+            if freeze_after_tc and t_c_v is not None and float(t_cur) <= float(t_c_v) + 1e-9:
+                cls_x_next = cls_frozen
+            else:
+                cls_x_next = cls_x_cur + cls_d_cur * dt + torch.sqrt(diffusion) * cls_deps
+
+            if return_trajectory:
+                traj.append(x_next.clone())
+
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                z_mid = x_next.clone()
+                cls_mid = cls_x_next.clone()
+
+    # last step
+    t_cur, t_next = t_steps[-2], t_steps[-1]
+    dt = t_next - t_cur
+    x_cur = x_next
+    cls_x_cur = cls_x_next
+
+    cls_model_input, cls_frozen = _cls_effective_and_freeze(
+        cls_x_cur, cls_frozen, t_cur, t_c_v, freeze_after_tc
+    )
+
+    if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+        model_input = torch.cat([x_cur] * 2, dim=0)
+        cls_model_input = torch.cat([cls_model_input] * 2, dim=0)
+        y_cur = torch.cat([y, y_null], dim=0)
+    else:
+        model_input = x_cur
+        y_cur = y
+    kwargs = dict(y=y_cur)
+    time_input = torch.ones(model_input.size(0)).to(
+        device=device, dtype=torch.float64
+        ) * t_cur
+    
+    # compute drift
+    v_cur, _, cls_v_cur = model(
+        model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs, cls_token=cls_model_input.to(dtype=_dtype)
+        )
+    v_cur = v_cur.to(torch.float64)
+    cls_v_cur = cls_v_cur.to(torch.float64)
+
+
+    s_cur = get_score_from_velocity(v_cur, model_input, time_input, path_type=path_type)
+    cls_s_cur = get_score_from_velocity(cls_v_cur, cls_model_input, time_input, path_type=path_type)
+
+    diffusion = compute_diffusion(t_cur)
+    d_cur = v_cur - 0.5 * diffusion * s_cur
+    cls_d_cur = cls_v_cur - 0.5 * diffusion * cls_s_cur  # d_cur [b, 4, 32 ,32]
+
+    if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+        d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+        d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+        cls_d_cur_cond, cls_d_cur_uncond = cls_d_cur.chunk(2)
+        if cls_cfg > 0:
+            cls_d_cur = cls_d_cur_uncond + cls_cfg * (cls_d_cur_cond - cls_d_cur_uncond)
+        else:
+            cls_d_cur = cls_d_cur_cond
+
+    mean_x = x_cur + dt * d_cur
+    if freeze_after_tc and t_c_v is not None and float(t_cur) <= float(t_c_v) + 1e-9:
+        cls_mean_x = cls_frozen
+    else:
+        cls_mean_x = cls_x_cur + dt * cls_d_cur
+
+    if return_trajectory:
+        traj.append(mean_x.clone())
+
+    if return_trajectory and return_mid_state:
+        return mean_x, z_mid, cls_mid, traj
+    if return_trajectory:
+        return mean_x, traj
+    if return_mid_state:
+        return mean_x, z_mid, cls_mid
+    return mean_x
+
+
+def euler_maruyama_image_noise_sampler(
+        model,
+        latents,
+        y,
+        num_steps=20,
+        heun=False,  # not used, just for compatability
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear",
+        cls_latents=None,
+        args=None,
+        return_mid_state=False,
+        t_mid=0.5,
+        t_c=None,
+        num_steps_before_tc=None,
+        num_steps_after_tc=None,
+        return_trajectory=False,
+        ):
+    """
+    EM 采样变体：仅图像 latent 引入随机扩散噪声，cls/token 通道不引入随机项（deterministic）。
+    """
+    if cfg_scale > 1.0:
+        y_null = torch.tensor([1000] * y.size(0), device=y.device)
+    _dtype = latents.dtype
+    cls_cfg = getattr(args, "cls_cfg_scale", 0) if args is not None else 0
+
+    t_steps = build_sampling_time_steps(
+        num_steps=num_steps,
+        t_c=t_c,
+        num_steps_before_tc=num_steps_before_tc,
+        num_steps_after_tc=num_steps_after_tc,
+    )
+    freeze_after_tc = _tc_segmented_freeze_cls(t_c, num_steps_before_tc, num_steps_after_tc)
+    t_c_v = float(t_c) if freeze_after_tc else None
+    x_next = latents.to(torch.float64)
+    cls_x_next = cls_latents.to(torch.float64)
+    device = x_next.device
+    z_mid = cls_mid = None
+    t_mid = float(t_mid)
+    cls_frozen = None
+    traj = [x_next.clone()] if return_trajectory else None
+
+    with torch.no_grad():
+        for t_cur, t_next in zip(t_steps[:-2], t_steps[1:-1]):
+            dt = t_next - t_cur
+            x_cur = x_next
+            cls_x_cur = cls_x_next
+
+            cls_model_input, cls_frozen = _cls_effective_and_freeze(
+                cls_x_cur, cls_frozen, t_cur, t_c_v, freeze_after_tc
+            )
+
+            tc, tn = float(t_cur), float(t_next)
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                if abs(tc - t_mid) < abs(tn - t_mid):
+                    z_mid = x_cur.clone()
+                    cls_mid = cls_model_input.clone()
+
+            if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                model_input = torch.cat([x_cur] * 2, dim=0)
+                cls_model_input = torch.cat([cls_model_input] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+            else:
+                model_input = x_cur
+                y_cur = y
+
+            kwargs = dict(y=y_cur)
+            time_input = torch.ones(model_input.size(0)).to(device=device, dtype=torch.float64) * t_cur
+            diffusion = compute_diffusion(t_cur)
+
+            eps_i = torch.randn_like(x_cur).to(device)
+            deps = eps_i * torch.sqrt(torch.abs(dt))
+
+            v_cur, _, cls_v_cur = model(
+                model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs, cls_token=cls_model_input.to(dtype=_dtype)
+            )
+            v_cur = v_cur.to(torch.float64)
+            cls_v_cur = cls_v_cur.to(torch.float64)
+
+            if add_img_noise:
+                s_cur = get_score_from_velocity(v_cur, model_input, time_input, path_type=path_type)
+                d_cur = v_cur - 0.5 * diffusion * s_cur
+
+                cls_s_cur = get_score_from_velocity(cls_v_cur, cls_model_input, time_input, path_type=path_type)
+                cls_d_cur = cls_v_cur - 0.5 * diffusion * cls_s_cur
+            else:
+                # t<=t_c 去随机阶段：与当前 ODE 逻辑一致，直接 d=v。
+                d_cur = v_cur
+                cls_d_cur = cls_v_cur
+
+            if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+                d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+                d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+                cls_d_cur_cond, cls_d_cur_uncond = cls_d_cur.chunk(2)
+                if cls_cfg > 0:
+                    cls_d_cur = cls_d_cur_uncond + cls_cfg * (cls_d_cur_cond - cls_d_cur_uncond)
+                else:
+                    cls_d_cur = cls_d_cur_cond
+
+            # 图像 latent 有随机扩散噪声；cls/token 仅走漂移（不加随机项）
+            x_next = x_cur + d_cur * dt + torch.sqrt(diffusion) * deps
+            if freeze_after_tc and t_c_v is not None and float(t_cur) <= float(t_c_v) + 1e-9:
+                cls_x_next = cls_frozen
+            else:
+                cls_x_next = cls_x_cur + cls_d_cur * dt
+            if return_trajectory:
+                traj.append(x_next.clone())
+
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                z_mid = x_next.clone()
+                cls_mid = cls_x_next.clone()
+
+    t_cur, t_next = t_steps[-2], t_steps[-1]
+    dt = t_next - t_cur
+    x_cur = x_next
+    cls_x_cur = cls_x_next
+
+    cls_model_input, cls_frozen = _cls_effective_and_freeze(
+        cls_x_cur, cls_frozen, t_cur, t_c_v, freeze_after_tc
+    )
+
+    if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+        model_input = torch.cat([x_cur] * 2, dim=0)
+        cls_model_input = torch.cat([cls_model_input] * 2, dim=0)
+        y_cur = torch.cat([y, y_null], dim=0)
+    else:
+        model_input = x_cur
+        y_cur = y
+    kwargs = dict(y=y_cur)
+    time_input = torch.ones(model_input.size(0)).to(
+        device=device, dtype=torch.float64
+    ) * t_cur
+
+    v_cur, _, cls_v_cur = model(
+        model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs, cls_token=cls_model_input.to(dtype=_dtype)
+    )
+    v_cur = v_cur.to(torch.float64)
+    cls_v_cur = cls_v_cur.to(torch.float64)
+
+    # 最后一步本身无随机项，也与 ODE 对齐使用 velocity 漂移。
+    d_cur = v_cur
+    cls_d_cur = cls_v_cur
+
+    if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+        d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+        d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+        cls_d_cur_cond, cls_d_cur_uncond = cls_d_cur.chunk(2)
+        if cls_cfg > 0:
+            cls_d_cur = cls_d_cur_uncond + cls_cfg * (cls_d_cur_cond - cls_d_cur_uncond)
+        else:
+            cls_d_cur = cls_d_cur_cond
+
+    mean_x = x_cur + dt * d_cur
+    if freeze_after_tc and t_c_v is not None and float(t_cur) <= float(t_c_v) + 1e-9:
+        cls_mean_x = cls_frozen
+    else:
+        cls_mean_x = cls_x_cur + dt * cls_d_cur
+
+    if return_trajectory and return_mid_state:
+        return mean_x, z_mid, cls_mid, traj
+    if return_trajectory:
+        return mean_x, traj
+    if return_mid_state:
+        return mean_x, z_mid, cls_mid
+    return mean_x
+
+
+def euler_maruyama_image_noise_before_tc_sampler(
+        model,
+        latents,
+        y,
+        num_steps=20,
+        heun=False,  # not used, just for compatability
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear",
+        cls_latents=None,
+        args=None,
+        return_mid_state=False,
+        t_mid=0.5,
+        t_c=None,
+        num_steps_before_tc=None,
+        num_steps_after_tc=None,
+        return_cls_final=False,
+        return_trajectory=False,
+        ):
+    """
+    EM 采样变体：
+    - 图像 latent 在 t > t_c 区间引入随机扩散噪声；
+    - 图像 latent 在 t <= t_c 区间不引入随机项（仅漂移）；
+    - cls/token 通道全程不引入随机项。
+    """
+    if cfg_scale > 1.0:
+        y_null = torch.tensor([1000] * y.size(0), device=y.device)
+    _dtype = latents.dtype
+    cls_cfg = getattr(args, "cls_cfg_scale", 0) if args is not None else 0
+
+    t_steps = build_sampling_time_steps(
+        num_steps=num_steps,
+        t_c=t_c,
+        num_steps_before_tc=num_steps_before_tc,
+        num_steps_after_tc=num_steps_after_tc,
+    )
+    freeze_after_tc = _tc_segmented_freeze_cls(t_c, num_steps_before_tc, num_steps_after_tc)
+    t_c_freeze = float(t_c) if freeze_after_tc else None
+    x_next = latents.to(torch.float64)
+    cls_x_next = cls_latents.to(torch.float64)
+    device = x_next.device
+    z_mid = cls_mid = None
+    t_mid = float(t_mid)
+    t_c_v = None if t_c is None else float(t_c)
+    cls_frozen = None
+    traj = [x_next.clone()] if return_trajectory else None
+
+    with torch.no_grad():
+        for t_cur, t_next in zip(t_steps[:-2], t_steps[1:-1]):
+            dt = t_next - t_cur
+            x_cur = x_next
+            cls_x_cur = cls_x_next
+
+            cls_model_input, cls_frozen = _cls_effective_and_freeze(
+                cls_x_cur, cls_frozen, t_cur, t_c_freeze, freeze_after_tc
+            )
+
+            tc, tn = float(t_cur), float(t_next)
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                if abs(tc - t_mid) < abs(tn - t_mid):
+                    z_mid = x_cur.clone()
+                    cls_mid = cls_model_input.clone()
+
+            if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                model_input = torch.cat([x_cur] * 2, dim=0)
+                cls_model_input = torch.cat([cls_model_input] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+            else:
+                model_input = x_cur
+                y_cur = y
+
+            kwargs = dict(y=y_cur)
+            time_input = torch.ones(model_input.size(0)).to(device=device, dtype=torch.float64) * t_cur
+            diffusion = compute_diffusion(t_cur)
+
+            # 跨过/进入 t_c 后关闭图像随机性；t>t_c 区间保留图像噪声
+            add_img_noise = True
+            if t_c_v is not None and float(t_next) <= t_c_v:
+                add_img_noise = False
+
+            eps_i = torch.randn_like(x_cur).to(device) if add_img_noise else torch.zeros_like(x_cur)
+            deps = eps_i * torch.sqrt(torch.abs(dt))
+
+            v_cur, _, cls_v_cur = model(
+                model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs, cls_token=cls_model_input.to(dtype=_dtype)
+            )
+            v_cur = v_cur.to(torch.float64)
+            cls_v_cur = cls_v_cur.to(torch.float64)
+
+            if add_img_noise:
+                s_cur = get_score_from_velocity(v_cur, model_input, time_input, path_type=path_type)
+                d_cur = v_cur - 0.5 * diffusion * s_cur
+
+                cls_s_cur = get_score_from_velocity(cls_v_cur, cls_model_input, time_input, path_type=path_type)
+                cls_d_cur = cls_v_cur - 0.5 * diffusion * cls_s_cur
+            else:
+                # t<=t_c 去随机段：使用显式欧拉 + velocity 漂移（不使用修正漂移项）
+                d_cur = v_cur
+                cls_d_cur = cls_v_cur
+
+            if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+                d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+                d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+                cls_d_cur_cond, cls_d_cur_uncond = cls_d_cur.chunk(2)
+                if cls_cfg > 0:
+                    cls_d_cur = cls_d_cur_uncond + cls_cfg * (cls_d_cur_cond - cls_d_cur_uncond)
+                else:
+                    cls_d_cur = cls_d_cur_cond
+
+            x_next = x_cur + d_cur * dt + torch.sqrt(diffusion) * deps
+            if freeze_after_tc and t_c_freeze is not None and float(t_cur) <= float(t_c_freeze) + 1e-9:
+                cls_x_next = cls_frozen
+            else:
+                cls_x_next = cls_x_cur + cls_d_cur * dt
+            if return_trajectory:
+                traj.append(x_next.clone())
+
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                z_mid = x_next.clone()
+                cls_mid = cls_x_next.clone()
+
+    t_cur, t_next = t_steps[-2], t_steps[-1]
+    dt = t_next - t_cur
+    x_cur = x_next
+    cls_x_cur = cls_x_next
+
+    cls_model_input, cls_frozen = _cls_effective_and_freeze(
+        cls_x_cur, cls_frozen, t_cur, t_c_freeze, freeze_after_tc
+    )
+
+    if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+        model_input = torch.cat([x_cur] * 2, dim=0)
+        cls_model_input = torch.cat([cls_model_input] * 2, dim=0)
+        y_cur = torch.cat([y, y_null], dim=0)
+    else:
+        model_input = x_cur
+        y_cur = y
+    kwargs = dict(y=y_cur)
+    time_input = torch.ones(model_input.size(0)).to(
+        device=device, dtype=torch.float64
+    ) * t_cur
+
+    v_cur, _, cls_v_cur = model(
+        model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs, cls_token=cls_model_input.to(dtype=_dtype)
+    )
+    v_cur = v_cur.to(torch.float64)
+    cls_v_cur = cls_v_cur.to(torch.float64)
+
+    # 最后一步无随机项，保持与 ODE 一致使用 d=v。
+    d_cur = v_cur
+    cls_d_cur = cls_v_cur
+
+    if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+        d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+        d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+        cls_d_cur_cond, cls_d_cur_uncond = cls_d_cur.chunk(2)
+        if cls_cfg > 0:
+            cls_d_cur = cls_d_cur_uncond + cls_cfg * (cls_d_cur_cond - cls_d_cur_uncond)
+        else:
+            cls_d_cur = cls_d_cur_cond
+
+    mean_x = x_cur + dt * d_cur
+    if freeze_after_tc and t_c_freeze is not None and float(t_cur) <= float(t_c_freeze) + 1e-9:
+        cls_mean_x = cls_frozen
+    else:
+        cls_mean_x = cls_x_cur + dt * cls_d_cur
+
+    if return_trajectory and return_mid_state and return_cls_final:
+        return mean_x, z_mid, cls_mid, cls_mean_x, traj
+    if return_trajectory and return_mid_state:
+        return mean_x, z_mid, cls_mid, traj
+    if return_trajectory and return_cls_final:
+        return mean_x, cls_mean_x, traj
+    if return_trajectory:
+        return mean_x, traj
+    if return_mid_state and return_cls_final:
+        return mean_x, z_mid, cls_mid, cls_mean_x
+    if return_mid_state:
+        return mean_x, z_mid, cls_mid
+    if return_cls_final:
+        return mean_x, cls_mean_x
+    return mean_x
+
+
+def euler_ode_sampler(
+        model,
+        latents,
+        y,
+        num_steps=20,
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear",
+        cls_latents=None,
+        args=None,
+        return_mid_state=False,
+        t_mid=0.5,
+        t_c=None,
+        num_steps_before_tc=None,
+        num_steps_after_tc=None,
+        return_trajectory=False,
+        ):
+    """
+    REG 的 ODE 入口：与 SDE 采样器解耦，直接委托 euler_sampler（linspace 1→0 或 t_c 分段，无 t_floor）。
+    """
+    return euler_sampler(
+        model,
+        latents,
+        y,
+        num_steps=num_steps,
+        heun=False,
+        cfg_scale=cfg_scale,
+        guidance_low=guidance_low,
+        guidance_high=guidance_high,
+        path_type=path_type,
+        cls_latents=cls_latents,
+        args=args,
+        return_mid_state=return_mid_state,
+        t_mid=t_mid,
+        t_c=t_c,
+        num_steps_before_tc=num_steps_before_tc,
+        num_steps_after_tc=num_steps_after_tc,
+        return_trajectory=return_trajectory,
+    )
+
+
+def euler_sampler(
+        model,
+        latents,
+        y,
+        num_steps=20,
+        heun=False,
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear",
+        cls_latents=None,
+        args=None,
+        return_mid_state=False,
+        t_mid=0.5,
+        t_c=None,
+        num_steps_before_tc=None,
+        num_steps_after_tc=None,
+        return_trajectory=False,
+        ):
+    """
+    轻量确定性漂移采样（与 glflow 同名同参的前缀兼容：model, latents, y, num_steps, heun, cfg, guidance, path_type, cls_latents, args）。
+
+    - 默认：linspace(1, 0, num_steps+1)，无 t_floor（与原先独立 ODE 一致）。
+    - 可选：同时传入 t_c、num_steps_before_tc、num_steps_after_tc 时，网格为 1→t_c→0；并与 EM 一致在 t≤t_c 段冻结 cls。
+    - 可选：return_mid_state / return_trajectory 供 train.py 与 sample_from_checkpoint 使用。
+
+    REG 的 SiT 需要 cls_token；cls_latents 不可为 None。heun 占位未使用。
+    """
+    if cls_latents is None:
+        raise ValueError(
+            "euler_sampler: 本仓库 REG SiT 需要 cls_token，请传入 cls_latents（例如高斯噪声或训练中的 cls 初值）。"
+        )
+    if cfg_scale > 1.0:
+        y_null = torch.full((y.size(0),), 1000, device=y.device, dtype=y.dtype)
+    else:
+        y_null = None
+    _dtype = latents.dtype
+    cls_cfg = getattr(args, "cls_cfg_scale", 0) if args is not None else 0
+    device = latents.device
+
+    t_steps = _build_euler_sampler_time_steps(
+        num_steps, t_c, num_steps_before_tc, num_steps_after_tc, device
+    )
+    freeze_after_tc = _tc_segmented_freeze_cls(t_c, num_steps_before_tc, num_steps_after_tc)
+    t_c_v = float(t_c) if freeze_after_tc else None
+
+    x_next = latents.to(torch.float64)
+    cls_x_next = cls_latents.to(torch.float64)
+    z_mid = cls_mid = None
+    t_mid = float(t_mid)
+    cls_frozen = None
+    traj = [x_next.clone()] if return_trajectory else None
+
+    with torch.no_grad():
+        for t_cur, t_next in zip(t_steps[:-1], t_steps[1:]):
+            dt = t_next - t_cur
+            x_cur = x_next
+            cls_x_cur = cls_x_next
+
+            cls_model_input, cls_frozen = _cls_effective_and_freeze(
+                cls_x_cur, cls_frozen, t_cur, t_c_v, freeze_after_tc
+            )
+
+            tc, tn = float(t_cur), float(t_next)
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                if abs(tc - t_mid) < abs(tn - t_mid):
+                    z_mid = x_cur.clone()
+                    cls_mid = cls_model_input.clone()
+
+            if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                model_input = torch.cat([x_cur] * 2, dim=0)
+                cls_model_input = torch.cat([cls_model_input] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+            else:
+                model_input = x_cur
+                y_cur = y
+
+            time_input = torch.ones(model_input.size(0), device=device, dtype=torch.float64) * t_cur
+
+            v_cur, _, cls_v_cur = model(
+                model_input.to(dtype=_dtype),
+                time_input.to(dtype=_dtype),
+                y_cur,
+                cls_token=cls_model_input.to(dtype=_dtype),
+            )
+            v_cur = v_cur.to(torch.float64)
+            cls_v_cur = cls_v_cur.to(torch.float64)
+
+            # ODE: follow velocity parameterization directly (d/dt x_t = v_t).
+            # This aligns with velocity training target and avoids extra v->score->drift conversion.
+            d_cur = v_cur
+            cls_d_cur = cls_v_cur
+
+            if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+                d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+                cls_d_cur_cond, cls_d_cur_uncond = cls_d_cur.chunk(2)
+                if cls_cfg > 0:
+                    cls_d_cur = cls_d_cur_uncond + cls_cfg * (cls_d_cur_cond - cls_d_cur_uncond)
+                else:
+                    cls_d_cur = cls_d_cur_cond
+
+            x_next = x_cur + dt * d_cur
+            if freeze_after_tc and t_c_v is not None and float(t_cur) <= float(t_c_v) + 1e-9:
+                cls_x_next = cls_frozen
+            else:
+                cls_x_next = cls_x_cur + dt * cls_d_cur
+
+            if return_trajectory:
+                traj.append(x_next.clone())
+
+            if return_mid_state and z_mid is None and tn <= t_mid <= tc:
+                z_mid = x_next.clone()
+                cls_mid = cls_x_next.clone()
+
+    if return_trajectory and return_mid_state:
+        return x_next, z_mid, cls_mid, traj
+    if return_trajectory:
+        return x_next, traj
+    if return_mid_state:
+        return x_next, z_mid, cls_mid
+    return x_next

back/samples_0.75_new.log

@@ -0,0 +1,11 @@
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] 
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] *****************************************
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] *****************************************
+时间网格：t_c=0.75, 步数 (1→t_c)=100, (t_c→0)=50
+Total number of images that will be sampled: 40192
+
+[rank3]:[W325 16:57:00.799344818 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 3]  using GPU 3 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+[rank1]:[W325 16:57:00.847229448 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 1]  using GPU 1 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+[rank0]:[W325 16:57:01.326116049 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+

conditional-flow-matching/.gitignore

@@ -0,0 +1,175 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+### VisualStudioCode
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+*.code-workspace
+**/.vscode
+
+# JetBrains
+.idea/
+
+# Lightning-Hydra-Template
+configs/local/default.yaml
+data/
+logs/
+wandb/
+.env
+.autoenv
+
+#Vim
+*.sw?
+
+# Slurm
+slurm*.out
+
+# Data and models
+*.pt
+*.h5
+*.h5ad
+*.tar
+*.tar.gz
+*.pkl
+*.npy
+*.npz
+*.csv
+
+# Images
+*.png
+*.svg
+*.gif
+*.jpg
+
+notebooks/figures/
+
+.DS_Store

conditional-flow-matching/.pre-commit-config.yaml

@@ -0,0 +1,99 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      # list of supported hooks: https://pre-commit.com/hooks.html
+      - id: trailing-whitespace
+        exclude: .svg$
+        require_serial: true
+      - id: end-of-file-fixer
+        require_serial: true
+      - id: check-docstring-first
+        require_serial: true
+      - id: check-yaml
+        require_serial: true
+      - id: debug-statements
+        require_serial: true
+      - id: detect-private-key
+        require_serial: true
+      - id: check-executables-have-shebangs
+        require_serial: true
+      - id: check-toml
+        require_serial: true
+      - id: check-case-conflict
+        require_serial: true
+      - id: check-added-large-files
+        require_serial: true
+
+  # python upgrading syntax to newer version
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.17.0
+    hooks:
+      - id: pyupgrade
+        require_serial: true
+        args: [--py38-plus]
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.5.5
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+
+  # python security linter
+  - repo: https://github.com/gitleaks/gitleaks
+    rev: v8.18.2
+    hooks:
+      - id: gitleaks
+
+  # yaml formatting
+  - repo: https://github.com/pre-commit/mirrors-prettier
+    rev: v3.0.0
+    hooks:
+      - id: prettier
+        require_serial: true
+        types: [yaml]
+
+  # shell scripts linter
+  - repo: https://github.com/shellcheck-py/shellcheck-py
+    rev: v0.9.0.5
+    hooks:
+      - id: shellcheck
+        require_serial: true
+        args: ["-e", "SC2102"]
+
+  - repo: https://github.com/pre-commit/mirrors-prettier
+    rev: v4.0.0-alpha.8
+    hooks:
+      - id: prettier
+        # To avoid conflicts, tell prettier to ignore file types
+        # that ruff already handles.
+        exclude_types: [python]
+
+  # word spelling linter
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.2.5
+    hooks:
+      - id: codespell
+        require_serial: true
+        args:
+          - --skip=logs/**,data/**,*.ipynb
+          - --ignore-words-list=ot,hist
+
+  # jupyter notebook linting
+  - repo: https://github.com/nbQA-dev/nbQA
+    rev: 1.7.0
+    hooks:
+      - id: nbqa-black
+        args: ["--line-length=99"]
+        require_serial: true
+      - id: nbqa-isort
+        args: ["--profile=black"]
+        require_serial: true
+      - id: nbqa-flake8
+        args:
+          [
+            "--extend-ignore=E203,E402,E501,F401,F841,F821,F403,F405,F811",
+            "--exclude=logs/*,data/*,notebooks/*",
+          ]
+        require_serial: true

conditional-flow-matching/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Alexander Tong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

conditional-flow-matching/README.md

@@ -0,0 +1,230 @@
+<div align="center">
+
+# TorchCFM: a Conditional Flow Matching library
+
+<!---[![Conference](http://img.shields.io/badge/AnyConference-year-4b44ce.svg)](https://papers.nips.cc/paper/2020) -->
+
+<!---[![contributors](https://img.shields.io/github/contributors/atong01/conditional-flow-matching.svg)](https://github.com/atong01/conditional-flow-matching/graphs/contributors) -->
+
+[![OT-CFM Preprint](http://img.shields.io/badge/paper-arxiv.2302.00482-B31B1B.svg)](https://arxiv.org/abs/2302.00482)
+[![SF2M Preprint](http://img.shields.io/badge/paper-arxiv.2307.03672-B31B1B.svg)](https://arxiv.org/abs/2307.03672)
+[![pytorch](https://img.shields.io/badge/PyTorch_1.8+-ee4c2c?logo=pytorch&logoColor=white)](https://pytorch.org/get-started/locally/)
+[![lightning](https://img.shields.io/badge/-Lightning_1.6+-792ee5?logo=pytorchlightning&logoColor=white)](https://pytorchlightning.ai/)
+[![hydra](https://img.shields.io/badge/Config-Hydra_1.2-89b8cd)](https://hydra.cc/)
+[![black](https://img.shields.io/badge/Code%20Style-Black-black.svg?labelColor=gray)](https://black.readthedocs.io/en/stable/)
+[![pre-commit](https://img.shields.io/badge/Pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://github.com/pre-commit/pre-commit)
+[![tests](https://github.com/atong01/conditional-flow-matching/actions/workflows/test.yaml/badge.svg)](https://github.com/atong01/conditional-flow-matching/actions/workflows/test.yaml)
+[![codecov](https://codecov.io/gh/atong01/conditional-flow-matching/branch/main/graph/badge.svg)](https://codecov.io/gh/atong01/conditional-flow-matching/)
+[![code-quality](https://github.com/atong01/conditional-flow-matching/actions/workflows/code-quality-main.yaml/badge.svg)](https://github.com/atong01/conditional-flow-matching/actions/workflows/code-quality-main.yaml)
+[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/atong01/conditional-flow-matching#license)
+<a href="https://github.com/ashleve/lightning-hydra-template"><img alt="Template" src="https://img.shields.io/badge/-Lightning--Hydra--Template-017F2F?style=flat&logo=github&labelColor=gray"></a>
+[![Downloads](https://static.pepy.tech/badge/torchcfm)](https://pepy.tech/project/torchcfm)
+[![Downloads](https://static.pepy.tech/badge/torchcfm/month)](https://pepy.tech/project/torchcfm)
+
+</div>
+
+## Description
+
+Conditional Flow Matching (CFM) is a fast way to train continuous normalizing flow (CNF) models. CFM is a simulation-free training objective for continuous normalizing flows that allows conditional generative modeling and speeds up training and inference. CFM's performance closes the gap between CNFs and diffusion models. To spread its use within the machine learning community, we have built a library focused on Flow Matching methods: TorchCFM. TorchCFM is a library showing how Flow Matching methods can be trained and used to deal with image generation, single-cell dynamics, tabular data and soon SO(3) data.
+
+<p align="center">
+<img src="assets/169_generated_samples_otcfm.png" width="600"/>
+<img src="assets/8gaussians-to-moons.gif" />
+</p>
+
+The density, vector field, and trajectories of simulation-free CNF training schemes: mapping 8 Gaussians to two moons (above) and a single Gaussian to two moons (below). Action matching with the same architecture (3x64 MLP with SeLU activations) underfits with the ReLU, SiLU, and SiLU activations as suggested in the [example code](https://github.com/necludov/jam), but it seems to fit better under our training setup (Action-Matching (Swish)).
+
+The models to produce the GIFs are stored in `examples/models` and can be visualized with this notebook: [![notebook](https://img.shields.io/static/v1?label=Run%20in&message=Google%20Colab&color=orange&logo=Google%20Cloud)](https://colab.research.google.com/github/atong01/conditional-flow-matching/blob/master/examples/notebooks/model-comparison-plotting.ipynb).
+
+We also have included an example of unconditional MNIST generation in `examples/notebooks/mnist_example.ipynb` for both deterministic and stochastic generation. [![notebook](https://img.shields.io/static/v1?label=Run%20in&message=Google%20Colab&color=orange&logo=Google%20Cloud)](https://colab.research.google.com/github/atong01/conditional-flow-matching/blob/master/examples/notebooks/mnist_example.ipynb).
+
+## The torchcfm Package
+
+In our version 1 update we have extracted implementations of the relevant flow matching variants into a package `torchcfm`. This allows abstraction of the choice of the conditional distribution `q(z)`. `torchcfm` supplies the following loss functions:
+
+- `ConditionalFlowMatcher`: $z = (x_0, x_1)$, $q(z) = q(x_0) q(x_1)$
+- `ExactOptimalTransportConditionalFlowMatcher`: $z = (x_0, x_1)$, $q(z) = \\pi(x_0, x_1)$ where $\\pi$ is an exact optimal transport joint. This is used in \[Tong et al. 2023a\] and \[Poolidan et al. 2023\] as "OT-CFM" and "Multisample FM with Batch OT" respectively.
+- `TargetConditionalFlowMatcher`: $z = x_1$, $q(z) = q(x_1)$ as defined in Lipman et al. 2023, learns a flow from a standard normal Gaussian to data using conditional flows which optimally transport the Gaussian to the datapoint (Note that this does not result in the marginal flow being optimal transport).
+- `SchrodingerBridgeConditionalFlowMatcher`: $z = (x_0, x_1)$, $q(z) = \\pi\_\\epsilon(x_0, x_1)$ where $\\pi\_\\epsilon$ is an entropically regularized OT plan, although in practice this is often approximated by a minibatch OT plan (See Tong et al. 2023b). The flow-matching variant of this where the marginals are equivalent to the Schrodinger Bridge marginals is known as `SB-CFM` \[Tong et al. 2023a\]. When the score is also known and the bridge is stochastic is called \[SF\]2M \[Tong et al. 2023b\]
+- `VariancePreservingConditionalFlowMatcher`: $z = (x_0, x_1)$ $q(z) = q(x_0) q(x_1)$ but with conditional Gaussian probability paths which preserve variance over time using a trigonometric interpolation as presented in \[Albergo et al. 2023a\].
+
+## How to cite
+
+This repository contains the code to reproduce the main experiments and illustrations of two preprints:
+
+- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://arxiv.org/abs/2302.00482). We introduce **Optimal Transport Conditional Flow Matching** (OT-CFM), a CFM variant that approximates the dynamical formulation of optimal transport (OT). Based on OT theory, OT-CFM leverages the static optimal transport plan as well as the optimal probability paths and vector fields to approximate dynamic OT.
+- [Simulation-free Schrödinger bridges via score and flow matching](https://arxiv.org/abs/2307.03672). We propose **Simulation-Free Score and Flow Matching** (\[SF\]<sup>2</sup>M). \[SF\]<sup>2</sup>M leverages OT-CFM as well as score-based methods to approximate Schrödinger bridges, a stochastic version of optimal transport.
+
+If you find this code useful in your research, please cite the following papers (expand for BibTeX):
+
+<details>
+<summary>
+A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, K. Fatras, G. Wolf, Y. Bengio. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport, 2023.
+</summary>
+
+```bibtex
+@article{tong2024improving,
+title={Improving and generalizing flow-based generative models with minibatch optimal transport},
+author={Alexander Tong and Kilian FATRAS and Nikolay Malkin and Guillaume Huguet and Yanlei Zhang and Jarrid Rector-Brooks and Guy Wolf and Yoshua Bengio},
+journal={Transactions on Machine Learning Research},
+issn={2835-8856},
+year={2024},
+url={https://openreview.net/forum?id=CD9Snc73AW},
+note={Expert Certification}
+}
+```
+
+</details>
+
+<details>
+<summary>
+A. Tong, N. Malkin, K. Fatras, L. Atanackovic, Y. Zhang, G. Huguet, G. Wolf, Y. Bengio. Simulation-Free Schrödinger Bridges via Score and Flow Matching, 2023.
+</summary>
+
+```bibtex
+@article{tong2023simulation,
+   title={Simulation-Free Schr{\"o}dinger Bridges via Score and Flow Matching},
+   author={Tong, Alexander and Malkin, Nikolay and Fatras, Kilian and Atanackovic, Lazar and Zhang, Yanlei and Huguet, Guillaume and Wolf, Guy and Bengio, Yoshua},
+   year={2023},
+   journal={arXiv preprint 2307.03672}
+}
+```
+
+</details>
+
+## V0 -> V1
+
+Major Changes:
+
+- **Added cifar10 examples with an FID of 3.5**
+- Added code for the new Simulation-free Score and Flow Matching (SF)2M preprint
+- Created `torchcfm` pip installable package
+- Moved `pytorch-lightning` implementation and experiments to `runner` directory
+- Moved `notebooks` -> `examples`
+- Added image generation implementation in both lightning and a notebook in `examples`
+
+## Implemented papers
+
+List of implemented papers:
+
+- Flow Matching for Generative Modeling (Lipman et al. 2023) [Paper](https://openreview.net/forum?id=PqvMRDCJT9t)
+- Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow (Liu et al. 2023) [Paper](https://openreview.net/forum?id=XVjTT1nw5z) [Code](https://github.com/gnobitab/RectifiedFlow.git)
+- Building Normalizing Flows with Stochastic Interpolants (Albergo et al. 2023a) [Paper](https://openreview.net/forum?id=li7qeBbCR1t)
+- Action Matching: Learning Stochastic Dynamics From Samples (Neklyudov et al. 2022) [Paper](https://arxiv.org/abs/2210.06662) [Code](https://github.com/necludov/jam)
+- Concurrent work to our OT-CFM method: Multisample Flow Matching: Straightening Flows with Minibatch Couplings (Pooladian et al. 2023) [Paper](https://arxiv.org/abs/2304.14772)
+- Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees (Jolicoeur-Martineau et al.) [Paper](https://arxiv.org/abs/2309.09968) [Code](https://github.com/SamsungSAILMontreal/ForestDiffusion)
+- Soon: SE(3)-Stochastic Flow Matching for Protein Backbone Generation (Bose et al.) [Paper](https://arxiv.org/abs/2310.02391)
+
+## How to run
+
+Run a simple minimal example here [![Run in Google Colab](https://img.shields.io/static/v1?label=Run%20in&message=Google%20Colab&color=orange&logo=Google%20Cloud)](https://colab.research.google.com/github/atong01/conditional-flow-matching/blob/master/examples/notebooks/training-8gaussians-to-moons.ipynb). Or install the more efficient code locally with these steps.
+
+TorchCFM is now on [pypi](https://pypi.org/project/torchcfm/)! You can install it with:
+
+```bash
+pip install torchcfm
+```
+
+To use the full library with the different examples, you can install dependencies:
+
+```bash
+# clone project
+git clone https://github.com/atong01/conditional-flow-matching.git
+cd conditional-flow-matching
+
+# [OPTIONAL] create conda environment
+conda create -n torchcfm python=3.10
+conda activate torchcfm
+
+# install pytorch according to instructions
+# https://pytorch.org/get-started/
+
+# install requirements
+pip install -r requirements.txt
+
+# install torchcfm
+pip install -e .
+```
+
+To run our jupyter notebooks, use the following commands after installing our package.
+
+```bash
+# install ipykernel
+conda install -c anaconda ipykernel
+
+# install conda env in jupyter notebook
+python -m ipykernel install --user --name=torchcfm
+
+# launch our notebooks with the torchcfm kernel
+```
+
+## Project Structure
+
+The directory structure looks like this:
+
+```
+
+│
+├── examples              <- Jupyter notebooks
+|   ├── cifar10           <- Cifar10 experiments
+│   ├── notebooks         <- Diverse examples with notebooks
+│
+│── runner                    <- Everything related to the original version (V0) of the library
+│
+|── torchcfm                  <- Code base of our Flow Matching methods
+|   ├── conditional_flow_matching.py      <- CFM classes
+│   ├── models                            <- Model architectures
+│   │   ├── models                           <- Models for 2D examples
+│   │   ├── Unet                             <- Unet models for image examples
+|
+├── .gitignore                <- List of files ignored by git
+├── .pre-commit-config.yaml   <- Configuration of pre-commit hooks for code formatting
+├── pyproject.toml            <- Configuration options for testing and linting
+├── requirements.txt          <- File for installing python dependencies
+├── setup.py                  <- File for installing project as a package
+└── README.md
+```
+
+## ❤️  Code Contributions
+
+This toolbox has been created and is maintained by
+
+- [Alexander Tong](http://alextong.net)
+- [Kilian Fatras](http://kilianfatras.github.io)
+
+It was initiated from a larger private codebase which loses the original commit history which contains work from other authors of the papers.
+
+Before making an issue, please verify that:
+
+- The problem still exists on the current `main` branch.
+- Your python dependencies are updated to recent versions.
+
+Suggestions for improvements are always welcome!
+
+## License
+
+Conditional-Flow-Matching is licensed under the MIT License.
+
+```
+MIT License
+
+Copyright (c) 2023 Alexander Tong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+```

conditional-flow-matching/pyproject.toml

@@ -0,0 +1,39 @@
+[tool.pytest.ini_options]
+addopts = [
+  "--color=yes",
+  "--durations=0",
+  "--strict-markers",
+  "--doctest-modules",
+]
+filterwarnings = [
+  "ignore::DeprecationWarning",
+  "ignore::UserWarning",
+]
+log_cli = "True"
+markers = [
+  "slow: slow tests",
+]
+minversion = "6.0"
+testpaths = "tests/"
+
+[tool.coverage.report]
+exclude_lines = [
+    "pragma: nocover",
+    "raise NotImplementedError",
+    "raise NotImplementedError()",
+    "if __name__ == .__main__.:",
+]
+
+[tool.ruff]
+line-length = 99
+
+[tool.ruff.lint]
+ignore = ["C901", "E501", "E741", "W605", "C408", "E402"]
+select = ["C", "E", "F", "I", "W"]
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["E402", "F401", "F403", "F811"]
+
+[tool.ruff.lint.isort]
+known-first-party = ["src"]
+known-third-party = ["torch", "transformers", "wandb"]

conditional-flow-matching/requirements.txt

@@ -0,0 +1,15 @@
+torch>=1.11.0
+torchvision>=0.11.0
+
+lightning-bolts
+matplotlib
+numpy
+scipy
+scikit-learn
+scprep
+scanpy
+torchdyn>=1.0.6   # 1.0.4 is broken on pypi
+pot
+torchdiffeq
+absl-py
+clean-fid

conditional-flow-matching/runner-requirements.txt

@@ -0,0 +1,52 @@
+# Note if using Conda it is recommended to install torch separately.
+# For most of testing the following commands were run to set up the environment
+# This was tested with torch==1.12.1
+# conda create -n ti-env python=3.10
+# conda activate ti-env
+# pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
+# pip install -r requirements.txt
+# --------- pytorch --------- #
+torch>=1.11.0,<2.0.0
+torchvision>=0.11.0
+pytorch-lightning==1.8.3.post2
+torchmetrics==0.11.0
+
+# --------- hydra --------- #
+hydra-core==1.2.0
+hydra-colorlog==1.2.0
+hydra-optuna-sweeper==1.2.0
+# hydra-submitit-launcher
+
+# --------- loggers --------- #
+wandb
+# neptune-client
+# mlflow
+# comet-ml
+
+# --------- others --------- #
+black
+isort
+flake8
+Flake8-pyproject  # for configuration via pyproject
+pyrootutils       # standardizing the project root setup
+pre-commit        # hooks for applying linters on commit
+rich              # beautiful text formatting in terminal
+pytest            # tests
+# sh              # for running bash commands in some tests (linux/macos only)
+
+
+# --------- pkg reqs -------- #
+lightning-bolts
+matplotlib
+numpy
+scipy
+scikit-learn
+scprep
+scanpy
+timm
+torchdyn>=1.0.5   # 1.0.4 is broken on pypi
+pot
+
+# --------- notebook reqs -------- #
+seaborn>=0.12.2
+pandas>=2.2.2

conditional-flow-matching/setup.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+
+import os
+
+from setuptools import find_packages, setup
+
+install_requires = [
+    "torch>=1.11.0",
+    "matplotlib",
+    "numpy",  # Due to pandas incompatibility
+    "scipy",
+    "scikit-learn",
+    "torchdyn>=1.0.6",
+    "pot",
+    "torchdiffeq",
+    "absl-py",
+    "pandas>=2.2.2",
+]
+
+version_py = os.path.join(os.path.dirname(__file__), "torchcfm", "version.py")
+version = open(version_py).read().strip().split("=")[-1].replace('"', "").strip()
+readme = open("README.md", encoding="utf8").read()
+setup(
+    name="torchcfm",
+    version=version,
+    description="Conditional Flow Matching for Fast Continuous Normalizing Flow Training.",
+    author="Alexander Tong, Kilian Fatras",
+    author_email="alexandertongdev@gmail.com",
+    url="https://github.com/atong01/conditional-flow-matching",
+    install_requires=install_requires,
+    license="MIT",
+    long_description=readme,
+    long_description_content_type="text/markdown",
+    packages=find_packages(exclude=["tests", "tests.*"]),
+    extras_require={"forest-flow": ["xgboost", "scikit-learn", "ForestDiffusion"]},
+)