Add files using upload-large-folder tool

New/REG/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Sihyun Yu
+
+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.

New/REG/dataset.py

@@ -0,0 +1,165 @@
+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):
+    """
+    New/REG 的数据集对齐：支持和 REG 相同的两种模式：
+    1) 仅加载 VAE latent + vae-sd features（与原本 New 行为一致）
+    2) 若提供 semantic_features_dir，则加载预处理 semantic cls token，并在 __getitem__ 返回 4 元组：
+       (raw_image_latent, x_latent, semantic_cls_token, y_label)
+    """
+
+    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")
+        self.semantic_features_dir = semantic_features_dir
+        self.use_preprocessed_semantic = False
+
+        # 1) 预处理 semantic 模式（与 REG/dataset.py 一致）
+        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.use_preprocessed_semantic = True
+
+        if self.use_preprocessed_semantic:
+            if self.semantic_features_dir is None:
+                raise ValueError("semantic_features_dir is None but use_preprocessed_semantic=True")
+
+            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")
+
+            # semantic 模式下无需额外维护 vae-sd feature 列表
+            return
+
+        # 2) 非预处理 semantic 模式：VAE latent + vae-sd features
+        self.features_dir = os.path.join(data_dir, "vae-sd")
+
+        # images
+        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
+        )
+
+        # features
+        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
+        )
+
+        # labels
+        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)
+
+            y = torch.tensor(self.labels[idx])
+            return (
+                torch.from_numpy(image).float(),
+                torch.from_numpy(image).float(),
+                torch.from_numpy(semantic_features).float(),
+                y,
+            )
+
+        # ---- non-preprocessed mode ----
+        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])

New/REG/evaluations/README.md

@@ -0,0 +1,72 @@
+# Evaluations
+
+To compare different generative models, we use FID, sFID, Precision, Recall, and Inception Score. These metrics can all be calculated using batches of samples, which we store in `.npz` (numpy) files.
+
+# Download batches
+
+We provide pre-computed sample batches for the reference datasets, our diffusion models, and several baselines we compare against. These are all stored in `.npz` format.
+
+Reference dataset batches contain pre-computed statistics over the whole dataset, as well as 10,000 images for computing Precision and Recall. All other batches contain 50,000 images which can be used to compute statistics and Precision/Recall.
+
+Here are links to download all of the sample and reference batches:
+
+ * LSUN
+   * LSUN bedroom: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/VIRTUAL_lsun_bedroom256.npz)
+     * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/admnet_dropout_lsun_bedroom.npz)
+     * [DDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/ddpm_lsun_bedroom.npz)
+     * [IDDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/iddpm_lsun_bedroom.npz)
+     * [StyleGAN](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/stylegan_lsun_bedroom.npz)
+   * LSUN cat: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/VIRTUAL_lsun_cat256.npz)
+     * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/admnet_dropout_lsun_cat.npz)
+     * [StyleGAN2](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/stylegan2_lsun_cat.npz)
+   * LSUN horse: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/VIRTUAL_lsun_horse256.npz)
+     * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/admnet_dropout_lsun_horse.npz)
+     * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/admnet_lsun_horse.npz)
+
+ * ImageNet
+   * ImageNet 64x64: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/VIRTUAL_imagenet64_labeled.npz)
+     * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/admnet_imagenet64.npz)
+     * [IDDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/iddpm_imagenet64.npz)
+     * [BigGAN](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/biggan_deep_imagenet64.npz)
+   * ImageNet 128x128: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/VIRTUAL_imagenet128_labeled.npz)
+     * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_imagenet128.npz)
+     * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_guided_imagenet128.npz)
+     * [ADM-G, 25 steps](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_guided_25step_imagenet128.npz)
+     * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/biggan_deep_trunc1_imagenet128.npz)
+   * ImageNet 256x256: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz)
+     * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_imagenet256.npz)
+     * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_imagenet256.npz)
+     * [ADM-G, 25 step](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_25step_imagenet256.npz)
+     * [ADM-G + ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_upsampled_imagenet256.npz)
+     * [ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_upsampled_imagenet256.npz)
+     * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/biggan_deep_trunc1_imagenet256.npz)
+   * ImageNet 512x512: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/VIRTUAL_imagenet512.npz)
+     * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_imagenet512.npz)
+     * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_imagenet512.npz)
+     * [ADM-G, 25 step](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_25step_imagenet512.npz)
+     * [ADM-G + ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_upsampled_imagenet512.npz)
+     * [ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_upsampled_imagenet512.npz)
+     * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/biggan_deep_trunc1_imagenet512.npz)
+
+# Run evaluations
+
+First, generate or download a batch of samples and download the corresponding reference batch for the given dataset. For this example, we'll use ImageNet 256x256, so the refernce batch is `VIRTUAL_imagenet256_labeled.npz` and we can use the sample batch `admnet_guided_upsampled_imagenet256.npz`.
+
+Next, run the `evaluator.py` script. The requirements of this script can be found in [requirements.txt](requirements.txt). Pass two arguments to the script: the reference batch and the sample batch. The script will download the InceptionV3 model used for evaluations into the current working directory (if it is not already present). This file is roughly 100MB.
+
+The output of the script will look something like this, where the first `...` is a bunch of verbose TensorFlow logging:
+
+```
+$ python evaluator.py VIRTUAL_imagenet256_labeled.npz admnet_guided_upsampled_imagenet256.npz
+...
+computing reference batch activations...
+computing/reading reference batch statistics...
+computing sample batch activations...
+computing/reading sample batch statistics...
+Computing evaluations...
+Inception Score: 215.8370361328125
+FID: 3.9425574129223264
+sFID: 6.140433703346162
+Precision: 0.8265
+Recall: 0.5309
+```

New/REG/evaluations/evaluator.py

@@ -0,0 +1,679 @@
+import argparse
+import io
+import os
+import random
+import warnings
+import zipfile
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from functools import partial
+from multiprocessing import cpu_count
+from multiprocessing.pool import ThreadPool
+from typing import Iterable, Optional, Tuple
+
+import numpy as np
+import requests
+import tensorflow.compat.v1 as tf
+from scipy import linalg
+from tqdm.auto import tqdm
+
+INCEPTION_V3_URL = "https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/classify_image_graph_def.pb"
+INCEPTION_V3_PATH = "classify_image_graph_def.pb"
+
+FID_POOL_NAME = "pool_3:0"
+FID_SPATIAL_NAME = "mixed_6/conv:0"
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ref_batch", help="path to reference batch npz file")
+    parser.add_argument("--sample_batch", help="path to sample batch npz file")
+    parser.add_argument("--save_path", help="path to sample batch npz file")
+    parser.add_argument("--cfg_cond", default=1, type=int)
+    parser.add_argument("--step", default=1, type=int)
+    parser.add_argument("--cfg", default=1.0, type=float)
+    parser.add_argument("--cls_cfg", default=1.0, type=float)
+    parser.add_argument("--gh", default=1.0, type=float)
+    parser.add_argument("--num_steps", default=250, type=int)
+    args = parser.parse_args()
+
+    if not os.path.exists(args.save_path):
+        os.mkdir(args.save_path)
+
+
+    config = tf.ConfigProto(
+        allow_soft_placement=True  # allows DecodeJpeg to run on CPU in Inception graph
+    )
+    config.gpu_options.allow_growth = True
+    evaluator = Evaluator(tf.Session(config=config))
+
+    print("warming up TensorFlow...")
+    # This will cause TF to print a bunch of verbose stuff now rather
+    # than after the next print(), to help prevent confusion.
+    evaluator.warmup()
+
+    print("computing reference batch activations...")
+    ref_acts = evaluator.read_activations(args.ref_batch)
+    print("computing/reading reference batch statistics...")
+    ref_stats, ref_stats_spatial = evaluator.read_statistics(args.ref_batch, ref_acts)
+
+    print("computing sample batch activations...")
+    sample_acts = evaluator.read_activations(args.sample_batch)
+    print("computing/reading sample batch statistics...")
+    sample_stats, sample_stats_spatial = evaluator.read_statistics(args.sample_batch, sample_acts)
+
+    print("Computing evaluations...")
+    Inception_Score = evaluator.compute_inception_score(sample_acts[0])
+    FID = sample_stats.frechet_distance(ref_stats)
+    sFID = sample_stats_spatial.frechet_distance(ref_stats_spatial)
+    prec, recall = evaluator.compute_prec_recall(ref_acts[0], sample_acts[0])
+
+    print("Inception Score:", Inception_Score)
+    print("FID:", FID)
+    print("sFID:", sFID)
+    print("Precision:", prec)
+    print("Recall:", recall)
+
+    if args.cfg_cond:
+        file_path = args.save_path + str(args.num_steps) + str(args.step) + str(args.cfg) + str(args.gh) + str(args.cls_cfg)+ "cfg_cond_true.txt"
+    else:
+        file_path = args.save_path + str(args.num_steps) + str(args.step) + str(args.cfg) + str(args.gh) + str(args.cls_cfg)+ "cfg_cond_false.txt"
+    with open(file_path, "w") as file:
+        file.write("Inception Score: {}\n".format(Inception_Score))
+        file.write("FID: {}\n".format(FID))
+        file.write("sFID: {}\n".format(sFID))
+        file.write("Precision: {}\n".format(prec))
+        file.write("Recall: {}\n".format(recall))
+
+
+class InvalidFIDException(Exception):
+    pass
+
+
+class FIDStatistics:
+    def __init__(self, mu: np.ndarray, sigma: np.ndarray):
+        self.mu = mu
+        self.sigma = sigma
+
+    def frechet_distance(self, other, eps=1e-6):
+        """
+        Compute the Frechet distance between two sets of statistics.
+        """
+        # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L132
+        mu1, sigma1 = self.mu, self.sigma
+        mu2, sigma2 = other.mu, other.sigma
+
+        mu1 = np.atleast_1d(mu1)
+        mu2 = np.atleast_1d(mu2)
+
+        sigma1 = np.atleast_2d(sigma1)
+        sigma2 = np.atleast_2d(sigma2)
+
+        assert (
+            mu1.shape == mu2.shape
+        ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}"
+        assert (
+            sigma1.shape == sigma2.shape
+        ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}"
+
+        diff = mu1 - mu2
+
+        # product might be almost singular
+        covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+        if not np.isfinite(covmean).all():
+            msg = (
+                "fid calculation produces singular product; adding %s to diagonal of cov estimates"
+                % eps
+            )
+            warnings.warn(msg)
+            offset = np.eye(sigma1.shape[0]) * eps
+            covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+        # numerical error might give slight imaginary component
+        if np.iscomplexobj(covmean):
+            if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+                m = np.max(np.abs(covmean.imag))
+                raise ValueError("Imaginary component {}".format(m))
+            covmean = covmean.real
+
+        tr_covmean = np.trace(covmean)
+
+        return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
+
+
+class Evaluator:
+    def __init__(
+        self,
+        session,
+        batch_size=64,
+        softmax_batch_size=512,
+    ):
+        self.sess = session
+        self.batch_size = batch_size
+        self.softmax_batch_size = softmax_batch_size
+        self.manifold_estimator = ManifoldEstimator(session)
+        with self.sess.graph.as_default():
+            self.image_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
+            self.softmax_input = tf.placeholder(tf.float32, shape=[None, 2048])
+            self.pool_features, self.spatial_features = _create_feature_graph(self.image_input)
+            self.softmax = _create_softmax_graph(self.softmax_input)
+
+    def warmup(self):
+        self.compute_activations(np.zeros([1, 8, 64, 64, 3]))
+
+    def read_activations(self, npz_path: str) -> Tuple[np.ndarray, np.ndarray]:
+        with open_npz_array(npz_path, "arr_0") as reader:
+            return self.compute_activations(reader.read_batches(self.batch_size))
+
+    def compute_activations(self, batches: Iterable[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Compute image features for downstream evals.
+
+        :param batches: a iterator over NHWC numpy arrays in [0, 255].
+        :return: a tuple of numpy arrays of shape [N x X], where X is a feature
+                 dimension. The tuple is (pool_3, spatial).
+        """
+        preds = []
+        spatial_preds = []
+        for batch in tqdm(batches):
+            batch = batch.astype(np.float32)
+            pred, spatial_pred = self.sess.run(
+                [self.pool_features, self.spatial_features], {self.image_input: batch}
+            )
+            preds.append(pred.reshape([pred.shape[0], -1]))
+            spatial_preds.append(spatial_pred.reshape([spatial_pred.shape[0], -1]))
+        return (
+            np.concatenate(preds, axis=0),
+            np.concatenate(spatial_preds, axis=0),
+        )
+
+    def read_statistics(
+        self, npz_path: str, activations: Tuple[np.ndarray, np.ndarray]
+    ) -> Tuple[FIDStatistics, FIDStatistics]:
+        obj = np.load(npz_path)
+        if "mu" in list(obj.keys()):
+            return FIDStatistics(obj["mu"], obj["sigma"]), FIDStatistics(
+                obj["mu_s"], obj["sigma_s"]
+            )
+        return tuple(self.compute_statistics(x) for x in activations)
+
+    def compute_statistics(self, activations: np.ndarray) -> FIDStatistics:
+        mu = np.mean(activations, axis=0)
+        sigma = np.cov(activations, rowvar=False)
+        return FIDStatistics(mu, sigma)
+
+    def compute_inception_score(self, activations: np.ndarray, split_size: int = 5000) -> float:
+        softmax_out = []
+        for i in range(0, len(activations), self.softmax_batch_size):
+            acts = activations[i : i + self.softmax_batch_size]
+            softmax_out.append(self.sess.run(self.softmax, feed_dict={self.softmax_input: acts}))
+        preds = np.concatenate(softmax_out, axis=0)
+        # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46
+        scores = []
+        for i in range(0, len(preds), split_size):
+            part = preds[i : i + split_size]
+            kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))
+            kl = np.mean(np.sum(kl, 1))
+            scores.append(np.exp(kl))
+        return float(np.mean(scores))
+
+    def compute_prec_recall(
+        self, activations_ref: np.ndarray, activations_sample: np.ndarray
+    ) -> Tuple[float, float]:
+        radii_1 = self.manifold_estimator.manifold_radii(activations_ref)
+        radii_2 = self.manifold_estimator.manifold_radii(activations_sample)
+        pr = self.manifold_estimator.evaluate_pr(
+            activations_ref, radii_1, activations_sample, radii_2
+        )
+        return (float(pr[0][0]), float(pr[1][0]))
+
+
+class ManifoldEstimator:
+    """
+    A helper for comparing manifolds of feature vectors.
+
+    Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L57
+    """
+
+    def __init__(
+        self,
+        session,
+        row_batch_size=10000,
+        col_batch_size=10000,
+        nhood_sizes=(3,),
+        clamp_to_percentile=None,
+        eps=1e-5,
+    ):
+        """
+        Estimate the manifold of given feature vectors.
+
+        :param session: the TensorFlow session.
+        :param row_batch_size: row batch size to compute pairwise distances
+                               (parameter to trade-off between memory usage and performance).
+        :param col_batch_size: column batch size to compute pairwise distances.
+        :param nhood_sizes: number of neighbors used to estimate the manifold.
+        :param clamp_to_percentile: prune hyperspheres that have radius larger than
+                                    the given percentile.
+        :param eps: small number for numerical stability.
+        """
+        self.distance_block = DistanceBlock(session)
+        self.row_batch_size = row_batch_size
+        self.col_batch_size = col_batch_size
+        self.nhood_sizes = nhood_sizes
+        self.num_nhoods = len(nhood_sizes)
+        self.clamp_to_percentile = clamp_to_percentile
+        self.eps = eps
+
+    def warmup(self):
+        feats, radii = (
+            np.zeros([1, 2048], dtype=np.float32),
+            np.zeros([1, 1], dtype=np.float32),
+        )
+        self.evaluate_pr(feats, radii, feats, radii)
+
+    def manifold_radii(self, features: np.ndarray) -> np.ndarray:
+        num_images = len(features)
+
+        # Estimate manifold of features by calculating distances to k-NN of each sample.
+        radii = np.zeros([num_images, self.num_nhoods], dtype=np.float32)
+        distance_batch = np.zeros([self.row_batch_size, num_images], dtype=np.float32)
+        seq = np.arange(max(self.nhood_sizes) + 1, dtype=np.int32)
+
+        for begin1 in range(0, num_images, self.row_batch_size):
+            end1 = min(begin1 + self.row_batch_size, num_images)
+            row_batch = features[begin1:end1]
+
+            for begin2 in range(0, num_images, self.col_batch_size):
+                end2 = min(begin2 + self.col_batch_size, num_images)
+                col_batch = features[begin2:end2]
+
+                # Compute distances between batches.
+                distance_batch[
+                    0 : end1 - begin1, begin2:end2
+                ] = self.distance_block.pairwise_distances(row_batch, col_batch)
+
+            # Find the k-nearest neighbor from the current batch.
+            radii[begin1:end1, :] = np.concatenate(
+                [
+                    x[:, self.nhood_sizes]
+                    for x in _numpy_partition(distance_batch[0 : end1 - begin1, :], seq, axis=1)
+                ],
+                axis=0,
+            )
+
+        if self.clamp_to_percentile is not None:
+            max_distances = np.percentile(radii, self.clamp_to_percentile, axis=0)
+            radii[radii > max_distances] = 0
+        return radii
+
+    def evaluate(self, features: np.ndarray, radii: np.ndarray, eval_features: np.ndarray):
+        """
+        Evaluate if new feature vectors are at the manifold.
+        """
+        num_eval_images = eval_features.shape[0]
+        num_ref_images = radii.shape[0]
+        distance_batch = np.zeros([self.row_batch_size, num_ref_images], dtype=np.float32)
+        batch_predictions = np.zeros([num_eval_images, self.num_nhoods], dtype=np.int32)
+        max_realism_score = np.zeros([num_eval_images], dtype=np.float32)
+        nearest_indices = np.zeros([num_eval_images], dtype=np.int32)
+
+        for begin1 in range(0, num_eval_images, self.row_batch_size):
+            end1 = min(begin1 + self.row_batch_size, num_eval_images)
+            feature_batch = eval_features[begin1:end1]
+
+            for begin2 in range(0, num_ref_images, self.col_batch_size):
+                end2 = min(begin2 + self.col_batch_size, num_ref_images)
+                ref_batch = features[begin2:end2]
+
+                distance_batch[
+                    0 : end1 - begin1, begin2:end2
+                ] = self.distance_block.pairwise_distances(feature_batch, ref_batch)
+
+            # From the minibatch of new feature vectors, determine if they are in the estimated manifold.
+            # If a feature vector is inside a hypersphere of some reference sample, then
+            # the new sample lies at the estimated manifold.
+            # The radii of the hyperspheres are determined from distances of neighborhood size k.
+            samples_in_manifold = distance_batch[0 : end1 - begin1, :, None] <= radii
+            batch_predictions[begin1:end1] = np.any(samples_in_manifold, axis=1).astype(np.int32)
+
+            max_realism_score[begin1:end1] = np.max(
+                radii[:, 0] / (distance_batch[0 : end1 - begin1, :] + self.eps), axis=1
+            )
+            nearest_indices[begin1:end1] = np.argmin(distance_batch[0 : end1 - begin1, :], axis=1)
+
+        return {
+            "fraction": float(np.mean(batch_predictions)),
+            "batch_predictions": batch_predictions,
+            "max_realisim_score": max_realism_score,
+            "nearest_indices": nearest_indices,
+        }
+
+    def evaluate_pr(
+        self,
+        features_1: np.ndarray,
+        radii_1: np.ndarray,
+        features_2: np.ndarray,
+        radii_2: np.ndarray,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Evaluate precision and recall efficiently.
+
+        :param features_1: [N1 x D] feature vectors for reference batch.
+        :param radii_1: [N1 x K1] radii for reference vectors.
+        :param features_2: [N2 x D] feature vectors for the other batch.
+        :param radii_2: [N x K2] radii for other vectors.
+        :return: a tuple of arrays for (precision, recall):
+                 - precision: an np.ndarray of length K1
+                 - recall: an np.ndarray of length K2
+        """
+        features_1_status = np.zeros([len(features_1), radii_2.shape[1]], dtype=np.bool_)
+        features_2_status = np.zeros([len(features_2), radii_1.shape[1]], dtype=np.bool_)
+        for begin_1 in range(0, len(features_1), self.row_batch_size):
+            end_1 = begin_1 + self.row_batch_size
+            batch_1 = features_1[begin_1:end_1]
+            for begin_2 in range(0, len(features_2), self.col_batch_size):
+                end_2 = begin_2 + self.col_batch_size
+                batch_2 = features_2[begin_2:end_2]
+                batch_1_in, batch_2_in = self.distance_block.less_thans(
+                    batch_1, radii_1[begin_1:end_1], batch_2, radii_2[begin_2:end_2]
+                )
+                features_1_status[begin_1:end_1] |= batch_1_in
+                features_2_status[begin_2:end_2] |= batch_2_in
+        return (
+            np.mean(features_2_status.astype(np.float64), axis=0),
+            np.mean(features_1_status.astype(np.float64), axis=0),
+        )
+
+
+class DistanceBlock:
+    """
+    Calculate pairwise distances between vectors.
+
+    Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L34
+    """
+
+    def __init__(self, session):
+        self.session = session
+
+        # Initialize TF graph to calculate pairwise distances.
+        with session.graph.as_default():
+            self._features_batch1 = tf.placeholder(tf.float32, shape=[None, None])
+            self._features_batch2 = tf.placeholder(tf.float32, shape=[None, None])
+            distance_block_16 = _batch_pairwise_distances(
+                tf.cast(self._features_batch1, tf.float16),
+                tf.cast(self._features_batch2, tf.float16),
+            )
+            self.distance_block = tf.cond(
+                tf.reduce_all(tf.math.is_finite(distance_block_16)),
+                lambda: tf.cast(distance_block_16, tf.float32),
+                lambda: _batch_pairwise_distances(self._features_batch1, self._features_batch2),
+            )
+
+            # Extra logic for less thans.
+            self._radii1 = tf.placeholder(tf.float32, shape=[None, None])
+            self._radii2 = tf.placeholder(tf.float32, shape=[None, None])
+            dist32 = tf.cast(self.distance_block, tf.float32)[..., None]
+            self._batch_1_in = tf.math.reduce_any(dist32 <= self._radii2, axis=1)
+            self._batch_2_in = tf.math.reduce_any(dist32 <= self._radii1[:, None], axis=0)
+
+    def pairwise_distances(self, U, V):
+        """
+        Evaluate pairwise distances between two batches of feature vectors.
+        """
+        return self.session.run(
+            self.distance_block,
+            feed_dict={self._features_batch1: U, self._features_batch2: V},
+        )
+
+    def less_thans(self, batch_1, radii_1, batch_2, radii_2):
+        return self.session.run(
+            [self._batch_1_in, self._batch_2_in],
+            feed_dict={
+                self._features_batch1: batch_1,
+                self._features_batch2: batch_2,
+                self._radii1: radii_1,
+                self._radii2: radii_2,
+            },
+        )
+
+
+def _batch_pairwise_distances(U, V):
+    """
+    Compute pairwise distances between two batches of feature vectors.
+    """
+    with tf.variable_scope("pairwise_dist_block"):
+        # Squared norms of each row in U and V.
+        norm_u = tf.reduce_sum(tf.square(U), 1)
+        norm_v = tf.reduce_sum(tf.square(V), 1)
+
+        # norm_u as a column and norm_v as a row vectors.
+        norm_u = tf.reshape(norm_u, [-1, 1])
+        norm_v = tf.reshape(norm_v, [1, -1])
+
+        # Pairwise squared Euclidean distances.
+        D = tf.maximum(norm_u - 2 * tf.matmul(U, V, False, True) + norm_v, 0.0)
+
+    return D
+
+
+class NpzArrayReader(ABC):
+    @abstractmethod
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        pass
+
+    @abstractmethod
+    def remaining(self) -> int:
+        pass
+
+    def read_batches(self, batch_size: int) -> Iterable[np.ndarray]:
+        def gen_fn():
+            while True:
+                batch = self.read_batch(batch_size)
+                if batch is None:
+                    break
+                yield batch
+
+        rem = self.remaining()
+        num_batches = rem // batch_size + int(rem % batch_size != 0)
+        return BatchIterator(gen_fn, num_batches)
+
+
+class BatchIterator:
+    def __init__(self, gen_fn, length):
+        self.gen_fn = gen_fn
+        self.length = length
+
+    def __len__(self):
+        return self.length
+
+    def __iter__(self):
+        return self.gen_fn()
+
+
+class StreamingNpzArrayReader(NpzArrayReader):
+    def __init__(self, arr_f, shape, dtype):
+        self.arr_f = arr_f
+        self.shape = shape
+        self.dtype = dtype
+        self.idx = 0
+
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        if self.idx >= self.shape[0]:
+            return None
+
+        bs = min(batch_size, self.shape[0] - self.idx)
+        self.idx += bs
+
+        if self.dtype.itemsize == 0:
+            return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype)
+
+        read_count = bs * np.prod(self.shape[1:])
+        read_size = int(read_count * self.dtype.itemsize)
+        data = _read_bytes(self.arr_f, read_size, "array data")
+        return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]])
+
+    def remaining(self) -> int:
+        return max(0, self.shape[0] - self.idx)
+
+
+class MemoryNpzArrayReader(NpzArrayReader):
+    def __init__(self, arr):
+        self.arr = arr
+        self.idx = 0
+
+    @classmethod
+    def load(cls, path: str, arr_name: str):
+        with open(path, "rb") as f:
+            arr = np.load(f)[arr_name]
+        return cls(arr)
+
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        if self.idx >= self.arr.shape[0]:
+            return None
+
+        res = self.arr[self.idx : self.idx + batch_size]
+        self.idx += batch_size
+        return res
+
+    def remaining(self) -> int:
+        return max(0, self.arr.shape[0] - self.idx)
+
+
+@contextmanager
+def open_npz_array(path: str, arr_name: str) -> NpzArrayReader:
+    with _open_npy_file(path, arr_name) as arr_f:
+        version = np.lib.format.read_magic(arr_f)
+        if version == (1, 0):
+            header = np.lib.format.read_array_header_1_0(arr_f)
+        elif version == (2, 0):
+            header = np.lib.format.read_array_header_2_0(arr_f)
+        else:
+            yield MemoryNpzArrayReader.load(path, arr_name)
+            return
+        shape, fortran, dtype = header
+        if fortran or dtype.hasobject:
+            yield MemoryNpzArrayReader.load(path, arr_name)
+        else:
+            yield StreamingNpzArrayReader(arr_f, shape, dtype)
+
+
+def _read_bytes(fp, size, error_template="ran out of data"):
+    """
+    Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886
+
+    Read from file-like object until size bytes are read.
+    Raises ValueError if not EOF is encountered before size bytes are read.
+    Non-blocking objects only supported if they derive from io objects.
+    Required as e.g. ZipExtFile in python 2.6 can return less data than
+    requested.
+    """
+    data = bytes()
+    while True:
+        # io files (default in python3) return None or raise on
+        # would-block, python2 file will truncate, probably nothing can be
+        # done about that.  note that regular files can't be non-blocking
+        try:
+            r = fp.read(size - len(data))
+            data += r
+            if len(r) == 0 or len(data) == size:
+                break
+        except io.BlockingIOError:
+            pass
+    if len(data) != size:
+        msg = "EOF: reading %s, expected %d bytes got %d"
+        raise ValueError(msg % (error_template, size, len(data)))
+    else:
+        return data
+
+
+@contextmanager
+def _open_npy_file(path: str, arr_name: str):
+    with open(path, "rb") as f:
+        with zipfile.ZipFile(f, "r") as zip_f:
+            if f"{arr_name}.npy" not in zip_f.namelist():
+                raise ValueError(f"missing {arr_name} in npz file")
+            with zip_f.open(f"{arr_name}.npy", "r") as arr_f:
+                yield arr_f
+
+
+def _download_inception_model():
+    if os.path.exists(INCEPTION_V3_PATH):
+        return
+    print("downloading InceptionV3 model...")
+    with requests.get(INCEPTION_V3_URL, stream=True) as r:
+        r.raise_for_status()
+        tmp_path = INCEPTION_V3_PATH + ".tmp"
+        with open(tmp_path, "wb") as f:
+            for chunk in tqdm(r.iter_content(chunk_size=8192)):
+                f.write(chunk)
+        os.rename(tmp_path, INCEPTION_V3_PATH)
+
+
+def _create_feature_graph(input_batch):
+    _download_inception_model()
+    prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}"
+    with open(INCEPTION_V3_PATH, "rb") as f:
+        graph_def = tf.GraphDef()
+        graph_def.ParseFromString(f.read())
+    pool3, spatial = tf.import_graph_def(
+        graph_def,
+        input_map={f"ExpandDims:0": input_batch},
+        return_elements=[FID_POOL_NAME, FID_SPATIAL_NAME],
+        name=prefix,
+    )
+    _update_shapes(pool3)
+    spatial = spatial[..., :7]
+    return pool3, spatial
+
+
+def _create_softmax_graph(input_batch):
+    _download_inception_model()
+    prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}"
+    with open(INCEPTION_V3_PATH, "rb") as f:
+        graph_def = tf.GraphDef()
+        graph_def.ParseFromString(f.read())
+    (matmul,) = tf.import_graph_def(
+        graph_def, return_elements=[f"softmax/logits/MatMul"], name=prefix
+    )
+    w = matmul.inputs[1]
+    logits = tf.matmul(input_batch, w)
+    return tf.nn.softmax(logits)
+
+
+def _update_shapes(pool3):
+    # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L50-L63
+    ops = pool3.graph.get_operations()
+    for op in ops:
+        for o in op.outputs:
+            shape = o.get_shape()
+            if shape._dims is not None:  # pylint: disable=protected-access
+                # shape = [s.value for s in shape] TF 1.x
+                shape = [s for s in shape]  # TF 2.x
+                new_shape = []
+                for j, s in enumerate(shape):
+                    if s == 1 and j == 0:
+                        new_shape.append(None)
+                    else:
+                        new_shape.append(s)
+                o.__dict__["_shape_val"] = tf.TensorShape(new_shape)
+    return pool3
+
+
+def _numpy_partition(arr, kth, **kwargs):
+    num_workers = min(cpu_count(), len(arr))
+    chunk_size = len(arr) // num_workers
+    extra = len(arr) % num_workers
+
+    start_idx = 0
+    batches = []
+    for i in range(num_workers):
+        size = chunk_size + (1 if i < extra else 0)
+        batches.append(arr[start_idx : start_idx + size])
+        start_idx += size
+
+    with ThreadPool(num_workers) as pool:
+        return list(pool.map(partial(np.partition, kth=kth, **kwargs), batches))
+
+
+if __name__ == "__main__":
+    main()

New/REG/evaluations/requirements.txt

@@ -0,0 +1,4 @@
+tensorflow-gpu>=2.0
+scipy
+requests
+tqdm

New/REG/generate_2400000_100.sh

@@ -0,0 +1,137 @@
+#!/usr/bin/env bash
+# ODE 采样（euler_sampler，与 samplers.py 中 REG ODE 一致）：分别用 5 / 10 / 50 步，输出到不同子目录。
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+cd "$SCRIPT_DIR"
+
+NUM_GPUS=1
+NUM_SAMPLES=100
+PER_PROC_BATCH_SIZE=10
+
+CFG_SCALE=1
+CLS_CFG_SCALE=1
+GH=0.85
+
+CKPT_PATH="/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/REG-XLarge-256/REG-XLarge-256/2400000.pt"
+#CKPT_PATH="/your_path/reg_xlarge_dinov2_base_align_8_cls/linear-dinov2-b-enc8/checkpoints/0040000.pt"
+
+# 根输出目录（可用环境变量覆盖）；其下会生成 ode_steps_5 / ode_steps_10 / ode_steps_50
+OUTPUT_BASE="${OUTPUT_BASE:-./reg_xl_2400000_ode_multistep}"
+NOISE_BANK_PATH="${NOISE_BANK_PATH:-${OUTPUT_BASE}/noise_bank.pt}"
+
+export NCCL_P2P_DISABLE=1
+export OUTPUT_BASE
+
+GLOBAL_BATCH_SIZE=$((PER_PROC_BATCH_SIZE * NUM_GPUS))
+TOTAL_SAMPLES=$((((NUM_SAMPLES + GLOBAL_BATCH_SIZE - 1) / GLOBAL_BATCH_SIZE) * GLOBAL_BATCH_SIZE))
+export NOISE_BANK_PATH TOTAL_SAMPLES
+
+REBUILD_NOISE_BANK=0
+if [[ ! -f "${NOISE_BANK_PATH}" ]]; then
+  REBUILD_NOISE_BANK=1
+else
+  EXISTING_SAMPLES=$(python - <<'PY'
+import os
+import torch
+p = os.environ["NOISE_BANK_PATH"]
+try:
+    d = torch.load(p, map_location="cpu", weights_only=True)
+except TypeError:
+    d = torch.load(p, map_location="cpu")
+print(int(d["z"].shape[0]))
+PY
+)
+  if [[ "${EXISTING_SAMPLES}" != "${TOTAL_SAMPLES}" ]]; then
+    echo "noise_bank size mismatch: existing=${EXISTING_SAMPLES}, required=${TOTAL_SAMPLES}. Rebuilding..."
+    REBUILD_NOISE_BANK=1
+  fi
+fi
+
+if [[ "${REBUILD_NOISE_BANK}" -eq 1 ]]; then
+  mkdir -p "$(dirname "${NOISE_BANK_PATH}")"
+  echo "Creating fixed noise bank: ${NOISE_BANK_PATH} (total_samples=${TOTAL_SAMPLES})"
+  python - <<'PY'
+import os
+import torch
+
+noise_path = os.environ["NOISE_BANK_PATH"]
+total_samples = int(os.environ["TOTAL_SAMPLES"])
+num_classes = 1000
+in_channels = 4
+latent_size = 32  # 256 / 8
+cls_dim = 768
+
+z = torch.randn(total_samples, in_channels, latent_size, latent_size)
+y = torch.randint(0, num_classes, (total_samples,), dtype=torch.long)
+cls_z = torch.randn(total_samples, cls_dim)
+
+torch.save({"z": z, "y": y, "cls_z": cls_z}, noise_path)
+print(f"Saved fixed noise bank to {noise_path}")
+PY
+fi
+
+for NUM_STEP in 5 10 20 50 100; do
+  RANDOM_PORT=$((RANDOM % 100 + 1200))
+  OUTPUT_DIR="${OUTPUT_BASE}/ode_steps_${NUM_STEP}/checkpoints"
+  echo "=== ODE num_steps=${NUM_STEP} -> ${OUTPUT_DIR} ==="
+
+  python -m torch.distributed.launch \
+    --master_port="${RANDOM_PORT}" \
+    --nproc_per_node="${NUM_GPUS}" \
+    generate.py \
+    --model "SiT-XL/2" \
+    --num-fid-samples "${NUM_SAMPLES}" \
+    --ckpt "${CKPT_PATH}" \
+    --path-type=linear \
+    --encoder-depth=8 \
+    --projector-embed-dims=768 \
+    --per-proc-batch-size="${PER_PROC_BATCH_SIZE}" \
+    --mode=ode \
+    --num-steps="${NUM_STEP}" \
+    --fixed-noise-file "${NOISE_BANK_PATH}" \
+    --cfg-scale="${CFG_SCALE}" \
+    --cls-cfg-scale="${CLS_CFG_SCALE}" \
+    --guidance-high="${GH}" \
+    --sample-dir "${OUTPUT_DIR}" \
+    --cls=768
+done
+
+echo "All ODE runs finished. Building side-by-side pairs..."
+
+python - <<'PY'
+import os
+from PIL import Image
+
+output_base = os.environ.get("OUTPUT_BASE", "./reg_xl_2400000_ode_multistep")
+steps = [5, 10, 50]
+
+def find_single_subdir(step):
+    root = os.path.join(output_base, f"ode_steps_{step}", "checkpoints")
+    if not os.path.isdir(root):
+        raise RuntimeError(f"Missing directory: {root}")
+    subdirs = [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
+    if len(subdirs) != 1:
+        raise RuntimeError(f"Expected exactly 1 run folder under {root}, got {len(subdirs)}: {subdirs}")
+    return os.path.join(root, subdirs[0])
+
+run_dirs = {s: find_single_subdir(s) for s in steps}
+pair_dir = os.path.join(output_base, "pair")
+os.makedirs(pair_dir, exist_ok=True)
+
+files = sorted([f for f in os.listdir(run_dirs[steps[0]]) if f.endswith(".png")])
+for name in files:
+    imgs = [Image.open(os.path.join(run_dirs[s], name)).convert("RGB") for s in steps]
+    w = sum(im.width for im in imgs)
+    h = max(im.height for im in imgs)
+    canvas = Image.new("RGB", (w, h))
+    x = 0
+    for im in imgs:
+        canvas.paste(im, (x, 0))
+        x += im.width
+    canvas.save(os.path.join(pair_dir, name))
+
+print(f"Saved paired comparisons to: {pair_dir}")
+PY
+
+echo "Done. Outputs under: ${OUTPUT_BASE}/ode_steps_{5,10,50}/checkpoints ; pairs under: ${OUTPUT_BASE}/pair"

New/REG/models/clip_vit.py

@@ -0,0 +1,426 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+import clip
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class UpdatedVisionTransformer(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x: torch.Tensor):
+        x = self.model.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.model.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.model.positional_embedding.to(x.dtype)
+        x = self.model.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.model.transformer(x)
+        x = x.permute(1, 0, 2)[:, 1:]  # LND -> NLD
+
+        # x = self.ln_post(x[:, 0, :])
+
+        # if self.proj is not None:
+        #     x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = UpdatedVisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

New/REG/models/jepa.py

@@ -0,0 +1,547 @@
+# 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.
+#
+
+import math
+from functools import partial
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def repeat_interleave_batch(x, B, repeat):
+    N = len(x) // B
+    x = torch.cat([
+        torch.cat([x[i*B:(i+1)*B] for _ in range(repeat)], dim=0)
+        for i in range(N)
+    ], dim=0)
+    return x
+
+def apply_masks(x, masks):
+    """
+    :param x: tensor of shape [B (batch-size), N (num-patches), D (feature-dim)]
+    :param masks: list of tensors containing indices of patches in [N] to keep
+    """
+    all_x = []
+    for m in masks:
+        mask_keep = m.unsqueeze(-1).repeat(1, 1, x.size(-1))
+        all_x += [torch.gather(x, dim=1, index=mask_keep)]
+    return torch.cat(all_x, dim=0)
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=float)
+    grid_w = np.arange(grid_size, dtype=float)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid length
+    return:
+    pos_embed: [grid_size, embed_dim] or [1+grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid = np.arange(grid_size, dtype=float)
+    pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega   # (D/2,)
+
+    pos = pos.reshape(-1)   # (M,)
+    out = np.einsum('m,d->md', pos, omega)   # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class ConvEmbed(nn.Module):
+    """
+    3x3 Convolution stems for ViT following ViTC models
+    """
+
+    def __init__(self, channels, strides, img_size=224, in_chans=3, batch_norm=True):
+        super().__init__()
+        # Build the stems
+        stem = []
+        channels = [in_chans] + channels
+        for i in range(len(channels) - 2):
+            stem += [nn.Conv2d(channels[i], channels[i+1], kernel_size=3,
+                               stride=strides[i], padding=1, bias=(not batch_norm))]
+            if batch_norm:
+                stem += [nn.BatchNorm2d(channels[i+1])]
+            stem += [nn.ReLU(inplace=True)]
+        stem += [nn.Conv2d(channels[-2], channels[-1], kernel_size=1, stride=strides[-1])]
+        self.stem = nn.Sequential(*stem)
+
+        # Comptute the number of patches
+        stride_prod = int(np.prod(strides))
+        self.num_patches = (img_size[0] // stride_prod)**2
+
+    def forward(self, x):
+        p = self.stem(x)
+        return p.flatten(2).transpose(1, 2)
+
+
+class VisionTransformerPredictor(nn.Module):
+    """ Vision Transformer """
+    def __init__(
+        self,
+        num_patches,
+        embed_dim=768,
+        predictor_embed_dim=384,
+        depth=6,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        init_std=0.02,
+        **kwargs
+    ):
+        super().__init__()
+        self.predictor_embed = nn.Linear(embed_dim, predictor_embed_dim, bias=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, predictor_embed_dim))
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        # --
+        self.predictor_pos_embed = nn.Parameter(torch.zeros(1, num_patches, predictor_embed_dim),
+                                                requires_grad=False)
+        predictor_pos_embed = get_2d_sincos_pos_embed(self.predictor_pos_embed.shape[-1],
+                                                      int(num_patches**.5),
+                                                      cls_token=False)
+        self.predictor_pos_embed.data.copy_(torch.from_numpy(predictor_pos_embed).float().unsqueeze(0))
+        # --
+        self.predictor_blocks = nn.ModuleList([
+            Block(
+                dim=predictor_embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.predictor_norm = norm_layer(predictor_embed_dim)
+        self.predictor_proj = nn.Linear(predictor_embed_dim, embed_dim, bias=True)
+        # ------
+        self.init_std = init_std
+        trunc_normal_(self.mask_token, std=self.init_std)
+        self.apply(self._init_weights)
+        self.fix_init_weight()
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.predictor_blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=self.init_std)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            trunc_normal_(m.weight, std=self.init_std)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x, masks_x, masks):
+        assert (masks is not None) and (masks_x is not None), 'Cannot run predictor without mask indices'
+
+        if not isinstance(masks_x, list):
+            masks_x = [masks_x]
+
+        if not isinstance(masks, list):
+            masks = [masks]
+
+        # -- Batch Size
+        B = len(x) // len(masks_x)
+
+        # -- map from encoder-dim to pedictor-dim
+        x = self.predictor_embed(x)
+
+        # -- add positional embedding to x tokens
+        x_pos_embed = self.predictor_pos_embed.repeat(B, 1, 1)
+        x += apply_masks(x_pos_embed, masks_x)
+
+        _, N_ctxt, D = x.shape
+
+        # -- concat mask tokens to x
+        pos_embs = self.predictor_pos_embed.repeat(B, 1, 1)
+        pos_embs = apply_masks(pos_embs, masks)
+        pos_embs = repeat_interleave_batch(pos_embs, B, repeat=len(masks_x))
+        # --
+        pred_tokens = self.mask_token.repeat(pos_embs.size(0), pos_embs.size(1), 1)
+        # --
+        pred_tokens += pos_embs
+        x = x.repeat(len(masks), 1, 1)
+        x = torch.cat([x, pred_tokens], dim=1)
+
+        # -- fwd prop
+        for blk in self.predictor_blocks:
+            x = blk(x)
+        x = self.predictor_norm(x)
+
+        # -- return preds for mask tokens
+        x = x[:, N_ctxt:]
+        x = self.predictor_proj(x)
+
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer """
+    def __init__(
+        self,
+        img_size=[224],
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        predictor_embed_dim=384,
+        depth=12,
+        predictor_depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        init_std=0.02,
+        **kwargs
+    ):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        # --
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0],
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+        # --
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim), requires_grad=False)
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1],
+                                            int(self.patch_embed.num_patches**.5),
+                                            cls_token=False)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        # --
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+        # ------
+        self.init_std = init_std
+        self.apply(self._init_weights)
+        self.fix_init_weight()
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=self.init_std)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            trunc_normal_(m.weight, std=self.init_std)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x, masks=None):
+        if masks is not None:
+            if not isinstance(masks, list):
+                masks = [masks]
+
+        # -- patchify x
+        x = self.patch_embed(x)
+        B, N, D = x.shape
+
+        # -- add positional embedding to x
+        pos_embed = self.interpolate_pos_encoding(x, self.pos_embed)
+        x = x + pos_embed
+
+        # -- mask x
+        if masks is not None:
+            x = apply_masks(x, masks)
+
+        # -- fwd prop
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        return x
+
+    def interpolate_pos_encoding(self, x, pos_embed):
+        npatch = x.shape[1] - 1
+        N = pos_embed.shape[1] - 1
+        if npatch == N:
+            return pos_embed
+        class_emb = pos_embed[:, 0]
+        pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        pos_embed = nn.functional.interpolate(
+            pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=math.sqrt(npatch / N),
+            mode='bicubic',
+        )
+        pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_emb.unsqueeze(0), pos_embed), dim=1)
+
+
+def vit_predictor(**kwargs):
+    model = VisionTransformerPredictor(
+        mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs)
+    return model
+
+
+def vit_tiny(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_small(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_large(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_huge(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_giant(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=1408, depth=40, num_heads=16, mlp_ratio=48/11,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+VIT_EMBED_DIMS = {
+    'vit_tiny': 192,
+    'vit_small': 384,
+    'vit_base': 768,
+    'vit_large': 1024,
+    'vit_huge': 1280,
+    'vit_giant': 1408,
+}

New/REG/models/mae_vit.py

@@ -0,0 +1,71 @@
+# 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.
+# --------------------------------------------------------
+# References:
+# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+import timm.models.vision_transformer
+
+
+class VisionTransformer(timm.models.vision_transformer.VisionTransformer):
+    """ Vision Transformer with support for global average pooling
+    """
+    def __init__(self, global_pool=False, **kwargs):
+        super(VisionTransformer, self).__init__(**kwargs)
+
+        self.global_pool = global_pool
+        if self.global_pool:
+            norm_layer = kwargs['norm_layer']
+            embed_dim = kwargs['embed_dim']
+            self.fc_norm = norm_layer(embed_dim)
+
+            del self.norm  # remove the original norm
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x = x[:, 1:, :] #.mean(dim=1)  # global pool without cls token
+
+        return x
+
+
+def vit_base_patch16(**kwargs):
+    model = VisionTransformer(
+        num_classes=0,
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_large_patch16(**kwargs):
+    model = VisionTransformer(
+        num_classes=0,
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_huge_patch14(**kwargs):
+    model = VisionTransformer(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model

New/REG/models/mocov3_vit.py

@@ -0,0 +1,207 @@
+# Copyright (c) Facebook, Inc. and its 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.
+
+import math
+import torch
+import torch.nn as nn
+from functools import partial, reduce
+from operator import mul
+
+from timm.layers.helpers import to_2tuple
+from timm.models.vision_transformer import VisionTransformer, _cfg
+from timm.models.vision_transformer import PatchEmbed
+
+__all__ = [
+    'vit_small', 
+    'vit_base',
+    'vit_large',
+    'vit_conv_small',
+    'vit_conv_base',
+]
+
+
+def patchify_avg(input_tensor, patch_size):
+    # Ensure input tensor is 4D: (batch_size, channels, height, width)
+    if input_tensor.dim() != 4:
+        raise ValueError("Input tensor must be 4D (batch_size, channels, height, width)")
+
+    # Get input tensor dimensions
+    batch_size, channels, height, width = input_tensor.shape
+
+    # Ensure patch_size is valid
+    patch_height, patch_width = patch_size, patch_size
+    if height % patch_height != 0 or width % patch_width != 0:
+        raise ValueError("Input tensor dimensions must be divisible by patch_size")
+
+    # Use unfold to create patches
+    patches = input_tensor.unfold(2, patch_height, patch_height).unfold(3, patch_width, patch_width)
+    
+    # Reshape patches to desired format: (batch_size, num_patches, channels)
+    patches = patches.contiguous().view(
+        batch_size, channels, -1, patch_height, patch_width
+        ).mean(dim=-1).mean(dim=-1)
+    patches = patches.permute(0, 2, 1).contiguous()
+
+    return patches
+
+
+
+class VisionTransformerMoCo(VisionTransformer):
+    def __init__(self, stop_grad_conv1=False, **kwargs):
+        super().__init__(**kwargs)
+        # Use fixed 2D sin-cos position embedding
+        self.build_2d_sincos_position_embedding()
+
+        # weight initialization
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Linear):
+                if 'qkv' in name:
+                    # treat the weights of Q, K, V separately
+                    val = math.sqrt(6. / float(m.weight.shape[0] // 3 + m.weight.shape[1]))
+                    nn.init.uniform_(m.weight, -val, val)
+                else:
+                    nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+        nn.init.normal_(self.cls_token, std=1e-6)
+
+        if isinstance(self.patch_embed, PatchEmbed):
+            # xavier_uniform initialization
+            val = math.sqrt(6. / float(3 * reduce(mul, self.patch_embed.patch_size, 1) + self.embed_dim))
+            nn.init.uniform_(self.patch_embed.proj.weight, -val, val)
+            nn.init.zeros_(self.patch_embed.proj.bias)
+
+            if stop_grad_conv1:
+                self.patch_embed.proj.weight.requires_grad = False
+                self.patch_embed.proj.bias.requires_grad = False
+
+    def build_2d_sincos_position_embedding(self, temperature=10000.):
+        h = self.patch_embed.img_size[0] // self.patch_embed.patch_size[0]
+        w = self.patch_embed.img_size[1] // self.patch_embed.patch_size[1]
+        grid_w = torch.arange(w, dtype=torch.float32)
+        grid_h = torch.arange(h, dtype=torch.float32)
+        grid_w, grid_h = torch.meshgrid(grid_w, grid_h)
+        assert self.embed_dim % 4 == 0, 'Embed dimension must be divisible by 4 for 2D sin-cos position embedding'
+        pos_dim = self.embed_dim // 4
+        omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
+        omega = 1. / (temperature**omega)
+        out_w = torch.einsum('m,d->md', [grid_w.flatten(), omega])
+        out_h = torch.einsum('m,d->md', [grid_h.flatten(), omega])
+        pos_emb = torch.cat([torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)], dim=1)[None, :, :]
+
+        # assert self.num_tokens == 1, 'Assuming one and only one token, [cls]'
+        pe_token = torch.zeros([1, 1, self.embed_dim], dtype=torch.float32)
+        self.pos_embed = nn.Parameter(torch.cat([pe_token, pos_emb], dim=1))
+        self.pos_embed.requires_grad = False
+
+    def forward_diffusion_output(self, x):
+        x = x.reshape(*x.shape[0:2], -1).permute(0, 2, 1)
+        x = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+        x = self.blocks(x)
+        x = self.norm(x)
+        return x
+
+class ConvStem(nn.Module):
+    """ 
+    ConvStem, from Early Convolutions Help Transformers See Better, Tete et al. https://arxiv.org/abs/2106.14881
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+
+        assert patch_size == 16, 'ConvStem only supports patch size of 16'
+        assert embed_dim % 8 == 0, 'Embed dimension must be divisible by 8 for ConvStem'
+
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        # build stem, similar to the design in https://arxiv.org/abs/2106.14881
+        stem = []
+        input_dim, output_dim = 3, embed_dim // 8
+        for l in range(4):
+            stem.append(nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=2, padding=1, bias=False))
+            stem.append(nn.BatchNorm2d(output_dim))
+            stem.append(nn.ReLU(inplace=True))
+            input_dim = output_dim
+            output_dim *= 2
+        stem.append(nn.Conv2d(input_dim, embed_dim, kernel_size=1))
+        self.proj = nn.Sequential(*stem)
+
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+def vit_small(**kwargs):
+    model = VisionTransformerMoCo(
+        img_size=256,
+        patch_size=16, embed_dim=384, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def vit_base(**kwargs):
+    model = VisionTransformerMoCo(
+        img_size=256,
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def vit_large(**kwargs):
+    model = VisionTransformerMoCo(
+        img_size=256,
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def vit_conv_small(**kwargs):
+    # minus one ViT block
+    model = VisionTransformerMoCo(
+        patch_size=16, embed_dim=384, depth=11, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), embed_layer=ConvStem, **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def vit_conv_base(**kwargs):
+    # minus one ViT block
+    model = VisionTransformerMoCo(
+        patch_size=16, embed_dim=768, depth=11, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), embed_layer=ConvStem, **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def build_mlp(num_layers, input_dim, mlp_dim, output_dim, last_bn=True):
+    mlp = []
+    for l in range(num_layers):
+        dim1 = input_dim if l == 0 else mlp_dim
+        dim2 = output_dim if l == num_layers - 1 else mlp_dim
+
+        mlp.append(nn.Linear(dim1, dim2, bias=False))
+
+        if l < num_layers - 1:
+            mlp.append(nn.BatchNorm1d(dim2))
+            mlp.append(nn.ReLU(inplace=True))
+        elif last_bn:
+            # follow SimCLR's design: https://github.com/google-research/simclr/blob/master/model_util.py#L157
+            # for simplicity, we further removed gamma in BN
+            mlp.append(nn.BatchNorm1d(dim2, affine=False))
+
+    return nn.Sequential(*mlp)

New/REG/models/sit.py

@@ -0,0 +1,420 @@
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# GLIDE: https://github.com/openai/glide-text2im
+# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
+
+
+def build_mlp(hidden_size, projector_dim, z_dim):
+    return nn.Sequential(
+                nn.Linear(hidden_size, projector_dim),
+                nn.SiLU(),
+                nn.Linear(projector_dim, projector_dim),
+                nn.SiLU(),
+                nn.Linear(projector_dim, z_dim),
+            )
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################            
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+    
+    @staticmethod
+    def positional_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        self.timestep_embedding = self.positional_embedding
+        t_freq = self.timestep_embedding(t, dim=self.frequency_embedding_size).to(t.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size, dropout_prob):
+        super().__init__()
+        use_cfg_embedding = dropout_prob > 0
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+    def token_drop(self, labels, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        labels = torch.where(drop_ids, self.num_classes, labels)
+        return labels
+
+    def forward(self, labels, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            labels = self.token_drop(labels, force_drop_ids)
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+
+#################################################################################
+#                                 Core SiT Model                                #
+#################################################################################
+
+class SiTBlock(nn.Module):
+    """
+    A SiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(
+            hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=block_kwargs["qk_norm"]
+            )
+        if "fused_attn" in block_kwargs.keys():
+            self.attn.fused_attn = block_kwargs["fused_attn"]
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0
+            )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaLN_modulation(c).chunk(6, dim=-1)
+        )
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of SiT.
+    """
+    def __init__(self, hidden_size, patch_size, out_channels, cls_token_dim):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.linear_cls = nn.Linear(hidden_size, cls_token_dim, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c, cls=None):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+
+        if cls is None:
+            x = self.linear(x)
+            return x, None
+        else:
+            cls_token = self.linear_cls(x[:, 0]).unsqueeze(1)
+            x = self.linear(x[:, 1:])
+            return x, cls_token.squeeze(1)
+
+
+class SiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        path_type='edm',
+        input_size=32,
+        patch_size=2,
+        in_channels=4,
+        hidden_size=1152,
+        decoder_hidden_size=768,
+        encoder_depth=8,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        class_dropout_prob=0.1,
+        num_classes=1000,
+        use_cfg=False,
+        z_dims=[768],
+        projector_dim=2048,
+        cls_token_dim=768,
+        **block_kwargs # fused_attn
+    ):
+        super().__init__()
+        self.path_type = path_type
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.use_cfg = use_cfg
+        self.num_classes = num_classes
+        self.z_dims = z_dims
+        self.encoder_depth = encoder_depth
+
+        self.x_embedder = PatchEmbed(
+            input_size, patch_size, in_channels, hidden_size, bias=True
+            )
+        self.t_embedder = TimestepEmbedder(hidden_size) # timestep embedding type
+        self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
+        num_patches = self.x_embedder.num_patches
+        # Will use fixed sin-cos embedding:
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches+1, hidden_size), requires_grad=False)
+
+        self.blocks = nn.ModuleList([
+            SiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **block_kwargs) for _ in range(depth)
+        ])
+        self.projectors = nn.ModuleList([
+            build_mlp(hidden_size, projector_dim, z_dim) for z_dim in z_dims
+            ])
+
+        z_dim = self.z_dims[0]
+        cls_token_dim = z_dim
+        self.final_layer = FinalLayer(decoder_hidden_size, patch_size, self.out_channels, cls_token_dim)
+
+
+        self.cls_projectors2 = nn.Linear(in_features=cls_token_dim, out_features=hidden_size, bias=True)
+        self.wg_norm = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(
+            self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5), cls_token=1, extra_tokens=1
+            )
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in SiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+        nn.init.constant_(self.final_layer.linear_cls.weight, 0)
+        nn.init.constant_(self.final_layer.linear_cls.bias, 0)
+
+    def unpatchify(self, x, patch_size=None):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, C, H, W)
+        """
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0] if patch_size is None else patch_size
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
+        return imgs
+    
+    def forward(self, x, t, y, return_logvar=False, cls_token=None):
+        """
+        Forward pass of SiT.
+        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        y: (N,) tensor of class labels
+        """
+
+        #cat with cls_token
+        x = self.x_embedder(x)   # (N, T, D), where T = H * W / patch_size ** 2
+        if cls_token is not None:
+            cls_token = self.cls_projectors2(cls_token)
+            cls_token = self.wg_norm(cls_token)
+            cls_token = cls_token.unsqueeze(1)  # [b, length, d]
+            x = torch.cat((cls_token, x), dim=1)
+            x = x + self.pos_embed
+        else:
+            exit()
+        N, T, D = x.shape
+
+        # timestep and class embedding
+        t_embed = self.t_embedder(t)                   # (N, D)
+        y = self.y_embedder(y, self.training)    # (N, D)
+        c = t_embed + y
+
+        for i, block in enumerate(self.blocks):
+            x = block(x, c)
+            if (i + 1) == self.encoder_depth:
+                zs = [projector(x.reshape(-1, D)).reshape(N, T, -1) for projector in self.projectors]
+
+        x, cls_token = self.final_layer(x, c, cls=cls_token)
+        x = self.unpatchify(x)
+
+        return x, zs, cls_token
+
+
+#################################################################################
+#                   Sine/Cosine Positional Embedding Functions                  #
+#################################################################################
+# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+#################################################################################
+#                                   SiT Configs                                  #
+#################################################################################
+
+def SiT_XL_2(**kwargs):
+    return SiT(depth=28, hidden_size=1152, decoder_hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_XL_4(**kwargs):
+    return SiT(depth=28, hidden_size=1152, decoder_hidden_size=1152, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_XL_8(**kwargs):
+    return SiT(depth=28, hidden_size=1152, decoder_hidden_size=1152, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_L_2(**kwargs):
+    return SiT(depth=24, hidden_size=1024, decoder_hidden_size=1024, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_L_4(**kwargs):
+    return SiT(depth=24, hidden_size=1024, decoder_hidden_size=1024, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_L_8(**kwargs):
+    return SiT(depth=24, hidden_size=1024, decoder_hidden_size=1024, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_B_2(**kwargs):
+    return SiT(depth=12, hidden_size=768, decoder_hidden_size=768, patch_size=2, num_heads=12, **kwargs)
+
+def SiT_B_4(**kwargs):
+    return SiT(depth=12, hidden_size=768, decoder_hidden_size=768, patch_size=4, num_heads=12, **kwargs)
+
+def SiT_B_8(**kwargs):
+    return SiT(depth=12, hidden_size=768, decoder_hidden_size=768, patch_size=8, num_heads=12, **kwargs)
+
+def SiT_S_2(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs)
+
+def SiT_S_4(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs)
+
+def SiT_S_8(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs)
+
+
+SiT_models = {
+    'SiT-XL/2': SiT_XL_2,  'SiT-XL/4': SiT_XL_4,  'SiT-XL/8': SiT_XL_8,
+    'SiT-L/2':  SiT_L_2,   'SiT-L/4':  SiT_L_4,   'SiT-L/8':  SiT_L_8,
+    'SiT-B/2':  SiT_B_2,   'SiT-B/4':  SiT_B_4,   'SiT-B/8':  SiT_B_8,
+    'SiT-S/2':  SiT_S_2,   'SiT-S/4':  SiT_S_4,   'SiT-S/8':  SiT_S_8,
+}
+

New/REG/preprocessing/README.md

@@ -0,0 +1,25 @@
+<h1 align="center"> Preprocessing Guide
+</h1>
+
+#### Dataset download
+
+We follow the preprocessing code used in [edm2](https://github.com/NVlabs/edm2). In this code we made a several edits: (1) we removed unncessary parts except preprocessing because this code is only used for preprocessing, (2) we use [-1, 1] range for an input to the stable diffusion VAE (similar to DiT or SiT) unlike edm2 that uses [0, 1] range, and (3) we consider preprocessing to 256x256 resolution (or 512x512 resolution).
+
+After downloading ImageNet, please run the following scripts (please update 256x256 to 512x512 if you want to do experiments on 512x512 resolution);
+
+Convert raw ImageNet data to a ZIP archive at 256x256 resolution
+```bash
+bash dataset_prepare_encode.sh
+```
+
+Convert the pixel data to VAE latents
+
+```bash
+bash dataset_prepare_convert.sh
+```
+
+Here,`YOUR_DOWNLOAD_PATH` is the directory that you downloaded the dataset, and `TARGET_PATH` is the directory that you will save the preprocessed images and corresponding compressed latent vectors. This directory will be used for your experiment scripts. 
+
+## Acknowledgement
+
+This code is mainly built upon [edm2](https://github.com/NVlabs/edm2) repository.

New/REG/preprocessing/dataset_image_encoder.py

@@ -0,0 +1,353 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Tool for creating ZIP/PNG based datasets."""
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+import functools
+import io
+import json
+import os
+import re
+import zipfile
+from pathlib import Path
+from typing import Callable, Optional, Tuple, Union
+import click
+import numpy as np
+import PIL.Image
+import torch
+from tqdm import tqdm
+
+from encoders import StabilityVAEEncoder
+from utils import load_encoders
+from torchvision.transforms import Normalize
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+CLIP_DEFAULT_MEAN = (0.48145466, 0.4578275, 0.40821073)
+CLIP_DEFAULT_STD = (0.26862954, 0.26130258, 0.27577711)
+
+def preprocess_raw_image(x, enc_type):
+    resolution = x.shape[-1]
+    if 'clip' in enc_type:
+        x = x / 255.
+        x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic')
+        x = Normalize(CLIP_DEFAULT_MEAN, CLIP_DEFAULT_STD)(x)
+    elif 'mocov3' in enc_type or 'mae' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+    elif 'dinov2' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+        x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic')
+    elif 'dinov1' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+    elif 'jepa' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+        x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic')
+
+    return x
+
+
+#----------------------------------------------------------------------------
+
+@dataclass
+class ImageEntry:
+    img: np.ndarray
+    label: Optional[int]
+
+#----------------------------------------------------------------------------
+# Parse a 'M,N' or 'MxN' integer tuple.
+# Example: '4x2' returns (4,2)
+
+def parse_tuple(s: str) -> Tuple[int, int]:
+    m = re.match(r'^(\d+)[x,](\d+)$', s)
+    if m:
+        return int(m.group(1)), int(m.group(2))
+    raise click.ClickException(f'cannot parse tuple {s}')
+
+#----------------------------------------------------------------------------
+
+def maybe_min(a: int, b: Optional[int]) -> int:
+    if b is not None:
+        return min(a, b)
+    return a
+
+#----------------------------------------------------------------------------
+
+def file_ext(name: Union[str, Path]) -> str:
+    return str(name).split('.')[-1]
+
+#----------------------------------------------------------------------------
+
+def is_image_ext(fname: Union[str, Path]) -> bool:
+    ext = file_ext(fname).lower()
+    return f'.{ext}' in PIL.Image.EXTENSION
+
+#----------------------------------------------------------------------------
+
+def open_image_folder(source_dir, *, max_images: Optional[int]) -> tuple[int, Iterator[ImageEntry]]:
+    input_images = []
+    def _recurse_dirs(root: str): # workaround Path().rglob() slowness
+        with os.scandir(root) as it:
+            for e in it:
+                if e.is_file():
+                    input_images.append(os.path.join(root, e.name))
+                elif e.is_dir():
+                    _recurse_dirs(os.path.join(root, e.name))
+    _recurse_dirs(source_dir)
+    input_images = sorted([f for f in input_images if is_image_ext(f)])
+
+    arch_fnames = {fname: os.path.relpath(fname, source_dir).replace('\\', '/') for fname in input_images}
+    max_idx = maybe_min(len(input_images), max_images)
+
+    # Load labels.
+    labels = dict()
+    meta_fname = os.path.join(source_dir, 'dataset.json')
+    if os.path.isfile(meta_fname):
+        with open(meta_fname, 'r') as file:
+            data = json.load(file)['labels']
+            if data is not None:
+                labels = {x[0]: x[1] for x in data}
+
+    # No labels available => determine from top-level directory names.
+    if len(labels) == 0:
+        toplevel_names = {arch_fname: arch_fname.split('/')[0] if '/' in arch_fname else '' for arch_fname in arch_fnames.values()}
+        toplevel_indices = {toplevel_name: idx for idx, toplevel_name in enumerate(sorted(set(toplevel_names.values())))}
+        if len(toplevel_indices) > 1:
+            labels = {arch_fname: toplevel_indices[toplevel_name] for arch_fname, toplevel_name in toplevel_names.items()}
+
+    def iterate_images():
+        for idx, fname in enumerate(input_images):
+            img = np.array(PIL.Image.open(fname).convert('RGB'))#.transpose(2, 0, 1)
+            yield ImageEntry(img=img, label=labels.get(arch_fnames[fname]))
+            if idx >= max_idx - 1:
+                break
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def open_image_zip(source, *, max_images: Optional[int]) -> tuple[int, Iterator[ImageEntry]]:
+    with zipfile.ZipFile(source, mode='r') as z:
+        input_images = [str(f) for f in sorted(z.namelist()) if is_image_ext(f)]
+        max_idx = maybe_min(len(input_images), max_images)
+
+        # Load labels.
+        labels = dict()
+        if 'dataset.json' in z.namelist():
+            with z.open('dataset.json', 'r') as file:
+                data = json.load(file)['labels']
+                if data is not None:
+                    labels = {x[0]: x[1] for x in data}
+
+    def iterate_images():
+        with zipfile.ZipFile(source, mode='r') as z:
+            for idx, fname in enumerate(input_images):
+                with z.open(fname, 'r') as file:
+                    img = np.array(PIL.Image.open(file).convert('RGB'))
+                yield ImageEntry(img=img, label=labels.get(fname))
+                if idx >= max_idx - 1:
+                    break
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def make_transform(
+    transform: Optional[str],
+    output_width: Optional[int],
+    output_height: Optional[int]
+) -> Callable[[np.ndarray], Optional[np.ndarray]]:
+    def scale(width, height, img):
+        w = img.shape[1]
+        h = img.shape[0]
+        if width == w and height == h:
+            return img
+        img = PIL.Image.fromarray(img, 'RGB')
+        ww = width if width is not None else w
+        hh = height if height is not None else h
+        img = img.resize((ww, hh), PIL.Image.Resampling.LANCZOS)
+        return np.array(img)
+
+    def center_crop(width, height, img):
+        crop = np.min(img.shape[:2])
+        img = img[(img.shape[0] - crop) // 2 : (img.shape[0] + crop) // 2, (img.shape[1] - crop) // 2 : (img.shape[1] + crop) // 2]
+        img = PIL.Image.fromarray(img, 'RGB')
+        img = img.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        return np.array(img)
+
+    def center_crop_wide(width, height, img):
+        ch = int(np.round(width * img.shape[0] / img.shape[1]))
+        if img.shape[1] < width or ch < height:
+            return None
+
+        img = img[(img.shape[0] - ch) // 2 : (img.shape[0] + ch) // 2]
+        img = PIL.Image.fromarray(img, 'RGB')
+        img = img.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        img = np.array(img)
+
+        canvas = np.zeros([width, width, 3], dtype=np.uint8)
+        canvas[(width - height) // 2 : (width + height) // 2, :] = img
+        return canvas
+
+    def center_crop_imagenet(image_size: int, arr: np.ndarray):
+        """
+        Center cropping implementation from ADM.
+        https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+        """
+        pil_image = PIL.Image.fromarray(arr)
+        while min(*pil_image.size) >= 2 * image_size:
+            new_size = tuple(x // 2 for x in pil_image.size)
+            assert len(new_size) == 2
+            pil_image = pil_image.resize(new_size, resample=PIL.Image.Resampling.BOX)
+
+        scale = image_size / min(*pil_image.size)
+        new_size = tuple(round(x * scale) for x in pil_image.size)
+        assert len(new_size) == 2
+        pil_image = pil_image.resize(new_size, resample=PIL.Image.Resampling.BICUBIC)
+
+        arr = np.array(pil_image)
+        crop_y = (arr.shape[0] - image_size) // 2
+        crop_x = (arr.shape[1] - image_size) // 2
+        return arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]
+
+    if transform is None:
+        return functools.partial(scale, output_width, output_height)
+    if transform == 'center-crop':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + 'transform')
+        return functools.partial(center_crop, output_width, output_height)
+    if transform == 'center-crop-wide':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + ' transform')
+        return functools.partial(center_crop_wide, output_width, output_height)
+    if transform == 'center-crop-dhariwal':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + ' transform')
+        if output_width != output_height:
+            raise click.ClickException('width and height must match in --resolution=WxH when using ' + transform + ' transform')
+        return functools.partial(center_crop_imagenet, output_width)
+    assert False, 'unknown transform'
+
+#----------------------------------------------------------------------------
+
+def open_dataset(source, *, max_images: Optional[int]):
+    if os.path.isdir(source):
+        return open_image_folder(source, max_images=max_images)
+    elif os.path.isfile(source):
+        if file_ext(source) == 'zip':
+            return open_image_zip(source, max_images=max_images)
+        else:
+            raise click.ClickException(f'Only zip archives are supported: {source}')
+    else:
+        raise click.ClickException(f'Missing input file or directory: {source}')
+
+#----------------------------------------------------------------------------
+
+def open_dest(dest: str) -> Tuple[str, Callable[[str, Union[bytes, str]], None], Callable[[], None]]:
+    dest_ext = file_ext(dest)
+
+    if dest_ext == 'zip':
+        if os.path.dirname(dest) != '':
+            os.makedirs(os.path.dirname(dest), exist_ok=True)
+        zf = zipfile.ZipFile(file=dest, mode='w', compression=zipfile.ZIP_STORED)
+        def zip_write_bytes(fname: str, data: Union[bytes, str]):
+            zf.writestr(fname, data)
+        return '', zip_write_bytes, zf.close
+    else:
+        # If the output folder already exists, check that is is
+        # empty.
+        #
+        # Note: creating the output directory is not strictly
+        # necessary as folder_write_bytes() also mkdirs, but it's better
+        # to give an error message earlier in case the dest folder
+        # somehow cannot be created.
+        if os.path.isdir(dest) and len(os.listdir(dest)) != 0:
+            raise click.ClickException('--dest folder must be empty')
+        os.makedirs(dest, exist_ok=True)
+
+        def folder_write_bytes(fname: str, data: Union[bytes, str]):
+            os.makedirs(os.path.dirname(fname), exist_ok=True)
+            with open(fname, 'wb') as fout:
+                if isinstance(data, str):
+                    data = data.encode('utf8')
+                fout.write(data)
+        return dest, folder_write_bytes, lambda: None
+
+#----------------------------------------------------------------------------
+
+@click.group()
+def cmdline():
+    '''Dataset processing tool for dataset image data conversion and VAE encode/decode preprocessing.'''
+    if os.environ.get('WORLD_SIZE', '1') != '1':
+        raise click.ClickException('Distributed execution is not supported.')
+
+
+#----------------------------------------------------------------------------
+
+
+
+@cmdline.command()
+@click.option('--source',     help='Input directory or archive name', metavar='PATH',   type=str, required=True)
+@click.option('--dest',       help='Output directory or archive name', metavar='PATH',  type=str, required=True)
+@click.option('--max-images', help='Maximum number of images to output', metavar='INT', type=int)
+@click.option('--enc-type',   help='Maximum number of images to output', metavar='PATH', type=str, default='dinov2-vit-b')
+@click.option('--resolution',   help='Maximum number of images to output', metavar='INT', type=int, default=256)
+
+def encode(
+    source: str,
+    dest: str,
+    max_images: Optional[int],
+    enc_type,
+    resolution
+):
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    encoder, encoder_type, architectures = load_encoders(enc_type, device, resolution)
+    encoder, encoder_type, architectures = encoder[0], encoder_type[0], architectures[0]
+    print("Encoder is over!!!")
+
+    """Encode pixel data to VAE latents."""
+    PIL.Image.init()
+    if dest == '':
+        raise click.ClickException('--dest output filename or directory must not be an empty string')
+
+    num_files, input_iter = open_dataset(source, max_images=max_images)
+    archive_root_dir, save_bytes, close_dest = open_dest(dest)
+    print("Data is over!!!")
+    labels = []
+
+    temp_list1 = []
+    temp_list2 = []
+    for idx, image in tqdm(enumerate(input_iter), total=num_files):
+        with torch.no_grad():
+            img_tensor = torch.tensor(image.img).to('cuda').permute(2, 0, 1).unsqueeze(0)
+            raw_image_ = preprocess_raw_image(img_tensor, encoder_type)
+            z = encoder.forward_features(raw_image_)
+            if 'dinov2' in encoder_type: z = z['x_norm_patchtokens']
+            temp_list1.append(z)
+            z = z.detach().cpu().numpy()
+            temp_list2.append(z)
+
+        idx_str = f'{idx:08d}'
+        archive_fname = f'{idx_str[:5]}/img-feature-{idx_str}.npy'
+
+        f = io.BytesIO()
+        np.save(f, z)
+        save_bytes(os.path.join(archive_root_dir, archive_fname), f.getvalue())
+        labels.append([archive_fname, image.label] if image.label is not None else None)
+
+
+    metadata = {'labels': labels if all(x is not None for x in labels) else None}
+    save_bytes(os.path.join(archive_root_dir, 'dataset.json'), json.dumps(metadata))
+    close_dest()
+
+if __name__ == "__main__":
+    cmdline()
+
+
+#----------------------------------------------------------------------------

New/REG/preprocessing/dataset_prepare_convert.sh

@@ -0,0 +1,11 @@
+
+
+
+
+
+#256
+python preprocessing/dataset_tools.py convert \
+    --source=/home/share/imagenet/train \
+    --dest=/home/share/imagenet_vae/imagenet_256_vae \
+    --resolution=256x256 \
+    --transform=center-crop-dhariwal

New/REG/preprocessing/dataset_prepare_encode.sh

@@ -0,0 +1,9 @@
+
+
+
+
+
+#256
+python preprocessing/dataset_tools.py encode \
+    --source=/home/share/imagenet_vae/imagenet_256_vae \
+    --dest=/home/share/imagenet_vae/vae-sd-256

New/REG/preprocessing/dataset_tools.py

@@ -0,0 +1,422 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Tool for creating ZIP/PNG based datasets."""
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+import functools
+import io
+import json
+import os
+import re
+import zipfile
+from pathlib import Path
+from typing import Callable, Optional, Tuple, Union
+import click
+import numpy as np
+import PIL.Image
+import torch
+from tqdm import tqdm
+
+from encoders import StabilityVAEEncoder
+
+#----------------------------------------------------------------------------
+
+@dataclass
+class ImageEntry:
+    img: np.ndarray
+    label: Optional[int]
+
+#----------------------------------------------------------------------------
+# Parse a 'M,N' or 'MxN' integer tuple.
+# Example: '4x2' returns (4,2)
+
+def parse_tuple(s: str) -> Tuple[int, int]:
+    m = re.match(r'^(\d+)[x,](\d+)$', s)
+    if m:
+        return int(m.group(1)), int(m.group(2))
+    raise click.ClickException(f'cannot parse tuple {s}')
+
+#----------------------------------------------------------------------------
+
+def maybe_min(a: int, b: Optional[int]) -> int:
+    if b is not None:
+        return min(a, b)
+    return a
+
+#----------------------------------------------------------------------------
+
+def file_ext(name: Union[str, Path]) -> str:
+    return str(name).split('.')[-1]
+
+#----------------------------------------------------------------------------
+
+def is_image_ext(fname: Union[str, Path]) -> bool:
+    ext = file_ext(fname).lower()
+    return f'.{ext}' in PIL.Image.EXTENSION
+
+#----------------------------------------------------------------------------
+
+def open_image_folder(source_dir, *, max_images: Optional[int]) -> tuple[int, Iterator[ImageEntry]]:
+    input_images = []
+    def _recurse_dirs(root: str): # workaround Path().rglob() slowness
+        with os.scandir(root) as it:
+            for e in it:
+                if e.is_file():
+                    input_images.append(os.path.join(root, e.name))
+                elif e.is_dir():
+                    _recurse_dirs(os.path.join(root, e.name))
+    _recurse_dirs(source_dir)
+    input_images = sorted([f for f in input_images if is_image_ext(f)])
+
+    arch_fnames = {fname: os.path.relpath(fname, source_dir).replace('\\', '/') for fname in input_images}
+    max_idx = maybe_min(len(input_images), max_images)
+
+    # Load labels.
+    labels = dict()
+    meta_fname = os.path.join(source_dir, 'dataset.json')
+    if os.path.isfile(meta_fname):
+        with open(meta_fname, 'r') as file:
+            data = json.load(file)['labels']
+            if data is not None:
+                labels = {x[0]: x[1] for x in data}
+
+    # No labels available => determine from top-level directory names.
+    if len(labels) == 0:
+        toplevel_names = {arch_fname: arch_fname.split('/')[0] if '/' in arch_fname else '' for arch_fname in arch_fnames.values()}
+        toplevel_indices = {toplevel_name: idx for idx, toplevel_name in enumerate(sorted(set(toplevel_names.values())))}
+        if len(toplevel_indices) > 1:
+            labels = {arch_fname: toplevel_indices[toplevel_name] for arch_fname, toplevel_name in toplevel_names.items()}
+
+    def iterate_images():
+        for idx, fname in enumerate(input_images):
+            img = np.array(PIL.Image.open(fname).convert('RGB'))
+            yield ImageEntry(img=img, label=labels.get(arch_fnames[fname]))
+            if idx >= max_idx - 1:
+                break
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def open_image_zip(source, *, max_images: Optional[int]) -> tuple[int, Iterator[ImageEntry]]:
+    with zipfile.ZipFile(source, mode='r') as z:
+        input_images = [str(f) for f in sorted(z.namelist()) if is_image_ext(f)]
+        max_idx = maybe_min(len(input_images), max_images)
+
+        # Load labels.
+        labels = dict()
+        if 'dataset.json' in z.namelist():
+            with z.open('dataset.json', 'r') as file:
+                data = json.load(file)['labels']
+                if data is not None:
+                    labels = {x[0]: x[1] for x in data}
+
+    def iterate_images():
+        with zipfile.ZipFile(source, mode='r') as z:
+            for idx, fname in enumerate(input_images):
+                with z.open(fname, 'r') as file:
+                    img = np.array(PIL.Image.open(file).convert('RGB'))
+                yield ImageEntry(img=img, label=labels.get(fname))
+                if idx >= max_idx - 1:
+                    break
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def make_transform(
+    transform: Optional[str],
+    output_width: Optional[int],
+    output_height: Optional[int]
+) -> Callable[[np.ndarray], Optional[np.ndarray]]:
+    def scale(width, height, img):
+        w = img.shape[1]
+        h = img.shape[0]
+        if width == w and height == h:
+            return img
+        img = PIL.Image.fromarray(img, 'RGB')
+        ww = width if width is not None else w
+        hh = height if height is not None else h
+        img = img.resize((ww, hh), PIL.Image.Resampling.LANCZOS)
+        return np.array(img)
+
+    def center_crop(width, height, img):
+        crop = np.min(img.shape[:2])
+        img = img[(img.shape[0] - crop) // 2 : (img.shape[0] + crop) // 2, (img.shape[1] - crop) // 2 : (img.shape[1] + crop) // 2]
+        img = PIL.Image.fromarray(img, 'RGB')
+        img = img.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        return np.array(img)
+
+    def center_crop_wide(width, height, img):
+        ch = int(np.round(width * img.shape[0] / img.shape[1]))
+        if img.shape[1] < width or ch < height:
+            return None
+
+        img = img[(img.shape[0] - ch) // 2 : (img.shape[0] + ch) // 2]
+        img = PIL.Image.fromarray(img, 'RGB')
+        img = img.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        img = np.array(img)
+
+        canvas = np.zeros([width, width, 3], dtype=np.uint8)
+        canvas[(width - height) // 2 : (width + height) // 2, :] = img
+        return canvas
+
+    def center_crop_imagenet(image_size: int, arr: np.ndarray):
+        """
+        Center cropping implementation from ADM.
+        https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+        """
+        pil_image = PIL.Image.fromarray(arr)
+        while min(*pil_image.size) >= 2 * image_size:
+            new_size = tuple(x // 2 for x in pil_image.size)
+            assert len(new_size) == 2
+            pil_image = pil_image.resize(new_size, resample=PIL.Image.Resampling.BOX)
+
+        scale = image_size / min(*pil_image.size)
+        new_size = tuple(round(x * scale) for x in pil_image.size)
+        assert len(new_size) == 2
+        pil_image = pil_image.resize(new_size, resample=PIL.Image.Resampling.BICUBIC)
+
+        arr = np.array(pil_image)
+        crop_y = (arr.shape[0] - image_size) // 2
+        crop_x = (arr.shape[1] - image_size) // 2
+        return arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]
+
+    if transform is None:
+        return functools.partial(scale, output_width, output_height)
+    if transform == 'center-crop':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + 'transform')
+        return functools.partial(center_crop, output_width, output_height)
+    if transform == 'center-crop-wide':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + ' transform')
+        return functools.partial(center_crop_wide, output_width, output_height)
+    if transform == 'center-crop-dhariwal':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + ' transform')
+        if output_width != output_height:
+            raise click.ClickException('width and height must match in --resolution=WxH when using ' + transform + ' transform')
+        return functools.partial(center_crop_imagenet, output_width)
+    assert False, 'unknown transform'
+
+#----------------------------------------------------------------------------
+
+def open_dataset(source, *, max_images: Optional[int]):
+    if os.path.isdir(source):
+        return open_image_folder(source, max_images=max_images)
+    elif os.path.isfile(source):
+        if file_ext(source) == 'zip':
+            return open_image_zip(source, max_images=max_images)
+        else:
+            raise click.ClickException(f'Only zip archives are supported: {source}')
+    else:
+        raise click.ClickException(f'Missing input file or directory: {source}')
+
+#----------------------------------------------------------------------------
+
+def open_dest(dest: str) -> Tuple[str, Callable[[str, Union[bytes, str]], None], Callable[[], None]]:
+    dest_ext = file_ext(dest)
+
+    if dest_ext == 'zip':
+        if os.path.dirname(dest) != '':
+            os.makedirs(os.path.dirname(dest), exist_ok=True)
+        zf = zipfile.ZipFile(file=dest, mode='w', compression=zipfile.ZIP_STORED)
+        def zip_write_bytes(fname: str, data: Union[bytes, str]):
+            zf.writestr(fname, data)
+        return '', zip_write_bytes, zf.close
+    else:
+        # If the output folder already exists, check that is is
+        # empty.
+        #
+        # Note: creating the output directory is not strictly
+        # necessary as folder_write_bytes() also mkdirs, but it's better
+        # to give an error message earlier in case the dest folder
+        # somehow cannot be created.
+        if os.path.isdir(dest) and len(os.listdir(dest)) != 0:
+            raise click.ClickException('--dest folder must be empty')
+        os.makedirs(dest, exist_ok=True)
+
+        def folder_write_bytes(fname: str, data: Union[bytes, str]):
+            os.makedirs(os.path.dirname(fname), exist_ok=True)
+            with open(fname, 'wb') as fout:
+                if isinstance(data, str):
+                    data = data.encode('utf8')
+                fout.write(data)
+        return dest, folder_write_bytes, lambda: None
+
+#----------------------------------------------------------------------------
+
+@click.group()
+def cmdline():
+    '''Dataset processing tool for dataset image data conversion and VAE encode/decode preprocessing.'''
+    if os.environ.get('WORLD_SIZE', '1') != '1':
+        raise click.ClickException('Distributed execution is not supported.')
+
+#----------------------------------------------------------------------------
+
+@cmdline.command()
+@click.option('--source',     help='Input directory or archive name', metavar='PATH',   type=str, required=True)
+@click.option('--dest',       help='Output directory or archive name', metavar='PATH',  type=str, required=True)
+@click.option('--max-images', help='Maximum number of images to output', metavar='INT', type=int)
+@click.option('--transform',  help='Input crop/resize mode', metavar='MODE',            type=click.Choice(['center-crop', 'center-crop-wide', 'center-crop-dhariwal']))
+@click.option('--resolution', help='Output resolution (e.g., 512x512)', metavar='WxH',  type=parse_tuple)
+
+def convert(
+    source: str,
+    dest: str,
+    max_images: Optional[int],
+    transform: Optional[str],
+    resolution: Optional[Tuple[int, int]]
+):
+    """Convert an image dataset into archive format for training.
+
+    Specifying the input images:
+
+    \b
+    --source path/                      Recursively load all images from path/
+    --source dataset.zip                Load all images from dataset.zip
+
+    Specifying the output format and path:
+
+    \b
+    --dest /path/to/dir                 Save output files under /path/to/dir
+    --dest /path/to/dataset.zip         Save output files into /path/to/dataset.zip
+
+    The output dataset format can be either an image folder or an uncompressed zip archive.
+    Zip archives makes it easier to move datasets around file servers and clusters, and may
+    offer better training performance on network file systems.
+
+    Images within the dataset archive will be stored as uncompressed PNG.
+    Uncompresed PNGs can be efficiently decoded in the training loop.
+
+    Class labels are stored in a file called 'dataset.json' that is stored at the
+    dataset root folder.  This file has the following structure:
+
+    \b
+    {
+        "labels": [
+            ["00000/img00000000.png",6],
+            ["00000/img00000001.png",9],
+            ... repeated for every image in the datase
+            ["00049/img00049999.png",1]
+        ]
+    }
+
+    If the 'dataset.json' file cannot be found, class labels are determined from
+    top-level directory names.
+
+    Image scale/crop and resolution requirements:
+
+    Output images must be square-shaped and they must all have the same power-of-two
+    dimensions.
+
+    To scale arbitrary input image size to a specific width and height, use the
+    --resolution option.  Output resolution will be either the original
+    input resolution (if resolution was not specified) or the one specified with
+    --resolution option.
+
+    The --transform=center-crop-dhariwal selects a crop/rescale mode that is intended
+    to exactly match with results obtained for ImageNet in common diffusion model literature:
+
+    \b
+    python dataset_tool.py convert --source=downloads/imagenet/ILSVRC/Data/CLS-LOC/train \\
+        --dest=datasets/img64.zip --resolution=64x64 --transform=center-crop-dhariwal
+    """
+    PIL.Image.init()
+    if dest == '':
+        raise click.ClickException('--dest output filename or directory must not be an empty string')
+    print("Begin!!!!!!!!")
+    num_files, input_iter = open_dataset(source, max_images=max_images)
+    print("open_dataset is over")
+    archive_root_dir, save_bytes, close_dest = open_dest(dest)
+    print("open_dest is over")
+    transform_image = make_transform(transform, *resolution if resolution is not None else (None, None))
+    dataset_attrs = None
+
+    labels = []
+    for idx, image in tqdm(enumerate(input_iter), total=num_files):
+        idx_str = f'{idx:08d}'
+        archive_fname = f'{idx_str[:5]}/img{idx_str}.png'
+
+        # Apply crop and resize.
+        img = transform_image(image.img)
+        if img is None:
+            continue
+
+        # Error check to require uniform image attributes across
+        # the whole dataset.
+        assert img.ndim == 3
+        cur_image_attrs = {'width': img.shape[1], 'height': img.shape[0]}
+        if dataset_attrs is None:
+            dataset_attrs = cur_image_attrs
+            width = dataset_attrs['width']
+            height = dataset_attrs['height']
+            if width != height:
+                raise click.ClickException(f'Image dimensions after scale and crop are required to be square.  Got {width}x{height}')
+            if width != 2 ** int(np.floor(np.log2(width))):
+                raise click.ClickException('Image width/height after scale and crop are required to be power-of-two')
+        elif dataset_attrs != cur_image_attrs:
+            err = [f'  dataset {k}/cur image {k}: {dataset_attrs[k]}/{cur_image_attrs[k]}' for k in dataset_attrs.keys()]
+            raise click.ClickException(f'Image {archive_fname} attributes must be equal across all images of the dataset.  Got:\n' + '\n'.join(err))
+
+        # Save the image as an uncompressed PNG.
+        img = PIL.Image.fromarray(img)
+        image_bits = io.BytesIO()
+        img.save(image_bits, format='png', compress_level=0, optimize=False)
+        save_bytes(os.path.join(archive_root_dir, archive_fname), image_bits.getbuffer())
+        labels.append([archive_fname, image.label] if image.label is not None else None)
+
+    metadata = {'labels': labels if all(x is not None for x in labels) else None}
+    save_bytes(os.path.join(archive_root_dir, 'dataset.json'), json.dumps(metadata))
+    close_dest()
+
+#----------------------------------------------------------------------------
+
+@cmdline.command()
+@click.option('--model-url',  help='VAE encoder model', metavar='URL',                  type=str, default='stabilityai/sd-vae-ft-mse', show_default=True)
+@click.option('--source',     help='Input directory or archive name', metavar='PATH',   type=str, required=True)
+@click.option('--dest',       help='Output directory or archive name', metavar='PATH',  type=str, required=True)
+@click.option('--max-images', help='Maximum number of images to output', metavar='INT', type=int)
+
+def encode(
+    model_url: str,
+    source: str,
+    dest: str,
+    max_images: Optional[int],
+):
+    """Encode pixel data to VAE latents."""
+    PIL.Image.init()
+    if dest == '':
+        raise click.ClickException('--dest output filename or directory must not be an empty string')
+
+    vae = StabilityVAEEncoder(vae_name=model_url, batch_size=1)
+    print("VAE is over!!!")
+    num_files, input_iter = open_dataset(source, max_images=max_images)
+    archive_root_dir, save_bytes, close_dest = open_dest(dest)
+    print("Data is over!!!")
+    labels = []
+    #device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    for idx, image in tqdm(enumerate(input_iter), total=num_files):
+        img_tensor = torch.tensor(image.img).to('cuda').permute(2, 0, 1).unsqueeze(0)
+        mean_std = vae.encode_pixels(img_tensor)[0].cpu()
+        idx_str = f'{idx:08d}'
+        archive_fname = f'{idx_str[:5]}/img-mean-std-{idx_str}.npy'
+
+        f = io.BytesIO()
+        np.save(f, mean_std)
+        save_bytes(os.path.join(archive_root_dir, archive_fname), f.getvalue())
+        labels.append([archive_fname, image.label] if image.label is not None else None)
+
+    metadata = {'labels': labels if all(x is not None for x in labels) else None}
+    save_bytes(os.path.join(archive_root_dir, 'dataset.json'), json.dumps(metadata))
+    close_dest()
+
+if __name__ == "__main__":
+    cmdline()
+
+#----------------------------------------------------------------------------

New/REG/preprocessing/dnnlib/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+from .util import EasyDict, make_cache_dir_path

New/REG/preprocessing/dnnlib/util.py

@@ -0,0 +1,485 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Miscellaneous utility classes and functions."""
+
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+import requests
+import html
+import hashlib
+import glob
+import tempfile
+import urllib
+import urllib.parse
+import uuid
+
+from typing import Any, Callable, BinaryIO, List, Tuple, Union, Optional
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: Optional[str] = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: Union[str, bytes]) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if isinstance(text, bytes):
+            text = text.decode()
+        if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+            self.file = None
+
+
+# Cache directories
+# ------------------------------------------------------------------------------------------
+
+_dnnlib_cache_dir = None
+
+def set_cache_dir(path: str) -> None:
+    global _dnnlib_cache_dir
+    _dnnlib_cache_dir = path
+
+def make_cache_dir_path(*paths: str) -> str:
+    if _dnnlib_cache_dir is not None:
+        return os.path.join(_dnnlib_cache_dir, *paths)
+    if 'DNNLIB_CACHE_DIR' in os.environ:
+        return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
+    if 'HOME' in os.environ:
+        return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
+    if 'USERPROFILE' in os.environ:
+        return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
+    return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def format_time_brief(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m".format(s // (60 * 60), (s // 60) % 60)
+    else:
+        return "{0}d {1:02}h".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24)
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name) # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: Union[str, Callable], **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name) if isinstance(func_name, str) else func_name
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def construct_class_by_name(*args, class_name: Union[str, type], **kwargs) -> Any:
+    """Finds the python class with the given name and constructs it with the given arguments."""
+    return call_func_by_name(*args, func_name=class_name, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    module = obj.__module__
+    if module == '__main__':
+        fname = sys.modules[module].__file__
+        assert fname is not None
+        module = os.path.splitext(os.path.basename(fname))[0]
+    return module + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: Optional[List[str]] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        os.makedirs(target_dir_name, exist_ok=True)
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    if allow_file_urls and obj.startswith('file://'):
+        return True
+    try:
+        res = urllib.parse.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = urllib.parse.urlparse(urllib.parse.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+# Note on static typing: a better API would be to split 'open_url' to 'openl_url' and
+# 'download_url' with separate return types (BinaryIO, str).  As the `return_filename=True`
+# case is somewhat uncommon, we just pretend like this function never returns a string
+# and type ignore return value for those cases.
+def open_url(url: str, cache_dir: Optional[str] = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> BinaryIO:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+    assert not (return_filename and (not cache))
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match('^[a-z]+://', url):
+        return url if return_filename else open(url, "rb") # type: ignore
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname()
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith('file://'):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r'^/[a-zA-Z]:', filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb") # type: ignore
+
+    assert is_url(url)
+
+    # Lookup from cache.
+    if cache_dir is None:
+        cache_dir = make_cache_dir_path('downloads')
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            filename = cache_files[0]
+            return filename if return_filename else open(filename, "rb") # type: ignore
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = urllib.parse.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive download quota exceeded -- please try again later")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    assert url_data is not None
+
+    # Save to cache.
+    if cache:
+        assert url_name is not None
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        safe_name = safe_name[:min(len(safe_name), 128)]
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+        if return_filename:
+            return cache_file # type: ignore
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)

New/REG/preprocessing/encoders.py

@@ -0,0 +1,103 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Converting between pixel and latent representations of image data."""
+
+import os
+import warnings
+import numpy as np
+import torch
+from torch_utils import persistence
+from torch_utils import misc
+
+warnings.filterwarnings('ignore', 'torch.utils._pytree._register_pytree_node is deprecated.')
+warnings.filterwarnings('ignore', '`resume_download` is deprecated')
+
+#----------------------------------------------------------------------------
+# Abstract base class for encoders/decoders that convert back and forth
+# between pixel and latent representations of image data.
+#
+# Logically, "raw pixels" are first encoded into "raw latents" that are
+# then further encoded into "final latents". Decoding, on the other hand,
+# goes directly from the final latents to raw pixels. The final latents are
+# used as inputs and outputs of the model, whereas the raw latents are
+# stored in the dataset. This separation provides added flexibility in terms
+# of performing just-in-time adjustments, such as data whitening, without
+# having to construct a new dataset.
+#
+# All image data is represented as PyTorch tensors in NCHW order.
+# Raw pixels are represented as 3-channel uint8.
+
+@persistence.persistent_class
+class Encoder:
+    def __init__(self):
+        pass
+
+    def init(self, device): # force lazy init to happen now
+        pass
+
+    def __getstate__(self):
+        return self.__dict__
+
+    def encode_pixels(self, x): # raw pixels => raw latents
+        raise NotImplementedError # to be overridden by subclass
+#----------------------------------------------------------------------------
+# Pre-trained VAE encoder from Stability AI.
+
+@persistence.persistent_class
+class StabilityVAEEncoder(Encoder):
+    def __init__(self,
+        vae_name    = 'stabilityai/sd-vae-ft-mse',  # Name of the VAE to use.
+        batch_size  = 8,                            # Batch size to use when running the VAE.
+    ):
+        super().__init__()
+        self.vae_name = vae_name
+        self.batch_size = int(batch_size)
+        self._vae = None
+
+    def init(self, device): # force lazy init to happen now
+        super().init(device)
+        if self._vae is None:
+            self._vae = load_stability_vae(self.vae_name, device=device)
+        else:
+            self._vae.to(device)
+
+    def __getstate__(self):
+        return dict(super().__getstate__(), _vae=None) # do not pickle the vae
+
+    def _run_vae_encoder(self, x):
+        d = self._vae.encode(x)['latent_dist']
+        return torch.cat([d.mean, d.std], dim=1)
+
+    def encode_pixels(self, x): # raw pixels => raw latents
+        self.init(x.device)
+        x = x.to(torch.float32) / 127.5 - 1
+        x = torch.cat([self._run_vae_encoder(batch) for batch in x.split(self.batch_size)])
+        return x
+
+#----------------------------------------------------------------------------
+
+def load_stability_vae(vae_name='stabilityai/sd-vae-ft-mse', device=torch.device('cpu')):
+    import dnnlib
+    cache_dir = dnnlib.make_cache_dir_path('diffusers')
+    os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'
+    os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
+    os.environ['HF_HOME'] = cache_dir
+
+
+    import diffusers # pip install diffusers # pyright: ignore [reportMissingImports]
+    try:
+        # First try with local_files_only to avoid consulting tfhub metadata if the model is already in cache.
+        vae = diffusers.models.AutoencoderKL.from_pretrained(
+            vae_name, cache_dir=cache_dir, local_files_only=True
+            )
+    except:
+        # Could not load the model from cache; try without local_files_only.
+        vae = diffusers.models.AutoencoderKL.from_pretrained(vae_name, cache_dir=cache_dir)
+    return vae.eval().requires_grad_(False).to(device)
+
+#----------------------------------------------------------------------------

New/REG/preprocessing/torch_utils/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+# empty

New/REG/preprocessing/torch_utils/distributed.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+import os
+import re
+import socket
+import torch
+import torch.distributed
+from . import training_stats
+
+_sync_device = None
+
+#----------------------------------------------------------------------------
+
+def init():
+    global _sync_device
+
+    if not torch.distributed.is_initialized():
+        # Setup some reasonable defaults for env-based distributed init if
+        # not set by the running environment.
+        if 'MASTER_ADDR' not in os.environ:
+            os.environ['MASTER_ADDR'] = 'localhost'
+        if 'MASTER_PORT' not in os.environ:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.bind(('', 0))
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            os.environ['MASTER_PORT'] = str(s.getsockname()[1])
+            s.close()
+        if 'RANK' not in os.environ:
+            os.environ['RANK'] = '0'
+        if 'LOCAL_RANK' not in os.environ:
+            os.environ['LOCAL_RANK'] = '0'
+        if 'WORLD_SIZE' not in os.environ:
+            os.environ['WORLD_SIZE'] = '1'
+        backend = 'gloo' if os.name == 'nt' else 'nccl'
+        torch.distributed.init_process_group(backend=backend, init_method='env://')
+        torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', '0')))
+
+    _sync_device = torch.device('cuda') if get_world_size() > 1 else None
+    training_stats.init_multiprocessing(rank=get_rank(), sync_device=_sync_device)
+
+#----------------------------------------------------------------------------
+
+def get_rank():
+    return torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+
+#----------------------------------------------------------------------------
+
+def get_world_size():
+    return torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1
+
+#----------------------------------------------------------------------------
+
+def should_stop():
+    return False
+
+#----------------------------------------------------------------------------
+
+def should_suspend():
+    return False
+
+#----------------------------------------------------------------------------
+
+def request_suspend():
+    pass
+
+#----------------------------------------------------------------------------
+
+def update_progress(cur, total):
+    pass
+
+#----------------------------------------------------------------------------
+
+def print0(*args, **kwargs):
+    if get_rank() == 0:
+        print(*args, **kwargs)
+
+#----------------------------------------------------------------------------
+
+class CheckpointIO:
+    def __init__(self, **kwargs):
+        self._state_objs = kwargs
+
+    def save(self, pt_path, verbose=True):
+        if verbose:
+            print0(f'Saving {pt_path} ... ', end='', flush=True)
+        data = dict()
+        for name, obj in self._state_objs.items():
+            if obj is None:
+                data[name] = None
+            elif isinstance(obj, dict):
+                data[name] = obj
+            elif hasattr(obj, 'state_dict'):
+                data[name] = obj.state_dict()
+            elif hasattr(obj, '__getstate__'):
+                data[name] = obj.__getstate__()
+            elif hasattr(obj, '__dict__'):
+                data[name] = obj.__dict__
+            else:
+                raise ValueError(f'Invalid state object of type {type(obj).__name__}')
+        if get_rank() == 0:
+            torch.save(data, pt_path)
+        if verbose:
+            print0('done')
+
+    def load(self, pt_path, verbose=True):
+        if verbose:
+            print0(f'Loading {pt_path} ... ', end='', flush=True)
+        data = torch.load(pt_path, map_location=torch.device('cpu'))
+        for name, obj in self._state_objs.items():
+            if obj is None:
+                pass
+            elif isinstance(obj, dict):
+                obj.clear()
+                obj.update(data[name])
+            elif hasattr(obj, 'load_state_dict'):
+                obj.load_state_dict(data[name])
+            elif hasattr(obj, '__setstate__'):
+                obj.__setstate__(data[name])
+            elif hasattr(obj, '__dict__'):
+                obj.__dict__.clear()
+                obj.__dict__.update(data[name])
+            else:
+                raise ValueError(f'Invalid state object of type {type(obj).__name__}')
+        if verbose:
+            print0('done')
+
+    def load_latest(self, run_dir, pattern=r'training-state-(\d+).pt', verbose=True):
+        fnames = [entry.name for entry in os.scandir(run_dir) if entry.is_file() and re.fullmatch(pattern, entry.name)]
+        if len(fnames) == 0:
+            return None
+        pt_path = os.path.join(run_dir, max(fnames, key=lambda x: float(re.fullmatch(pattern, x).group(1))))
+        self.load(pt_path, verbose=verbose)
+        return pt_path
+
+#----------------------------------------------------------------------------

New/REG/preprocessing/torch_utils/misc.py

@@ -0,0 +1,277 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+import re
+import contextlib
+import functools
+import numpy as np
+import torch
+import warnings
+import dnnlib
+
+#----------------------------------------------------------------------------
+# Re-seed torch & numpy random generators based on the given arguments.
+
+def set_random_seed(*args):
+    seed = hash(args) % (1 << 31)
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+#----------------------------------------------------------------------------
+# Cached construction of constant tensors. Avoids CPU=>GPU copy when the
+# same constant is used multiple times.
+
+_constant_cache = dict()
+
+def constant(value, shape=None, dtype=None, device=None, memory_format=None):
+    value = np.asarray(value)
+    if shape is not None:
+        shape = tuple(shape)
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if device is None:
+        device = torch.device('cpu')
+    if memory_format is None:
+        memory_format = torch.contiguous_format
+
+    key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format)
+    tensor = _constant_cache.get(key, None)
+    if tensor is None:
+        tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device)
+        if shape is not None:
+            tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape))
+        tensor = tensor.contiguous(memory_format=memory_format)
+        _constant_cache[key] = tensor
+    return tensor
+
+#----------------------------------------------------------------------------
+# Variant of constant() that inherits dtype and device from the given
+# reference tensor by default.
+
+def const_like(ref, value, shape=None, dtype=None, device=None, memory_format=None):
+    if dtype is None:
+        dtype = ref.dtype
+    if device is None:
+        device = ref.device
+    return constant(value, shape=shape, dtype=dtype, device=device, memory_format=memory_format)
+
+#----------------------------------------------------------------------------
+# Cached construction of temporary tensors in pinned CPU memory.
+
+@functools.lru_cache(None)
+def pinned_buf(shape, dtype):
+    return torch.empty(shape, dtype=dtype).pin_memory()
+
+#----------------------------------------------------------------------------
+# Symbolic assert.
+
+try:
+    symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access
+except AttributeError:
+    symbolic_assert = torch.Assert # 1.7.0
+
+#----------------------------------------------------------------------------
+# Context manager to temporarily suppress known warnings in torch.jit.trace().
+# Note: Cannot use catch_warnings because of https://bugs.python.org/issue29672
+
+@contextlib.contextmanager
+def suppress_tracer_warnings():
+    flt = ('ignore', None, torch.jit.TracerWarning, None, 0)
+    warnings.filters.insert(0, flt)
+    yield
+    warnings.filters.remove(flt)
+
+#----------------------------------------------------------------------------
+# Assert that the shape of a tensor matches the given list of integers.
+# None indicates that the size of a dimension is allowed to vary.
+# Performs symbolic assertion when used in torch.jit.trace().
+
+def assert_shape(tensor, ref_shape):
+    if tensor.ndim != len(ref_shape):
+        raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}')
+    for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)):
+        if ref_size is None:
+            pass
+        elif isinstance(ref_size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}')
+        elif isinstance(size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}')
+        elif size != ref_size:
+            raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')
+
+#----------------------------------------------------------------------------
+# Function decorator that calls torch.autograd.profiler.record_function().
+
+def profiled_function(fn):
+    def decorator(*args, **kwargs):
+        with torch.autograd.profiler.record_function(fn.__name__):
+            return fn(*args, **kwargs)
+    decorator.__name__ = fn.__name__
+    return decorator
+
+#----------------------------------------------------------------------------
+# Sampler for torch.utils.data.DataLoader that loops over the dataset
+# indefinitely, shuffling items as it goes.
+
+class InfiniteSampler(torch.utils.data.Sampler):
+    def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed=0, start_idx=0):
+        assert len(dataset) > 0
+        assert num_replicas > 0
+        assert 0 <= rank < num_replicas
+        warnings.filterwarnings('ignore', '`data_source` argument is not used and will be removed')
+        super().__init__(dataset)
+        self.dataset_size = len(dataset)
+        self.start_idx = start_idx + rank
+        self.stride = num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+
+    def __iter__(self):
+        idx = self.start_idx
+        epoch = None
+        while True:
+            if epoch != idx // self.dataset_size:
+                epoch = idx // self.dataset_size
+                order = np.arange(self.dataset_size)
+                if self.shuffle:
+                    np.random.RandomState(hash((self.seed, epoch)) % (1 << 31)).shuffle(order)
+            yield int(order[idx % self.dataset_size])
+            idx += self.stride
+
+#----------------------------------------------------------------------------
+# Utilities for operating with torch.nn.Module parameters and buffers.
+
+def params_and_buffers(module):
+    assert isinstance(module, torch.nn.Module)
+    return list(module.parameters()) + list(module.buffers())
+
+def named_params_and_buffers(module):
+    assert isinstance(module, torch.nn.Module)
+    return list(module.named_parameters()) + list(module.named_buffers())
+
+@torch.no_grad()
+def copy_params_and_buffers(src_module, dst_module, require_all=False):
+    assert isinstance(src_module, torch.nn.Module)
+    assert isinstance(dst_module, torch.nn.Module)
+    src_tensors = dict(named_params_and_buffers(src_module))
+    for name, tensor in named_params_and_buffers(dst_module):
+        assert (name in src_tensors) or (not require_all)
+        if name in src_tensors:
+            tensor.copy_(src_tensors[name])
+
+#----------------------------------------------------------------------------
+# Context manager for easily enabling/disabling DistributedDataParallel
+# synchronization.
+
+@contextlib.contextmanager
+def ddp_sync(module, sync):
+    assert isinstance(module, torch.nn.Module)
+    if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel):
+        yield
+    else:
+        with module.no_sync():
+            yield
+
+#----------------------------------------------------------------------------
+# Check DistributedDataParallel consistency across processes.
+
+def check_ddp_consistency(module, ignore_regex=None):
+    assert isinstance(module, torch.nn.Module)
+    for name, tensor in named_params_and_buffers(module):
+        fullname = type(module).__name__ + '.' + name
+        if ignore_regex is not None and re.fullmatch(ignore_regex, fullname):
+            continue
+        tensor = tensor.detach()
+        if tensor.is_floating_point():
+            tensor = torch.nan_to_num(tensor)
+        other = tensor.clone()
+        torch.distributed.broadcast(tensor=other, src=0)
+        assert (tensor == other).all(), fullname
+
+#----------------------------------------------------------------------------
+# Print summary table of module hierarchy.
+
+@torch.no_grad()
+def print_module_summary(module, inputs, max_nesting=3, skip_redundant=True):
+    assert isinstance(module, torch.nn.Module)
+    assert not isinstance(module, torch.jit.ScriptModule)
+    assert isinstance(inputs, (tuple, list))
+
+    # Register hooks.
+    entries = []
+    nesting = [0]
+    def pre_hook(_mod, _inputs):
+        nesting[0] += 1
+    def post_hook(mod, _inputs, outputs):
+        nesting[0] -= 1
+        if nesting[0] <= max_nesting:
+            outputs = list(outputs) if isinstance(outputs, (tuple, list)) else [outputs]
+            outputs = [t for t in outputs if isinstance(t, torch.Tensor)]
+            entries.append(dnnlib.EasyDict(mod=mod, outputs=outputs))
+    hooks = [mod.register_forward_pre_hook(pre_hook) for mod in module.modules()]
+    hooks += [mod.register_forward_hook(post_hook) for mod in module.modules()]
+
+    # Run module.
+    outputs = module(*inputs)
+    for hook in hooks:
+        hook.remove()
+
+    # Identify unique outputs, parameters, and buffers.
+    tensors_seen = set()
+    for e in entries:
+        e.unique_params = [t for t in e.mod.parameters() if id(t) not in tensors_seen]
+        e.unique_buffers = [t for t in e.mod.buffers() if id(t) not in tensors_seen]
+        e.unique_outputs = [t for t in e.outputs if id(t) not in tensors_seen]
+        tensors_seen |= {id(t) for t in e.unique_params + e.unique_buffers + e.unique_outputs}
+
+    # Filter out redundant entries.
+    if skip_redundant:
+        entries = [e for e in entries if len(e.unique_params) or len(e.unique_buffers) or len(e.unique_outputs)]
+
+    # Construct table.
+    rows = [[type(module).__name__, 'Parameters', 'Buffers', 'Output shape', 'Datatype']]
+    rows += [['---'] * len(rows[0])]
+    param_total = 0
+    buffer_total = 0
+    submodule_names = {mod: name for name, mod in module.named_modules()}
+    for e in entries:
+        name = '<top-level>' if e.mod is module else submodule_names[e.mod]
+        param_size = sum(t.numel() for t in e.unique_params)
+        buffer_size = sum(t.numel() for t in e.unique_buffers)
+        output_shapes = [str(list(t.shape)) for t in e.outputs]
+        output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs]
+        rows += [[
+            name + (':0' if len(e.outputs) >= 2 else ''),
+            str(param_size) if param_size else '-',
+            str(buffer_size) if buffer_size else '-',
+            (output_shapes + ['-'])[0],
+            (output_dtypes + ['-'])[0],
+        ]]
+        for idx in range(1, len(e.outputs)):
+            rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]]
+        param_total += param_size
+        buffer_total += buffer_size
+    rows += [['---'] * len(rows[0])]
+    rows += [['Total', str(param_total), str(buffer_total), '-', '-']]
+
+    # Print table.
+    widths = [max(len(cell) for cell in column) for column in zip(*rows)]
+    print()
+    for row in rows:
+        print('  '.join(cell + ' ' * (width - len(cell)) for cell, width in zip(row, widths)))
+    print()
+
+#----------------------------------------------------------------------------
+# Tile a batch of images into a 2D grid.
+
+def tile_images(x, w, h):
+    assert x.ndim == 4 # NCHW => CHW
+    return x.reshape(h, w, *x.shape[1:]).permute(2, 0, 3, 1, 4).reshape(x.shape[1], h * x.shape[2], w * x.shape[3])
+
+#----------------------------------------------------------------------------

New/REG/preprocessing/torch_utils/persistence.py

@@ -0,0 +1,257 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Facilities for pickling Python code alongside other data.
+
+The pickled code is automatically imported into a separate Python module
+during unpickling. This way, any previously exported pickles will remain
+usable even if the original code is no longer available, or if the current
+version of the code is not consistent with what was originally pickled."""
+
+import sys
+import pickle
+import io
+import inspect
+import copy
+import uuid
+import types
+import functools
+import dnnlib
+
+#----------------------------------------------------------------------------
+
+_version            = 6         # internal version number
+_decorators         = set()     # {decorator_class, ...}
+_import_hooks       = []        # [hook_function, ...]
+_module_to_src_dict = dict()    # {module: src, ...}
+_src_to_module_dict = dict()    # {src: module, ...}
+
+#----------------------------------------------------------------------------
+
+def persistent_class(orig_class):
+    r"""Class decorator that extends a given class to save its source code
+    when pickled.
+
+    Example:
+
+        from torch_utils import persistence
+
+        @persistence.persistent_class
+        class MyNetwork(torch.nn.Module):
+            def __init__(self, num_inputs, num_outputs):
+                super().__init__()
+                self.fc = MyLayer(num_inputs, num_outputs)
+                ...
+
+        @persistence.persistent_class
+        class MyLayer(torch.nn.Module):
+            ...
+
+    When pickled, any instance of `MyNetwork` and `MyLayer` will save its
+    source code alongside other internal state (e.g., parameters, buffers,
+    and submodules). This way, any previously exported pickle will remain
+    usable even if the class definitions have been modified or are no
+    longer available.
+
+    The decorator saves the source code of the entire Python module
+    containing the decorated class. It does *not* save the source code of
+    any imported modules. Thus, the imported modules must be available
+    during unpickling, also including `torch_utils.persistence` itself.
+
+    It is ok to call functions defined in the same module from the
+    decorated class. However, if the decorated class depends on other
+    classes defined in the same module, they must be decorated as well.
+    This is illustrated in the above example in the case of `MyLayer`.
+
+    It is also possible to employ the decorator just-in-time before
+    calling the constructor. For example:
+
+        cls = MyLayer
+        if want_to_make_it_persistent:
+            cls = persistence.persistent_class(cls)
+        layer = cls(num_inputs, num_outputs)
+
+    As an additional feature, the decorator also keeps track of the
+    arguments that were used to construct each instance of the decorated
+    class. The arguments can be queried via `obj.init_args` and
+    `obj.init_kwargs`, and they are automatically pickled alongside other
+    object state. This feature can be disabled on a per-instance basis
+    by setting `self._record_init_args = False` in the constructor.
+
+    A typical use case is to first unpickle a previous instance of a
+    persistent class, and then upgrade it to use the latest version of
+    the source code:
+
+        with open('old_pickle.pkl', 'rb') as f:
+            old_net = pickle.load(f)
+        new_net = MyNetwork(*old_obj.init_args, **old_obj.init_kwargs)
+        misc.copy_params_and_buffers(old_net, new_net, require_all=True)
+    """
+    assert isinstance(orig_class, type)
+    if is_persistent(orig_class):
+        return orig_class
+
+    assert orig_class.__module__ in sys.modules
+    orig_module = sys.modules[orig_class.__module__]
+    orig_module_src = _module_to_src(orig_module)
+
+    @functools.wraps(orig_class, updated=())
+    class Decorator(orig_class):
+        _orig_module_src = orig_module_src
+        _orig_class_name = orig_class.__name__
+
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            record_init_args = getattr(self, '_record_init_args', True)
+            self._init_args = copy.deepcopy(args) if record_init_args else None
+            self._init_kwargs = copy.deepcopy(kwargs) if record_init_args else None
+            assert orig_class.__name__ in orig_module.__dict__
+            _check_pickleable(self.__reduce__())
+
+        @property
+        def init_args(self):
+            assert self._init_args is not None
+            return copy.deepcopy(self._init_args)
+
+        @property
+        def init_kwargs(self):
+            assert self._init_kwargs is not None
+            return dnnlib.EasyDict(copy.deepcopy(self._init_kwargs))
+
+        def __reduce__(self):
+            fields = list(super().__reduce__())
+            fields += [None] * max(3 - len(fields), 0)
+            if fields[0] is not _reconstruct_persistent_obj:
+                meta = dict(type='class', version=_version, module_src=self._orig_module_src, class_name=self._orig_class_name, state=fields[2])
+                fields[0] = _reconstruct_persistent_obj # reconstruct func
+                fields[1] = (meta,) # reconstruct args
+                fields[2] = None # state dict
+            return tuple(fields)
+
+    _decorators.add(Decorator)
+    return Decorator
+
+#----------------------------------------------------------------------------
+
+def is_persistent(obj):
+    r"""Test whether the given object or class is persistent, i.e.,
+    whether it will save its source code when pickled.
+    """
+    try:
+        if obj in _decorators:
+            return True
+    except TypeError:
+        pass
+    return type(obj) in _decorators # pylint: disable=unidiomatic-typecheck
+
+#----------------------------------------------------------------------------
+
+def import_hook(hook):
+    r"""Register an import hook that is called whenever a persistent object
+    is being unpickled. A typical use case is to patch the pickled source
+    code to avoid errors and inconsistencies when the API of some imported
+    module has changed.
+
+    The hook should have the following signature:
+
+        hook(meta) -> modified meta
+
+    `meta` is an instance of `dnnlib.EasyDict` with the following fields:
+
+        type:       Type of the persistent object, e.g. `'class'`.
+        version:    Internal version number of `torch_utils.persistence`.
+        module_src  Original source code of the Python module.
+        class_name: Class name in the original Python module.
+        state:      Internal state of the object.
+
+    Example:
+
+        @persistence.import_hook
+        def wreck_my_network(meta):
+            if meta.class_name == 'MyNetwork':
+                print('MyNetwork is being imported. I will wreck it!')
+                meta.module_src = meta.module_src.replace("True", "False")
+            return meta
+    """
+    assert callable(hook)
+    _import_hooks.append(hook)
+
+#----------------------------------------------------------------------------
+
+def _reconstruct_persistent_obj(meta):
+    r"""Hook that is called internally by the `pickle` module to unpickle
+    a persistent object.
+    """
+    meta = dnnlib.EasyDict(meta)
+    meta.state = dnnlib.EasyDict(meta.state)
+    for hook in _import_hooks:
+        meta = hook(meta)
+        assert meta is not None
+
+    assert meta.version == _version
+    module = _src_to_module(meta.module_src)
+
+    assert meta.type == 'class'
+    orig_class = module.__dict__[meta.class_name]
+    decorator_class = persistent_class(orig_class)
+    obj = decorator_class.__new__(decorator_class)
+
+    setstate = getattr(obj, '__setstate__', None)
+    if callable(setstate):
+        setstate(meta.state) # pylint: disable=not-callable
+    else:
+        obj.__dict__.update(meta.state)
+    return obj
+
+#----------------------------------------------------------------------------
+
+def _module_to_src(module):
+    r"""Query the source code of a given Python module.
+    """
+    src = _module_to_src_dict.get(module, None)
+    if src is None:
+        src = inspect.getsource(module)
+        _module_to_src_dict[module] = src
+        _src_to_module_dict[src] = module
+    return src
+
+def _src_to_module(src):
+    r"""Get or create a Python module for the given source code.
+    """
+    module = _src_to_module_dict.get(src, None)
+    if module is None:
+        module_name = "_imported_module_" + uuid.uuid4().hex
+        module = types.ModuleType(module_name)
+        sys.modules[module_name] = module
+        _module_to_src_dict[module] = src
+        _src_to_module_dict[src] = module
+        exec(src, module.__dict__) # pylint: disable=exec-used
+    return module
+
+#----------------------------------------------------------------------------
+
+def _check_pickleable(obj):
+    r"""Check that the given object is pickleable, raising an exception if
+    it is not. This function is expected to be considerably more efficient
+    than actually pickling the object.
+    """
+    def recurse(obj):
+        if isinstance(obj, (list, tuple, set)):
+            return [recurse(x) for x in obj]
+        if isinstance(obj, dict):
+            return [[recurse(x), recurse(y)] for x, y in obj.items()]
+        if isinstance(obj, (str, int, float, bool, bytes, bytearray)):
+            return None # Python primitive types are pickleable.
+        if f'{type(obj).__module__}.{type(obj).__name__}' in ['numpy.ndarray', 'torch.Tensor', 'torch.nn.parameter.Parameter']:
+            return None # NumPy arrays and PyTorch tensors are pickleable.
+        if is_persistent(obj):
+            return None # Persistent objects are pickleable, by virtue of the constructor check.
+        return obj
+    with io.BytesIO() as f:
+        pickle.dump(recurse(obj), f)
+
+#----------------------------------------------------------------------------

New/REG/preprocessing/torch_utils/training_stats.py

@@ -0,0 +1,283 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Facilities for reporting and collecting training statistics across
+multiple processes and devices. The interface is designed to minimize
+synchronization overhead as well as the amount of boilerplate in user
+code."""
+
+import re
+import numpy as np
+import torch
+import dnnlib
+
+from . import misc
+
+#----------------------------------------------------------------------------
+
+_num_moments    = 3             # [num_scalars, sum_of_scalars, sum_of_squares]
+_reduce_dtype   = torch.float32 # Data type to use for initial per-tensor reduction.
+_counter_dtype  = torch.float64 # Data type to use for the internal counters.
+_rank           = 0             # Rank of the current process.
+_sync_device    = None          # Device to use for multiprocess communication. None = single-process.
+_sync_called    = False         # Has _sync() been called yet?
+_counters       = dict()        # Running counters on each device, updated by report(): name => device => torch.Tensor
+_cumulative     = dict()        # Cumulative counters on the CPU, updated by _sync(): name => torch.Tensor
+
+#----------------------------------------------------------------------------
+
+def init_multiprocessing(rank, sync_device):
+    r"""Initializes `torch_utils.training_stats` for collecting statistics
+    across multiple processes.
+
+    This function must be called after
+    `torch.distributed.init_process_group()` and before `Collector.update()`.
+    The call is not necessary if multi-process collection is not needed.
+
+    Args:
+        rank:           Rank of the current process.
+        sync_device:    PyTorch device to use for inter-process
+                        communication, or None to disable multi-process
+                        collection. Typically `torch.device('cuda', rank)`.
+    """
+    global _rank, _sync_device
+    assert not _sync_called
+    _rank = rank
+    _sync_device = sync_device
+
+#----------------------------------------------------------------------------
+
+@misc.profiled_function
+def report(name, value):
+    r"""Broadcasts the given set of scalars to all interested instances of
+    `Collector`, across device and process boundaries. NaNs and Infs are
+    ignored.
+
+    This function is expected to be extremely cheap and can be safely
+    called from anywhere in the training loop, loss function, or inside a
+    `torch.nn.Module`.
+
+    Warning: The current implementation expects the set of unique names to
+    be consistent across processes. Please make sure that `report()` is
+    called at least once for each unique name by each process, and in the
+    same order. If a given process has no scalars to broadcast, it can do
+    `report(name, [])` (empty list).
+
+    Args:
+        name:   Arbitrary string specifying the name of the statistic.
+                Averages are accumulated separately for each unique name.
+        value:  Arbitrary set of scalars. Can be a list, tuple,
+                NumPy array, PyTorch tensor, or Python scalar.
+
+    Returns:
+        The same `value` that was passed in.
+    """
+    if name not in _counters:
+        _counters[name] = dict()
+
+    elems = torch.as_tensor(value)
+    if elems.numel() == 0:
+        return value
+
+    elems = elems.detach().flatten().to(_reduce_dtype)
+    square = elems.square()
+    finite = square.isfinite()
+    moments = torch.stack([
+        finite.sum(dtype=_reduce_dtype),
+        torch.where(finite, elems, 0).sum(),
+        torch.where(finite, square, 0).sum(),
+    ])
+    assert moments.ndim == 1 and moments.shape[0] == _num_moments
+    moments = moments.to(_counter_dtype)
+
+    device = moments.device
+    if device not in _counters[name]:
+        _counters[name][device] = torch.zeros_like(moments)
+    _counters[name][device].add_(moments)
+    return value
+
+#----------------------------------------------------------------------------
+
+def report0(name, value):
+    r"""Broadcasts the given set of scalars by the first process (`rank = 0`),
+    but ignores any scalars provided by the other processes.
+    See `report()` for further details.
+    """
+    report(name, value if _rank == 0 else [])
+    return value
+
+#----------------------------------------------------------------------------
+
+class Collector:
+    r"""Collects the scalars broadcasted by `report()` and `report0()` and
+    computes their long-term averages (mean and standard deviation) over
+    user-defined periods of time.
+
+    The averages are first collected into internal counters that are not
+    directly visible to the user. They are then copied to the user-visible
+    state as a result of calling `update()` and can then be queried using
+    `mean()`, `std()`, `as_dict()`, etc. Calling `update()` also resets the
+    internal counters for the next round, so that the user-visible state
+    effectively reflects averages collected between the last two calls to
+    `update()`.
+
+    Args:
+        regex:          Regular expression defining which statistics to
+                        collect. The default is to collect everything.
+        keep_previous:  Whether to retain the previous averages if no
+                        scalars were collected on a given round
+                        (default: False).
+    """
+    def __init__(self, regex='.*', keep_previous=False):
+        self._regex = re.compile(regex)
+        self._keep_previous = keep_previous
+        self._cumulative = dict()
+        self._moments = dict()
+        self.update()
+        self._moments.clear()
+
+    def names(self):
+        r"""Returns the names of all statistics broadcasted so far that
+        match the regular expression specified at construction time.
+        """
+        return [name for name in _counters if self._regex.fullmatch(name)]
+
+    def update(self):
+        r"""Copies current values of the internal counters to the
+        user-visible state and resets them for the next round.
+
+        If `keep_previous=True` was specified at construction time, the
+        operation is skipped for statistics that have received no scalars
+        since the last update, retaining their previous averages.
+
+        This method performs a number of GPU-to-CPU transfers and one
+        `torch.distributed.all_reduce()`. It is intended to be called
+        periodically in the main training loop, typically once every
+        N training steps.
+        """
+        if not self._keep_previous:
+            self._moments.clear()
+        for name, cumulative in _sync(self.names()):
+            if name not in self._cumulative:
+                self._cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
+            delta = cumulative - self._cumulative[name]
+            self._cumulative[name].copy_(cumulative)
+            if float(delta[0]) != 0:
+                self._moments[name] = delta
+
+    def _get_delta(self, name):
+        r"""Returns the raw moments that were accumulated for the given
+        statistic between the last two calls to `update()`, or zero if
+        no scalars were collected.
+        """
+        assert self._regex.fullmatch(name)
+        if name not in self._moments:
+            self._moments[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
+        return self._moments[name]
+
+    def num(self, name):
+        r"""Returns the number of scalars that were accumulated for the given
+        statistic between the last two calls to `update()`, or zero if
+        no scalars were collected.
+        """
+        delta = self._get_delta(name)
+        return int(delta[0])
+
+    def mean(self, name):
+        r"""Returns the mean of the scalars that were accumulated for the
+        given statistic between the last two calls to `update()`, or NaN if
+        no scalars were collected.
+        """
+        delta = self._get_delta(name)
+        if int(delta[0]) == 0:
+            return float('nan')
+        return float(delta[1] / delta[0])
+
+    def std(self, name):
+        r"""Returns the standard deviation of the scalars that were
+        accumulated for the given statistic between the last two calls to
+        `update()`, or NaN if no scalars were collected.
+        """
+        delta = self._get_delta(name)
+        if int(delta[0]) == 0 or not np.isfinite(float(delta[1])):
+            return float('nan')
+        if int(delta[0]) == 1:
+            return float(0)
+        mean = float(delta[1] / delta[0])
+        raw_var = float(delta[2] / delta[0])
+        return np.sqrt(max(raw_var - np.square(mean), 0))
+
+    def as_dict(self):
+        r"""Returns the averages accumulated between the last two calls to
+        `update()` as an `dnnlib.EasyDict`. The contents are as follows:
+
+            dnnlib.EasyDict(
+                NAME = dnnlib.EasyDict(num=FLOAT, mean=FLOAT, std=FLOAT),
+                ...
+            )
+        """
+        stats = dnnlib.EasyDict()
+        for name in self.names():
+            stats[name] = dnnlib.EasyDict(num=self.num(name), mean=self.mean(name), std=self.std(name))
+        return stats
+
+    def __getitem__(self, name):
+        r"""Convenience getter.
+        `collector[name]` is a synonym for `collector.mean(name)`.
+        """
+        return self.mean(name)
+
+#----------------------------------------------------------------------------
+
+def _sync(names):
+    r"""Synchronize the global cumulative counters across devices and
+    processes. Called internally by `Collector.update()`.
+    """
+    if len(names) == 0:
+        return []
+    global _sync_called
+    _sync_called = True
+
+    # Check that all ranks have the same set of names.
+    if _sync_device is not None:
+        value = hash(tuple(tuple(ord(char) for char in name) for name in names))
+        other = torch.as_tensor(value, dtype=torch.int64, device=_sync_device)
+        torch.distributed.broadcast(tensor=other, src=0)
+        if value != int(other.cpu()):
+            raise ValueError('Training statistics are inconsistent between ranks')
+
+    # Collect deltas within current rank.
+    deltas = []
+    device = _sync_device if _sync_device is not None else torch.device('cpu')
+    for name in names:
+        delta = torch.zeros([_num_moments], dtype=_counter_dtype, device=device)
+        for counter in _counters[name].values():
+            delta.add_(counter.to(device))
+            counter.copy_(torch.zeros_like(counter))
+        deltas.append(delta)
+    deltas = torch.stack(deltas)
+
+    # Sum deltas across ranks.
+    if _sync_device is not None:
+        torch.distributed.all_reduce(deltas)
+
+    # Update cumulative values.
+    deltas = deltas.cpu()
+    for idx, name in enumerate(names):
+        if name not in _cumulative:
+            _cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
+        _cumulative[name].add_(deltas[idx])
+
+    # Return name-value pairs.
+    return [(name, _cumulative[name]) for name in names]
+
+#----------------------------------------------------------------------------
+# Convenience.
+
+default_collector = Collector()
+
+#----------------------------------------------------------------------------

New/REG/samplers.py

@@ -0,0 +1,264 @@
+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 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,
+        ):
+    # setup conditioning
+    if cfg_scale > 1.0:
+        y_null = torch.tensor([1000] * y.size(0), device=y.device)
+        #[1000, 1000]
+    _dtype = latents.dtype
+
+
+    t_steps = torch.linspace(1., 0.04, num_steps, dtype=torch.float64)
+    t_steps = torch.cat([t_steps, torch.tensor([0.], dtype=torch.float64)])
+    x_next = latents.to(torch.float64)
+    cls_x_next = cls_latents.to(torch.float64)
+    device = x_next.device
+    z_mid, cls_mid = None, None
+    t_mid = float(t_mid)
+
+
+    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
+            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_x_cur.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_x_cur] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+            else:
+                model_input = x_cur
+                cls_model_input = cls_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)
+            cls_eps_i = torch.randn_like(cls_x_cur).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 args.cls_cfg_scale >0:
+                    cls_d_cur = cls_d_cur_uncond + args.cls_cfg_scale * (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
+            cls_x_next = cls_x_cur + cls_d_cur * dt + torch.sqrt(diffusion) * cls_deps
+            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
+
+    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_x_cur] * 2, dim=0)
+        y_cur = torch.cat([y, y_null], dim=0)
+    else:
+        model_input = x_cur
+        cls_model_input = cls_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 args.cls_cfg_scale > 0:
+            cls_d_cur = cls_d_cur_uncond + args.cls_cfg_scale * (cls_d_cur_cond - cls_d_cur_uncond)
+        else:
+            cls_d_cur = cls_d_cur_cond
+
+    mean_x = x_cur + dt * d_cur
+    cls_mean_x = cls_x_cur + dt * cls_d_cur
+
+    if return_mid_state:
+        return mean_x, z_mid, cls_mean_x if cls_mid is None else cls_mid
+    return mean_x
+
+
+def euler_sampler(
+        model,
+        latents,
+        y,
+        num_steps=20,
+        heun=False,  # not used; only for compatibility with caller
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear",  # not used for ODE (velocity parameterization directly)
+        cls_latents=None,
+        args=None
+        ):
+    """
+    REG 的 ODE 采样器：确定性（不注入扩散噪声）。
+
+    这里按照 REG/SiT 的 velocity 参数化直接做 ODE：
+      d/dt x_t = v_t
+    因此不需要把 velocity 再转成 score 再转 drift。
+    """
+    # setup conditioning
+    if cfg_scale > 1.0:
+        y_null = torch.tensor([1000] * y.size(0), device=y.device)
+    _dtype = latents.dtype
+
+    cls_cfg_scale = getattr(args, "cls_cfg_scale", 0) if args is not None else 0
+
+    # ODE 时间网格：默认从 t=1 到 t=0
+    t_steps = torch.linspace(1.0, 0.0, int(num_steps) + 1, dtype=torch.float64, device=latents.device)
+
+    x_next = latents.to(torch.float64)
+    cls_x_next = cls_latents.to(torch.float64)
+    device = x_next.device
+
+    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
+
+            # classifier-free guidance（只在指定时间段启用）
+            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_x_cur] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+            else:
+                model_input = x_cur
+                cls_model_input = cls_x_cur
+                y_cur = y
+
+            kwargs = dict(y=y_cur)
+            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),
+                **kwargs,
+                cls_token=cls_model_input.to(dtype=_dtype),
+            )
+
+            # ODE：velocity 参数化直接作为导数
+            d_cur = v_cur.to(torch.float64)
+            cls_d_cur = cls_v_cur.to(torch.float64)
+
+            # 指定时间段内进行 guidance 合成
+            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_scale > 0:
+                    cls_d_cur = cls_d_cur_uncond + cls_cfg_scale * (cls_d_cur_cond - cls_d_cur_uncond)
+                else:
+                    cls_d_cur = cls_d_cur_cond
+
+            x_next = x_cur + dt * d_cur
+            cls_x_next = cls_x_cur + dt * cls_d_cur
+
+    return x_next

New/REG/train.py

@@ -0,0 +1,573 @@
+import argparse
+import copy
+from copy import deepcopy
+import logging
+import os
+from pathlib import Path
+from collections import OrderedDict
+import json
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+
+from models.sit import SiT_models
+from loss import SILoss
+from utils import load_encoders
+
+from dataset import CustomDataset
+from diffusers.models import AutoencoderKL
+from PIL import Image
+
+from samplers import euler_maruyama_sampler
+# import wandb_utils
+import wandb
+import math
+from torchvision.utils import make_grid
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from torchvision.transforms import Normalize
+
+logger = get_logger(__name__)
+
+CLIP_DEFAULT_MEAN = (0.48145466, 0.4578275, 0.40821073)
+CLIP_DEFAULT_STD = (0.26862954, 0.26130258, 0.27577711)
+
+
+
+def preprocess_raw_image(x, enc_type):
+    resolution = x.shape[-1]
+    if 'clip' in enc_type:
+        x = x / 255.
+        x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic')
+        x = Normalize(CLIP_DEFAULT_MEAN, CLIP_DEFAULT_STD)(x)
+    elif 'mocov3' in enc_type or 'mae' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+    elif 'dinov2' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+        x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic')
+    elif 'dinov1' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+    elif 'jepa' in enc_type:
+        x = x / 255.
+        x = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)
+        x = torch.nn.functional.interpolate(x, 224 * (resolution // 256), mode='bicubic')
+
+    return x
+
+
+def array2grid(x):
+    nrow = round(math.sqrt(x.size(0)))
+    x = make_grid(x.clamp(0, 1), nrow=nrow, value_range=(0, 1))
+    x = x.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
+    return x
+
+
+@torch.no_grad()
+def sample_posterior(moments, latents_scale=1., latents_bias=0.):
+    device = moments.device
+    
+    mean, std = torch.chunk(moments, 2, dim=1)
+    z = mean + std * torch.randn_like(mean)
+    z = (z * latents_scale + latents_bias) 
+    return z 
+
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        name = name.replace("module.", "")
+        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+def create_logger(logging_dir):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    logging.basicConfig(
+        level=logging.INFO,
+        format='[\033[34m%(asctime)s\033[0m] %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S',
+        handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+    )
+    logger = logging.getLogger(__name__)
+    return logger
+
+
+def requires_grad(model, flag=True):
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+#################################################################################
+#                                  Training Loop                                #
+#################################################################################
+
+def main(args):    
+    # set accelerator
+    logging_dir = Path(args.output_dir, args.logging_dir)
+    accelerator_project_config = ProjectConfiguration(
+        project_dir=args.output_dir, logging_dir=logging_dir
+        )
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+        kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)]
+    )
+
+    if accelerator.is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
+        save_dir = os.path.join(args.output_dir, args.exp_name)
+        os.makedirs(save_dir, exist_ok=True)
+        args_dict = vars(args)
+        # Save to a JSON file
+        json_dir = os.path.join(save_dir, "args.json")
+        with open(json_dir, 'w') as f:
+            json.dump(args_dict, f, indent=4)
+        checkpoint_dir = f"{save_dir}/checkpoints"  # Stores saved model checkpoints
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        logger = create_logger(save_dir)
+        logger.info(f"Experiment directory created at {save_dir}")
+    device = accelerator.device
+    if torch.backends.mps.is_available():
+        accelerator.native_amp = False    
+    if args.seed is not None:
+        set_seed(args.seed + accelerator.process_index)
+    
+    # Create model:
+    assert args.resolution % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
+    latent_size = args.resolution // 8
+
+    if args.enc_type != None:
+        encoders, encoder_types, architectures = load_encoders(
+            args.enc_type, device, args.resolution
+            )
+    else:
+        raise NotImplementedError()
+    z_dims = [encoder.embed_dim for encoder in encoders] if args.enc_type != 'None' else [0]
+    block_kwargs = {"fused_attn": args.fused_attn, "qk_norm": args.qk_norm}
+    model = SiT_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes,
+        use_cfg = (args.cfg_prob > 0),
+        z_dims = z_dims,
+        encoder_depth=args.encoder_depth,
+        **block_kwargs
+    )
+
+    model = model.to(device)
+    ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+    requires_grad(ema, False)
+    
+    latents_scale = torch.tensor(
+        [0.18215, 0.18215, 0.18215, 0.18215]
+        ).view(1, 4, 1, 1).to(device)
+    latents_bias = torch.tensor(
+        [0., 0., 0., 0.]
+        ).view(1, 4, 1, 1).to(device)
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device)
+    vae.eval()
+
+    # create loss function
+    loss_fn = SILoss(
+        prediction=args.prediction,
+        path_type=args.path_type, 
+        encoders=encoders,
+        accelerator=accelerator,
+        latents_scale=latents_scale,
+        latents_bias=latents_bias,
+        weighting=args.weighting
+    )
+    if accelerator.is_main_process:
+        logger.info(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
+    
+    # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )    
+    
+    # Setup data:
+    train_dataset = CustomDataset(
+        args.data_dir, semantic_features_dir=args.semantic_features_dir
+    )
+    local_batch_size = int(args.batch_size // accelerator.num_processes)
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_size=local_batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True
+    )
+    if accelerator.is_main_process:
+        logger.info(f"Dataset contains {len(train_dataset):,} images ({args.data_dir})")
+    
+    # Prepare models for training:
+    update_ema(ema, model, decay=0)  # Ensure EMA is initialized with synced weights
+    model.train()  # important! This enables embedding dropout for classifier-free guidance
+    ema.eval()  # EMA model should always be in eval mode
+    
+    # resume:
+    global_step = 0
+    if args.resume_step > 0:
+        ckpt_name = str(args.resume_step).zfill(7) +'.pt'
+        ckpt = torch.load(
+            f'{os.path.join(args.output_dir, args.exp_name)}/checkpoints/{ckpt_name}',
+            map_location='cpu',
+            )
+        model.load_state_dict(ckpt['model'])
+        ema.load_state_dict(ckpt['ema'])
+        optimizer.load_state_dict(ckpt['opt'])
+        global_step = ckpt['steps']
+
+    model, optimizer, train_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader
+    )
+
+    if accelerator.is_main_process:
+        tracker_config = vars(copy.deepcopy(args))
+        accelerator.init_trackers(
+            project_name="REG",
+            config=tracker_config,
+            init_kwargs={
+                "wandb": {"name": f"{args.exp_name}"}
+            },
+        )
+
+        
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=global_step,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=not accelerator.is_local_main_process,
+    )
+
+    # Labels to condition the model with (feel free to change):
+    sample_batch_size = 64 // accelerator.num_processes
+    first_batch = next(iter(train_dataloader))
+    preprocessed_semantic = len(first_batch) == 4
+    if preprocessed_semantic:
+        gt_raw_images, gt_xs, _r_pre, _y = first_batch
+    else:
+        gt_raw_images, gt_xs, _y = first_batch
+    # 仅在“非预处理 semantic 模式”下，raw_image 是 RGB 图（分辨率应与 args.resolution 对齐）。
+    if not preprocessed_semantic:
+        assert gt_raw_images.shape[-1] == args.resolution
+    gt_xs = gt_xs[:sample_batch_size]
+    gt_xs = sample_posterior(
+        gt_xs.to(device), latents_scale=latents_scale, latents_bias=latents_bias
+        )
+    ys = torch.randint(1000, size=(sample_batch_size,), device=device)
+    ys = ys.to(device)
+    # Create sampling noise:
+    n = ys.size(0)
+    xT = torch.randn((n, 4, latent_size, latent_size), device=device)
+        
+    for epoch in range(args.epochs):
+        model.train()
+        for batch in train_dataloader:
+            if len(batch) == 4:
+                raw_image, x, r_preprocessed, y = batch
+                r_preprocessed = r_preprocessed.to(device).float()
+            else:
+                raw_image, x, y = batch
+                r_preprocessed = None
+
+            raw_image = raw_image.to(device)
+            x = x.squeeze(dim=1).to(device)
+            y = y.to(device)
+
+            z = None
+            if args.legacy:
+                # In our early experiments, we accidentally apply label dropping twice: 
+                # once in train.py and once in sit.py. 
+                # We keep this option for exact reproducibility with previous runs.
+                drop_ids = torch.rand(y.shape[0], device=y.device) < args.cfg_prob
+                labels = torch.where(drop_ids, args.num_classes, y)
+            else:
+                labels = y
+            with torch.no_grad():
+                x = sample_posterior(x, latents_scale=latents_scale, latents_bias=latents_bias)
+                zs = []
+                if r_preprocessed is not None:
+                    # 预处理 semantic 模式：直接用 cls token 构造 dense tokens。
+                    cls_token = r_preprocessed
+                    while cls_token.dim() > 2:
+                        cls_token = cls_token.squeeze(1)
+
+                    base_m = model.module if hasattr(model, "module") else model
+                    n_pad = base_m.x_embedder.num_patches
+                    zs = [
+                        torch.cat(
+                            [
+                                cls_token.unsqueeze(1),
+                                cls_token.unsqueeze(1).expand(-1, n_pad, -1),
+                            ],
+                            dim=1,
+                        )
+                    ]
+                else:
+                    # 在线 encoder 模式：与原 New/REG 行为一致
+                    with accelerator.autocast():
+                        for encoder, encoder_type, arch in zip(
+                            encoders, encoder_types, architectures
+                        ):
+                            raw_image_ = preprocess_raw_image(raw_image, encoder_type)
+                            z = encoder.forward_features(raw_image_)
+                            if "dinov2" in encoder_type:
+                                dense_z = z["x_norm_patchtokens"]
+                                cls_token = z["x_norm_clstoken"]
+                                dense_z = torch.cat(
+                                    [cls_token.unsqueeze(1), dense_z], dim=1
+                                )
+                            else:
+                                exit()
+                            zs.append(dense_z)
+
+            with accelerator.accumulate(model):
+                model_kwargs = dict(y=labels)
+                loss1, proj_loss1, time_input, noises, loss2 = loss_fn(model, x, model_kwargs, zs=zs,
+                                                                       cls_token=cls_token,
+                                                                       time_input=None, noises=None)
+                loss_mean = loss1.mean()
+                loss_mean_cls = loss2.mean() * args.cls
+                proj_loss_mean = proj_loss1.mean() * args.proj_coeff
+                loss = loss_mean + proj_loss_mean + loss_mean_cls
+
+
+                ## optimization
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    params_to_clip = model.parameters()
+                    grad_norm = accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                if accelerator.sync_gradients:
+                    update_ema(ema, model) # change ema function
+            
+            ### enter
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+            if global_step % args.checkpointing_steps == 0 and global_step > 0:
+                if accelerator.is_main_process:
+                    checkpoint = {
+                        "model": model.module.state_dict(),
+                        "ema": ema.state_dict(),
+                        "opt": optimizer.state_dict(),
+                        "args": args,
+                        "steps": global_step,
+                    }
+                    checkpoint_path = f"{checkpoint_dir}/{global_step:07d}.pt"
+                    torch.save(checkpoint, checkpoint_path)
+                    logger.info(f"Saved checkpoint to {checkpoint_path}")
+
+            if (global_step == 1 or (global_step % args.sampling_steps == 0 and global_step > 0)):
+                t_mid_vis = float(args.t_c)
+                tc_tag = f"{t_mid_vis:.4f}".rstrip("0").rstrip(".").replace(".", "_")
+                logging.info(
+                    f"Generating EMA samples (Euler-Maruyama; t≈{t_mid_vis:g} -> t=0)..."
+                )
+                ema.eval()
+                with torch.no_grad():
+                    latent_size = args.resolution // 8
+                    n_samples = min(16, args.batch_size)
+                    base_model = model.module if hasattr(model, "module") else model
+                    cls_dim = base_model.z_dims[0]
+                    shared_seed = torch.randint(0, 2**32, (1,), device=device).item()
+                    torch.manual_seed(shared_seed)
+                    z_init = torch.randn(n_samples, base_model.in_channels, latent_size, latent_size, device=device)
+                    torch.manual_seed(shared_seed)
+                    cls_init = torch.randn(n_samples, cls_dim, device=device)
+                    y_samples = torch.randint(0, args.num_classes, (n_samples,), device=device)
+
+                    z_0, z_mid, _ = euler_maruyama_sampler(
+                        ema,
+                        z_init,
+                        y_samples,
+                        num_steps=50,
+                        cfg_scale=1.0,
+                        guidance_low=0.0,
+                        guidance_high=1.0,
+                        path_type=args.path_type,
+                        cls_latents=cls_init,
+                        args=args,
+                        return_mid_state=True,
+                        t_mid=t_mid_vis,
+                    )
+
+                    samples_root = os.path.join(args.output_dir, args.exp_name, "samples")
+                    t0_dir = os.path.join(samples_root, "t0")
+                    t_mid_dir = os.path.join(samples_root, f"t0_{tc_tag}")
+                    os.makedirs(t0_dir, exist_ok=True)
+                    os.makedirs(t_mid_dir, exist_ok=True)
+
+                    z_f = z_0.to(dtype=torch.float32)
+                    samples_final = vae.decode((z_f - latents_bias) / latents_scale).sample
+                    samples_final = (samples_final + 1) / 2.0
+                    samples_final = samples_final.clamp(0, 1)
+                    grid_final = array2grid(samples_final)
+                    Image.fromarray(grid_final).save(
+                        os.path.join(t0_dir, f"step_{global_step:07d}_t0.png")
+                    )
+
+                    if z_mid is not None:
+                        z_m = z_mid.to(dtype=torch.float32)
+                        samples_mid = vae.decode((z_m - latents_bias) / latents_scale).sample
+                        samples_mid = (samples_mid + 1) / 2.0
+                        samples_mid = samples_mid.clamp(0, 1)
+                        grid_mid = array2grid(samples_mid)
+                        Image.fromarray(grid_mid).save(
+                            os.path.join(t_mid_dir, f"step_{global_step:07d}_t0_{tc_tag}.png")
+                        )
+                    else:
+                        logging.warning(
+                            f"Sampling time grid did not bracket t_mid={t_mid_vis:g}; "
+                            f"skip t0_{tc_tag} image this step."
+                        )
+
+                    del z_init, cls_init, y_samples, z_0
+                    if z_mid is not None:
+                        del z_mid
+                    del samples_final, grid_final
+                    if "samples_mid" in locals():
+                        del samples_mid, grid_mid
+                    torch.cuda.empty_cache()
+
+            logs = {
+                "loss_final": accelerator.gather(loss).mean().detach().item(),
+                "loss_mean": accelerator.gather(loss_mean).mean().detach().item(),
+                "proj_loss": accelerator.gather(proj_loss_mean).mean().detach().item(),
+                "loss_mean_cls": accelerator.gather(loss_mean_cls).mean().detach().item(),
+                "grad_norm": accelerator.gather(grad_norm).mean().detach().item()
+            }
+
+            log_message = ", ".join(f"{key}: {value:.6f}" for key, value in logs.items())
+            logging.info(f"Step: {global_step}, Training Logs: {log_message}")
+
+            progress_bar.set_postfix(**logs)
+            accelerator.log(logs, step=global_step)
+
+            if global_step >= args.max_train_steps:
+                break
+        if global_step >= args.max_train_steps:
+            break
+
+    model.eval()  # important! This disables randomized embedding dropout
+    # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
+    
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        logger.info("Done!")
+    accelerator.end_training()
+
+def parse_args(input_args=None):
+    parser = argparse.ArgumentParser(description="Training")
+
+    # logging:
+    parser.add_argument("--output-dir", type=str, default="exps")
+    parser.add_argument("--exp-name", type=str, required=True)
+    parser.add_argument("--logging-dir", type=str, default="logs")
+    parser.add_argument("--report-to", type=str, default="wandb")
+    parser.add_argument("--sampling-steps", type=int, default=10000)
+    parser.add_argument("--resume-step", type=int, default=0)
+
+    # model
+    parser.add_argument("--model", type=str)
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--encoder-depth", type=int, default=8)
+    parser.add_argument("--fused-attn", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--qk-norm",  action=argparse.BooleanOptionalAction, default=False)
+    parser.add_argument("--ops-head", type=int, default=16)
+
+    # dataset
+    parser.add_argument("--data-dir", type=str, default="../data/imagenet256")
+    parser.add_argument(
+        "--semantic-features-dir",
+        type=str,
+        default=None,
+        help="预处理 semantic features 目录（与 REG/dataset.py 语义相同），仅影响数据加载方式，不引入 t_c。",
+    )
+    parser.add_argument("--resolution", type=int, choices=[256, 512], default=256)
+    parser.add_argument("--batch-size", type=int, default=8)#256
+
+    # precision
+    parser.add_argument("--allow-tf32", action="store_true")
+    parser.add_argument("--mixed-precision", type=str, default="fp16", choices=["no", "fp16", "bf16"])
+
+    # optimization
+    parser.add_argument("--epochs", type=int, default=1400)
+    parser.add_argument("--max-train-steps", type=int, default=1000000)
+    parser.add_argument("--checkpointing-steps", type=int, default=10000)
+    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
+    parser.add_argument("--learning-rate", type=float, default=1e-4)
+    parser.add_argument("--adam-beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam-beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam-weight-decay", type=float, default=0., help="Weight decay to use.")
+    parser.add_argument("--adam-epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max-grad-norm", default=1.0, type=float, help="Max gradient norm.")
+
+    # seed
+    parser.add_argument("--seed", type=int, default=0)
+
+    # cpu
+    parser.add_argument("--num-workers", type=int, default=4)
+
+    # loss
+    parser.add_argument("--path-type", type=str, default="linear", choices=["linear", "cosine"])
+    parser.add_argument("--prediction", type=str, default="v", choices=["v"]) # currently we only support v-prediction
+    parser.add_argument("--cfg-prob", type=float, default=0.1)
+    parser.add_argument("--enc-type", type=str, default='dinov2-vit-b')
+    parser.add_argument("--proj-coeff", type=float, default=0.5)
+    parser.add_argument("--weighting", default="uniform", type=str, help="Max gradient norm.")
+    parser.add_argument("--legacy", action=argparse.BooleanOptionalAction, default=False)
+    parser.add_argument("--cls", type=float, default=0.03)
+    parser.add_argument(
+        "--t-c",
+        type=float,
+        default=0.5,
+        help="训练中采样时保存的中间时刻 t（用于输出 t0 与 t0_tc 对比图）。",
+    )
+    if input_args is not None:
+        args = parser.parse_args(input_args)
+    else:
+        args = parser.parse_args()
+
+    return args
+
+if __name__ == "__main__":
+    args = parse_args()
+    
+    main(args)

New/REG/train.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+# New/REG/train.sh：只对齐数据加载相关目录/参数；不引入 t_c/ot-cls。
+
+set -euo pipefail
+
+NUM_GPUS=4
+
+# ------------------------- 按需修改这些路径 -------------------------
+# New/REG/dataset.py 期望 data_dir 下至少有：
+#   - imagenet_256_vae/
+#   - vae-sd/（非预处理 semantic 模式需要）
+# 若启用预处理语义：还需要
+#   - imagenet_256_features/dinov2-vit-b_tmp/gpu0/dataset.json + 对应 .npy
+DATA_DIR="/gemini/space/zhaozy/dataset/Imagenet/imagenet_256"
+SEMANTIC_FEATURES_DIR="/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0"
+
+OUTPUT_BASE_DIR="/your_path/reg_xlarge_dinov2_base_align_8_cls"
+EXP_NAME="linear-dinov2-b-enc8"
+
+# ------------------------- 不建议改动以下逻辑 -------------------------
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+TRAIN_PY="${SCRIPT_DIR}/train.py"
+
+nohup 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 \
+  --cls=0.03 \
+  --output-dir="${OUTPUT_BASE_DIR}" \
+  --exp-name="${EXP_NAME}" \
+  --batch-size=256 \
+  --data-dir="${DATA_DIR}" \
+  --semantic-features-dir="${SEMANTIC_FEATURES_DIR}" \
+   > jsflow-experiment.log 2>&1 & 

New/REG/utils.py

@@ -0,0 +1,225 @@
+import os
+from torchvision.datasets.utils import download_url
+import torch
+import torchvision.models as torchvision_models
+import timm
+from models import mocov3_vit
+import math
+import warnings
+
+
+# code from SiT repository
+pretrained_models = {'last.pt'}
+
+def download_model(model_name):
+    """
+    Downloads a pre-trained SiT model from the web.
+    """
+    assert model_name in pretrained_models
+    local_path = f'pretrained_models/{model_name}'
+    if not os.path.isfile(local_path):
+        os.makedirs('pretrained_models', exist_ok=True)
+        web_path = f'https://www.dl.dropboxusercontent.com/scl/fi/cxedbs4da5ugjq5wg3zrg/last.pt?rlkey=8otgrdkno0nd89po3dpwngwcc&st=apcc645o&dl=0'
+        download_url(web_path, 'pretrained_models', filename=model_name)
+    model = torch.load(local_path, map_location=lambda storage, loc: storage)
+    return model
+
+def fix_mocov3_state_dict(state_dict):
+    for k in list(state_dict.keys()):
+        # retain only base_encoder up to before the embedding layer
+        if k.startswith('module.base_encoder'):
+            # fix naming bug in checkpoint
+            new_k = k[len("module.base_encoder."):]
+            if "blocks.13.norm13" in new_k:
+                new_k = new_k.replace("norm13", "norm1")
+            if "blocks.13.mlp.fc13" in k:
+                new_k = new_k.replace("fc13", "fc1")
+            if "blocks.14.norm14" in k:
+                new_k = new_k.replace("norm14", "norm2")
+            if "blocks.14.mlp.fc14" in k:
+                new_k = new_k.replace("fc14", "fc2")
+            # remove prefix
+            if 'head' not in new_k and new_k.split('.')[0] != 'fc':
+                state_dict[new_k] = state_dict[k]
+        # delete renamed or unused k
+        del state_dict[k]
+    if 'pos_embed' in state_dict.keys():
+        state_dict['pos_embed'] = timm.layers.pos_embed.resample_abs_pos_embed(
+            state_dict['pos_embed'], [16, 16],
+        )
+    return state_dict
+
+@torch.no_grad()
+def load_encoders(enc_type, device, resolution=256):
+    assert (resolution == 256) or (resolution == 512)
+    
+    enc_names = enc_type.split(',')
+    encoders, architectures, encoder_types = [], [], []
+    for enc_name in enc_names:
+        encoder_type, architecture, model_config = enc_name.split('-')
+        # Currently, we only support 512x512 experiments with DINOv2 encoders.
+        if resolution == 512:
+            if encoder_type != 'dinov2':
+                raise NotImplementedError(
+                    "Currently, we only support 512x512 experiments with DINOv2 encoders."
+                    )
+
+        architectures.append(architecture)
+        encoder_types.append(encoder_type)
+        if encoder_type == 'mocov3':
+            if architecture == 'vit':
+                if model_config == 's':
+                    encoder = mocov3_vit.vit_small()
+                elif model_config == 'b':
+                    encoder = mocov3_vit.vit_base()
+                elif model_config == 'l':
+                    encoder = mocov3_vit.vit_large()
+                ckpt = torch.load(f'./ckpts/mocov3_vit{model_config}.pth')
+                state_dict = fix_mocov3_state_dict(ckpt['state_dict'])
+                del encoder.head
+                encoder.load_state_dict(state_dict, strict=True)
+                encoder.head = torch.nn.Identity()
+            elif architecture == 'resnet':
+                raise NotImplementedError()
+ 
+            encoder = encoder.to(device)
+            encoder.eval()
+
+        elif 'dinov2' in encoder_type:
+            import timm
+            if 'reg' in encoder_type:
+                try:
+                    encoder = torch.hub.load('your_path/.cache/torch/hub/facebookresearch_dinov2_main',
+                                            f'dinov2_vit{model_config}14_reg', source='local')
+                except:
+                    encoder = torch.hub.load('facebookresearch/dinov2', f'dinov2_vit{model_config}14_reg')
+            else:
+                try:
+                    encoder = torch.hub.load('your_path/.cache/torch/hub/facebookresearch_dinov2_main',
+                                             f'dinov2_vit{model_config}14', source='local')
+                except:
+                    encoder = torch.hub.load('facebookresearch/dinov2', f'dinov2_vit{model_config}14')
+
+            print(f"Now you are using the {enc_name} as the aligning model")
+            del encoder.head
+            patch_resolution = 16 * (resolution // 256)
+            encoder.pos_embed.data = timm.layers.pos_embed.resample_abs_pos_embed(
+                encoder.pos_embed.data, [patch_resolution, patch_resolution],
+            )
+            encoder.head = torch.nn.Identity()
+            encoder = encoder.to(device)
+            encoder.eval()
+        
+        elif 'dinov1' == encoder_type:
+            import timm
+            from models import dinov1
+            encoder = dinov1.vit_base()
+            ckpt =  torch.load(f'./ckpts/dinov1_vit{model_config}.pth') 
+            if 'pos_embed' in ckpt.keys():
+                ckpt['pos_embed'] = timm.layers.pos_embed.resample_abs_pos_embed(
+                    ckpt['pos_embed'], [16, 16],
+                )
+            del encoder.head
+            encoder.head = torch.nn.Identity()
+            encoder.load_state_dict(ckpt, strict=True)
+            encoder = encoder.to(device)
+            encoder.forward_features = encoder.forward
+            encoder.eval()
+
+        elif encoder_type == 'clip':
+            import clip
+            from models.clip_vit import UpdatedVisionTransformer
+            encoder_ = clip.load(f"ViT-{model_config}/14", device='cpu')[0].visual
+            encoder = UpdatedVisionTransformer(encoder_).to(device)
+             #.to(device)
+            encoder.embed_dim = encoder.model.transformer.width
+            encoder.forward_features = encoder.forward
+            encoder.eval()
+        
+        elif encoder_type == 'mae':
+            from models.mae_vit import vit_large_patch16
+            import timm
+            kwargs = dict(img_size=256)
+            encoder = vit_large_patch16(**kwargs).to(device)
+            with open(f"ckpts/mae_vit{model_config}.pth", "rb") as f:
+                state_dict = torch.load(f)
+            if 'pos_embed' in state_dict["model"].keys():
+                state_dict["model"]['pos_embed'] = timm.layers.pos_embed.resample_abs_pos_embed(
+                    state_dict["model"]['pos_embed'], [16, 16],
+                )
+            encoder.load_state_dict(state_dict["model"])
+
+            encoder.pos_embed.data = timm.layers.pos_embed.resample_abs_pos_embed(
+                encoder.pos_embed.data, [16, 16],
+            )
+
+        elif encoder_type == 'jepa':
+            from models.jepa import vit_huge
+            kwargs = dict(img_size=[224, 224], patch_size=14)
+            encoder = vit_huge(**kwargs).to(device)
+            with open(f"ckpts/ijepa_vit{model_config}.pth", "rb") as f:
+                state_dict = torch.load(f, map_location=device)
+            new_state_dict = dict()
+            for key, value in state_dict['encoder'].items():
+                new_state_dict[key[7:]] = value
+            encoder.load_state_dict(new_state_dict)
+            encoder.forward_features = encoder.forward
+
+        encoders.append(encoder)
+    
+    return encoders, encoder_types, architectures
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def load_legacy_checkpoints(state_dict, encoder_depth):
+    new_state_dict = dict()
+    for key, value in state_dict.items():
+        if 'decoder_blocks' in key:
+            parts =key.split('.')
+            new_idx = int(parts[1]) + encoder_depth
+            parts[0] = 'blocks'
+            parts[1] = str(new_idx)
+            new_key = '.'.join(parts)
+            new_state_dict[new_key] = value
+        else:
+            new_state_dict[key] = value
+    return new_state_dict

New/REG/wandb/debug-internal.log

@@ -0,0 +1,7 @@
+{"time":"2026-03-26T13:08:48.077546544+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-26T13:08:49.00834408+08:00","level":"INFO","msg":"stream: created new stream","id":"0e5vs4f8"}
+{"time":"2026-03-26T13:08:49.008541708+08:00","level":"INFO","msg":"handler: started","stream_id":"0e5vs4f8"}
+{"time":"2026-03-26T13:08:49.009420311+08:00","level":"INFO","msg":"stream: started","id":"0e5vs4f8"}
+{"time":"2026-03-26T13:08:49.00944+08:00","level":"INFO","msg":"writer: started","stream_id":"0e5vs4f8"}
+{"time":"2026-03-26T13:08:49.009442857+08:00","level":"INFO","msg":"sender: started","stream_id":"0e5vs4f8"}
+{"time":"2026-03-27T00:25:27.785418795+08:00","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/0e5vs4f8/file_stream","body":"\n<html><head>\n<meta http-equiv=\"content-type\" content=\"text/html;charset=utf-8\">\n<title>502 Server Error</title>\n</head>\n<body text=#000000 bgcolor=#ffffff>\n<h1>Error: Server Error</h1>\n<h2>The server encountered a temporary error and could not complete your request.<p>Please try again in 30 seconds.</h2>\n<h2></h2>\n</body></html>\n"}

New/REG/wandb/debug.log

@@ -0,0 +1,20 @@
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_setup.py:_flush():81] Configure stats pid to 576284
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_init.py:setup_run_log_directory():717] Logging user logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG/wandb/run-20260326_130847-0e5vs4f8/logs/debug.log
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG/wandb/run-20260326_130847-0e5vs4f8/logs/debug-internal.log
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_init.py:init():844] calling init triggers
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-26 13:08:47,761 INFO    MainThread:576284 [wandb_init.py:init():892] starting backend
+2026-03-26 13:08:48,060 INFO    MainThread:576284 [wandb_init.py:init():895] sending inform_init request
+2026-03-26 13:08:48,073 INFO    MainThread:576284 [wandb_init.py:init():903] backend started and connected
+2026-03-26 13:08:48,075 INFO    MainThread:576284 [wandb_init.py:init():973] updated telemetry
+2026-03-26 13:08:48,089 INFO    MainThread:576284 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-26 13:08:49,730 INFO    MainThread:576284 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-26 13:08:49,821 INFO    MainThread:576284 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-26 13:08:49,821 INFO    MainThread:576284 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-26 13:08:49,821 INFO    MainThread:576284 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-26 13:08:49,821 INFO    MainThread:576284 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-26 13:08:49,826 INFO    MainThread:576284 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-26 13:08:49,826 INFO    MainThread:576284 [wandb_run.py:_config_callback():1403] config_cb None None {'output_dir': '/your_path/reg_xlarge_dinov2_base_align_8_cls', 'exp_name': 'linear-dinov2-b-enc8', 'logging_dir': 'logs', 'report_to': 'wandb', 'sampling_steps': 10000, 'resume_step': 0, 'train_generate_step': 2000, 'train_generate_num_steps': 2000, 'train_generate_num_images': 16, 'train_generate_mode': 'sde', 'train_generate_cfg_scale': 2.3, 'train_generate_cls_cfg_scale': 2.3, 'train_generate_guidance_high': 0.85, 'train_generate_vae': 'ema', 'model': 'SiT-XL/2', 'num_classes': 1000, 'encoder_depth': 8, 'fused_attn': True, 'qk_norm': False, 'ops_head': 16, 'data_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256', 'semantic_features_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0', 'resolution': 256, 'batch_size': 256, 'allow_tf32': True, 'mixed_precision': 'fp16', 'epochs': 1400, 'max_train_steps': 1000000, 'checkpointing_steps': 10000, 'gradient_accumulation_steps': 1, 'learning_rate': 0.0001, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_weight_decay': 0.0, 'adam_epsilon': 1e-08, 'max_grad_norm': 1.0, 'seed': 0, 'num_workers': 4, 'path_type': 'linear', 'prediction': 'v', 'cfg_prob': 0.1, 'enc_type': 'dinov2-vit-b', 'proj_coeff': 0.5, 'weighting': 'uniform', 'legacy': False, 'cls': 0.03}

New/REG/wandb/run-20260326_123101-m3lli51t/files/requirements.txt

@@ -0,0 +1,168 @@
+dill==0.3.8
+mkl-service==2.4.0
+mpmath==1.3.0
+typing_extensions==4.12.2
+urllib3==2.3.0
+torch==2.5.1
+ptyprocess==0.7.0
+traitlets==5.14.3
+pyasn1==0.6.1
+opencv-python-headless==4.12.0.88
+nest-asyncio==1.6.0
+kiwisolver==1.4.8
+click==8.2.1
+fire==0.7.1
+diffusers==0.35.1
+accelerate==1.7.0
+ipykernel==6.29.5
+peft==0.17.1
+attrs==24.3.0
+six==1.17.0
+numpy==2.0.1
+yarl==1.18.0
+huggingface_hub==0.34.4
+Bottleneck==1.4.2
+numexpr==2.11.0
+dataclasses==0.6
+typing-inspection==0.4.1
+safetensors==0.5.3
+pyparsing==3.2.3
+psutil==7.0.0
+imageio==2.37.0
+debugpy==1.8.14
+cycler==0.12.1
+pyasn1_modules==0.4.2
+matplotlib-inline==0.1.7
+matplotlib==3.10.3
+jedi==0.19.2
+tokenizers==0.21.2
+seaborn==0.13.2
+timm==1.0.15
+aiohappyeyeballs==2.6.1
+hf-xet==1.1.8
+multidict==6.1.0
+tqdm==4.67.1
+wheel==0.45.1
+simsimd==6.5.1
+sentencepiece==0.2.1
+grpcio==1.74.0
+asttokens==3.0.0
+absl-py==2.3.1
+stack-data==0.6.3
+pandas==2.3.0
+importlib_metadata==8.7.0
+pytorch-image-generation-metrics==0.6.1
+frozenlist==1.5.0
+MarkupSafe==3.0.2
+setuptools==78.1.1
+multiprocess==0.70.15
+pip==25.1
+requests==2.32.3
+mkl_random==1.2.8
+tensorboard-plugin-wit==1.8.1
+ExifRead-nocycle==3.0.1
+webdataset==0.2.111
+threadpoolctl==3.6.0
+pyarrow==21.0.0
+executing==2.2.0
+decorator==5.2.1
+contourpy==1.3.2
+annotated-types==0.7.0
+scikit-learn==1.7.1
+jupyter_client==8.6.3
+albumentations==1.4.24
+wandb==0.25.0
+certifi==2025.8.3
+idna==3.7
+xxhash==3.5.0
+Jinja2==3.1.6
+python-dateutil==2.9.0.post0
+aiosignal==1.4.0
+triton==3.1.0
+torchvision==0.20.1
+stringzilla==3.12.6
+pure_eval==0.2.3
+braceexpand==0.1.7
+zipp==3.22.0
+oauthlib==3.3.1
+Markdown==3.8.2
+fsspec==2025.3.0
+fonttools==4.58.2
+comm==0.2.2
+ipython==9.3.0
+img2dataset==1.47.0
+networkx==3.4.2
+PySocks==1.7.1
+tzdata==2025.2
+smmap==5.0.2
+mkl_fft==1.3.11
+sentry-sdk==2.29.1
+Pygments==2.19.1
+pexpect==4.9.0
+ftfy==6.3.1
+einops==0.8.1
+requests-oauthlib==2.0.0
+gitdb==4.0.12
+albucore==0.0.23
+torchdiffeq==0.2.5
+GitPython==3.1.44
+bitsandbytes==0.47.0
+pytorch-fid==0.3.0
+clean-fid==0.1.35
+pytorch-gan-metrics==0.5.4
+Brotli==1.0.9
+charset-normalizer==3.3.2
+gmpy2==2.2.1
+pillow==11.1.0
+PyYAML==6.0.2
+tornado==6.5.1
+termcolor==3.1.0
+setproctitle==1.3.6
+scipy==1.15.3
+regex==2024.11.6
+protobuf==6.31.1
+platformdirs==4.3.8
+joblib==1.5.1
+cachetools==4.2.4
+ipython_pygments_lexers==1.1.1
+google-auth==1.35.0
+transformers==4.53.2
+torch-fidelity==0.3.0
+tensorboard==2.4.0
+filelock==3.17.0
+packaging==25.0
+propcache==0.3.1
+pytz==2025.2
+aiohttp==3.11.10
+wcwidth==0.2.13
+clip==0.2.0
+Werkzeug==3.1.3
+tensorboard-data-server==0.6.1
+sympy==1.13.1
+pyzmq==26.4.0
+pydantic_core==2.33.2
+prompt_toolkit==3.0.51
+parso==0.8.4
+docker-pycreds==0.4.0
+rsa==4.9.1
+pydantic==2.11.5
+jupyter_core==5.8.1
+google-auth-oauthlib==0.4.6
+datasets==4.0.0
+torch-tb-profiler==0.4.3
+autocommand==2.2.2
+backports.tarfile==1.2.0
+importlib_metadata==8.0.0
+jaraco.collections==5.1.0
+jaraco.context==5.3.0
+jaraco.functools==4.0.1
+more-itertools==10.3.0
+packaging==24.2
+platformdirs==4.2.2
+typeguard==4.3.0
+inflect==7.3.1
+jaraco.text==3.12.1
+tomli==2.0.1
+typing_extensions==4.12.2
+wheel==0.45.1
+zipp==3.19.2

New/REG/wandb/run-20260326_123101-m3lli51t/files/wandb-metadata.json

@@ -0,0 +1,80 @@
+{
+  "os":  "Linux-5.15.0-94-generic-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-03-26T04:31:01.251691Z",
+  "args":  [
+    "--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",
+    "--cls=0.03",
+    "--output-dir=/your_path/reg_xlarge_dinov2_base_align_8_cls",
+    "--exp-name=linear-dinov2-b-enc8",
+    "--batch-size=256",
+    "--data-dir=/gemini/space/zhaozy/dataset/Imagenet/imagenet_256",
+    "--semantic-features-dir=/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0"
+  ],
+  "program":  "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG/train.py",
+  "codePath":  "train.py",
+  "codePathLocal":  "train.py",
+  "git":  {
+    "remote":  "https://github.com/Martinser/REG.git",
+    "commit":  "021ea2e50c38c5803bd9afff16316958a01fbd1d"
+  },
+  "email":  "2365972933@qq.com",
+  "root":  "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG",
+  "host":  "24c964746905d416ce09d045f9a06f23-taskrole1-0",
+  "executable":  "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/bin/python",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  4,
+  "disk":  {
+    "/":  {
+      "total":  "3838880616448",
+      "used":  "367215697920"
+    }
+  },
+  "memory":  {
+    "total":  "2164115296256"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-757303bb-4ec2-808b-a17f-95f6f5bad6dc"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-a09f2421-99e6-a72e-63bd-fd7452510758"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-9c670cc7-60a8-17f8-9b39-7ced3744976d"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-e6b1d8da-68d7-ed83-90d0-a4dedf33120e"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "shmu4nsng7t9zxqf30qry601bzvrritx"
+}

New/REG/wandb/run-20260326_123101-m3lli51t/logs/debug-internal.log

@@ -0,0 +1,8 @@
+{"time":"2026-03-26T12:31:01.594683189+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-26T12:31:02.390583334+08:00","level":"INFO","msg":"stream: created new stream","id":"m3lli51t"}
+{"time":"2026-03-26T12:31:02.390723765+08:00","level":"INFO","msg":"handler: started","stream_id":"m3lli51t"}
+{"time":"2026-03-26T12:31:02.39153628+08:00","level":"INFO","msg":"stream: started","id":"m3lli51t"}
+{"time":"2026-03-26T12:31:02.391552999+08:00","level":"INFO","msg":"writer: started","stream_id":"m3lli51t"}
+{"time":"2026-03-26T12:31:02.391598449+08:00","level":"INFO","msg":"sender: started","stream_id":"m3lli51t"}
+{"time":"2026-03-26T12:31:25.872492492+08:00","level":"INFO","msg":"api: retrying error","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/2365972933-teleai/REG/m3lli51t/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20260326%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20260326T043108Z&X-Goog-Expires=86399&X-Goog-Signature=b7efae9bb95ca96e0208d793a895c400eb9f42714e7f79e1ccff1fb97664053ab3d4a38f98cba45f7336911ceef6bb02894228ed5cf3b4dffdd02698b5a31a7faa483fd85396d291642957c4420cf884bf125233c9c9cec3ce04029a41984d86f1db898738ae442aae0015236bafa9031b20eedd7a081f77402ea66e0e54e4f65adadb2c11f9850863e10e75f24b9438f3cf3823f5bdf9d938245fc91ee6680fd95fa221fc69e487ffde730b8fc8bcd8fa79ae686fb7bd8ae6c1bd4b25d8eef9f64931d06496a6c52b7ed9dbe4a0716ed4f7fbb833b8470f3ae8fde3da8f15dc226596019add90ffbc322acc6147ab83fd50eca1690a3060e947a76b97d42b47&X-Goog-SignedHeaders=host&X-User=2365972933\": read tcp 172.20.98.27:43208->142.250.199.219:443: read: connection reset by peer"}
+{"time":"2026-03-26T12:31:27.355342484+08:00","level":"INFO","msg":"api: retrying error","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/2365972933-teleai/REG/m3lli51t/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20260326%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20260326T043108Z&X-Goog-Expires=86399&X-Goog-Signature=6f5856f5b1dd6e9d2efd01092fa98674135a31cd288253e9013a5ba3024eebbd56748db3da929cb5342625eb9500b726db70e1b5e2434471f98fded30c38a53884c36993b0f6ad68e20df49dc9642e3f3251903b352dd72aad76871227838ccfca1a4097f4bddd24e22f39bdd506334722f15d05be1b2fab112907648c2cad8774ae360e76c816113cdb33b19a2a0fa98a4cf50f21b74fcb4d6e392257126d83935963e0430a4a48a58e2354d47db9580855540e2f2cfceee9b9d8e6d53cc3a72ed15a4b4341a08b25e608511d4961afcfaf8231b8b25e8a7d732c7e9fd007db6fbd2197776041c8fd046bd2a8a638599df98923c5ddb5f7bead505453ccd11f&X-Goog-SignedHeaders=host&X-User=2365972933\": read tcp 172.20.98.27:43200->142.250.199.219:443: read: connection reset by peer"}

New/REG/wandb/run-20260326_123101-m3lli51t/logs/debug.log

@@ -0,0 +1,20 @@
+2026-03-26 12:31:01,272 INFO    MainThread:571602 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-26 12:31:01,272 INFO    MainThread:571602 [wandb_setup.py:_flush():81] Configure stats pid to 571602
+2026-03-26 12:31:01,272 INFO    MainThread:571602 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-26 12:31:01,272 INFO    MainThread:571602 [wandb_init.py:setup_run_log_directory():717] Logging user logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG/wandb/run-20260326_123101-m3lli51t/logs/debug.log
+2026-03-26 12:31:01,273 INFO    MainThread:571602 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG/wandb/run-20260326_123101-m3lli51t/logs/debug-internal.log
+2026-03-26 12:31:01,273 INFO    MainThread:571602 [wandb_init.py:init():844] calling init triggers
+2026-03-26 12:31:01,273 INFO    MainThread:571602 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-26 12:31:01,273 INFO    MainThread:571602 [wandb_init.py:init():892] starting backend
+2026-03-26 12:31:01,576 INFO    MainThread:571602 [wandb_init.py:init():895] sending inform_init request
+2026-03-26 12:31:01,591 INFO    MainThread:571602 [wandb_init.py:init():903] backend started and connected
+2026-03-26 12:31:01,593 INFO    MainThread:571602 [wandb_init.py:init():973] updated telemetry
+2026-03-26 12:31:01,607 INFO    MainThread:571602 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-26 12:31:03,274 INFO    MainThread:571602 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-26 12:31:03,366 INFO    MainThread:571602 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-26 12:31:03,366 INFO    MainThread:571602 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-26 12:31:03,366 INFO    MainThread:571602 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-26 12:31:03,366 INFO    MainThread:571602 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-26 12:31:03,372 INFO    MainThread:571602 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-26 12:31:03,372 INFO    MainThread:571602 [wandb_run.py:_config_callback():1403] config_cb None None {'output_dir': '/your_path/reg_xlarge_dinov2_base_align_8_cls', 'exp_name': 'linear-dinov2-b-enc8', 'logging_dir': 'logs', 'report_to': 'wandb', 'sampling_steps': 10000, 'resume_step': 0, 'model': 'SiT-XL/2', 'num_classes': 1000, 'encoder_depth': 8, 'fused_attn': True, 'qk_norm': False, 'ops_head': 16, 'data_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256', 'semantic_features_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0', 'resolution': 256, 'batch_size': 256, 'allow_tf32': True, 'mixed_precision': 'fp16', 'epochs': 1400, 'max_train_steps': 1000000, 'checkpointing_steps': 10000, 'gradient_accumulation_steps': 1, 'learning_rate': 0.0001, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_weight_decay': 0.0, 'adam_epsilon': 1e-08, 'max_grad_norm': 1.0, 'seed': 0, 'num_workers': 4, 'path_type': 'linear', 'prediction': 'v', 'cfg_prob': 0.1, 'enc_type': 'dinov2-vit-b', 'proj_coeff': 0.5, 'weighting': 'uniform', 'legacy': False, 'cls': 0.03}

New/REG/wandb/run-20260326_130847-0e5vs4f8/files/requirements.txt

@@ -0,0 +1,168 @@
+dill==0.3.8
+mkl-service==2.4.0
+mpmath==1.3.0
+typing_extensions==4.12.2
+urllib3==2.3.0
+torch==2.5.1
+ptyprocess==0.7.0
+traitlets==5.14.3
+pyasn1==0.6.1
+opencv-python-headless==4.12.0.88
+nest-asyncio==1.6.0
+kiwisolver==1.4.8
+click==8.2.1
+fire==0.7.1
+diffusers==0.35.1
+accelerate==1.7.0
+ipykernel==6.29.5
+peft==0.17.1
+attrs==24.3.0
+six==1.17.0
+numpy==2.0.1
+yarl==1.18.0
+huggingface_hub==0.34.4
+Bottleneck==1.4.2
+numexpr==2.11.0
+dataclasses==0.6
+typing-inspection==0.4.1
+safetensors==0.5.3
+pyparsing==3.2.3
+psutil==7.0.0
+imageio==2.37.0
+debugpy==1.8.14
+cycler==0.12.1
+pyasn1_modules==0.4.2
+matplotlib-inline==0.1.7
+matplotlib==3.10.3
+jedi==0.19.2
+tokenizers==0.21.2
+seaborn==0.13.2
+timm==1.0.15
+aiohappyeyeballs==2.6.1
+hf-xet==1.1.8
+multidict==6.1.0
+tqdm==4.67.1
+wheel==0.45.1
+simsimd==6.5.1
+sentencepiece==0.2.1
+grpcio==1.74.0
+asttokens==3.0.0
+absl-py==2.3.1
+stack-data==0.6.3
+pandas==2.3.0
+importlib_metadata==8.7.0
+pytorch-image-generation-metrics==0.6.1
+frozenlist==1.5.0
+MarkupSafe==3.0.2
+setuptools==78.1.1
+multiprocess==0.70.15
+pip==25.1
+requests==2.32.3
+mkl_random==1.2.8
+tensorboard-plugin-wit==1.8.1
+ExifRead-nocycle==3.0.1
+webdataset==0.2.111
+threadpoolctl==3.6.0
+pyarrow==21.0.0
+executing==2.2.0
+decorator==5.2.1
+contourpy==1.3.2
+annotated-types==0.7.0
+scikit-learn==1.7.1
+jupyter_client==8.6.3
+albumentations==1.4.24
+wandb==0.25.0
+certifi==2025.8.3
+idna==3.7
+xxhash==3.5.0
+Jinja2==3.1.6
+python-dateutil==2.9.0.post0
+aiosignal==1.4.0
+triton==3.1.0
+torchvision==0.20.1
+stringzilla==3.12.6
+pure_eval==0.2.3
+braceexpand==0.1.7
+zipp==3.22.0
+oauthlib==3.3.1
+Markdown==3.8.2
+fsspec==2025.3.0
+fonttools==4.58.2
+comm==0.2.2
+ipython==9.3.0
+img2dataset==1.47.0
+networkx==3.4.2
+PySocks==1.7.1
+tzdata==2025.2
+smmap==5.0.2
+mkl_fft==1.3.11
+sentry-sdk==2.29.1
+Pygments==2.19.1
+pexpect==4.9.0
+ftfy==6.3.1
+einops==0.8.1
+requests-oauthlib==2.0.0
+gitdb==4.0.12
+albucore==0.0.23
+torchdiffeq==0.2.5
+GitPython==3.1.44
+bitsandbytes==0.47.0
+pytorch-fid==0.3.0
+clean-fid==0.1.35
+pytorch-gan-metrics==0.5.4
+Brotli==1.0.9
+charset-normalizer==3.3.2
+gmpy2==2.2.1
+pillow==11.1.0
+PyYAML==6.0.2
+tornado==6.5.1
+termcolor==3.1.0
+setproctitle==1.3.6
+scipy==1.15.3
+regex==2024.11.6
+protobuf==6.31.1
+platformdirs==4.3.8
+joblib==1.5.1
+cachetools==4.2.4
+ipython_pygments_lexers==1.1.1
+google-auth==1.35.0
+transformers==4.53.2
+torch-fidelity==0.3.0
+tensorboard==2.4.0
+filelock==3.17.0
+packaging==25.0
+propcache==0.3.1
+pytz==2025.2
+aiohttp==3.11.10
+wcwidth==0.2.13
+clip==0.2.0
+Werkzeug==3.1.3
+tensorboard-data-server==0.6.1
+sympy==1.13.1
+pyzmq==26.4.0
+pydantic_core==2.33.2
+prompt_toolkit==3.0.51
+parso==0.8.4
+docker-pycreds==0.4.0
+rsa==4.9.1
+pydantic==2.11.5
+jupyter_core==5.8.1
+google-auth-oauthlib==0.4.6
+datasets==4.0.0
+torch-tb-profiler==0.4.3
+autocommand==2.2.2
+backports.tarfile==1.2.0
+importlib_metadata==8.0.0
+jaraco.collections==5.1.0
+jaraco.context==5.3.0
+jaraco.functools==4.0.1
+more-itertools==10.3.0
+packaging==24.2
+platformdirs==4.2.2
+typeguard==4.3.0
+inflect==7.3.1
+jaraco.text==3.12.1
+tomli==2.0.1
+typing_extensions==4.12.2
+wheel==0.45.1
+zipp==3.19.2

New/REG/wandb/run-20260326_130847-0e5vs4f8/files/wandb-metadata.json

@@ -0,0 +1,80 @@
+{
+  "os":  "Linux-5.15.0-94-generic-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-03-26T05:08:47.738290Z",
+  "args":  [
+    "--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",
+    "--cls=0.03",
+    "--output-dir=/your_path/reg_xlarge_dinov2_base_align_8_cls",
+    "--exp-name=linear-dinov2-b-enc8",
+    "--batch-size=256",
+    "--data-dir=/gemini/space/zhaozy/dataset/Imagenet/imagenet_256",
+    "--semantic-features-dir=/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0"
+  ],
+  "program":  "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG/train.py",
+  "codePath":  "train.py",
+  "codePathLocal":  "train.py",
+  "git":  {
+    "remote":  "https://github.com/Martinser/REG.git",
+    "commit":  "021ea2e50c38c5803bd9afff16316958a01fbd1d"
+  },
+  "email":  "2365972933@qq.com",
+  "root":  "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/New/REG",
+  "host":  "24c964746905d416ce09d045f9a06f23-taskrole1-0",
+  "executable":  "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/bin/python",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  4,
+  "disk":  {
+    "/":  {
+      "total":  "3838880616448",
+      "used":  "367217483776"
+    }
+  },
+  "memory":  {
+    "total":  "2164115296256"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-757303bb-4ec2-808b-a17f-95f6f5bad6dc"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-a09f2421-99e6-a72e-63bd-fd7452510758"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-9c670cc7-60a8-17f8-9b39-7ced3744976d"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-e6b1d8da-68d7-ed83-90d0-a4dedf33120e"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "zx05al0qjdz8zwayv7jbdzbhr8qwhpnu"
+}