Add files using upload-large-folder tool

REG copy/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.

REG copy/README.md

@@ -0,0 +1,156 @@
+<p align="center">
+  <h1 align="center">Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think (NeurIPS 2025 Oral)
+</h1>
+  <p align="center">
+      <a href='https://github.com/Martinser' style='text-decoration: none' >Ge Wu</a><sup>1</sup>&emsp;
+      <a href='https://github.com/ShenZhang-Shin' style='text-decoration: none' >Shen Zhang</a><sup>3</sup>&emsp;
+      <a href='' style='text-decoration: none' >Ruijing Shi</a><sup>1</sup>&emsp;
+      <a href='https://shgao.site/' style='text-decoration: none' >Shanghua Gao</a><sup>4</sup>&emsp;
+      <a href='https://zhenyuanchenai.github.io/' style='text-decoration: none' >Zhenyuan Chen</a><sup>1</sup>&emsp;
+      <a href='https://scholar.google.com/citations?user=6Z66DAwAAAAJ&hl=en' style='text-decoration: none' >Lei Wang</a><sup>1</sup>&emsp;     
+      <a href='https://www.zhihu.com/people/chen-zhao-wei-16-2' style='text-decoration: none' >Zhaowei Chen</a><sup>3</sup>&emsp;
+      <a href='https://gao-hongcheng.github.io/' style='text-decoration: none' >Hongcheng Gao</a><sup>5</sup>&emsp;
+      <a href='https://scholar.google.com/citations?view_op=list_works&hl=zh-CN&hl=zh-CN&user=0xP6bxcAAAAJ' style='text-decoration: none' >Yao Tang</a><sup>3</sup>&emsp;
+      <a href='https://scholar.google.com/citations?user=6CIDtZQAAAAJ&hl=en' style='text-decoration: none' >Jian Yang</a><sup>1</sup>&emsp;
+      <a href='https://mmcheng.net/cmm/' style='text-decoration: none' >Ming-Ming Cheng</a><sup>1,2</sup>&emsp;
+      <a href='https://implus.github.io/' style='text-decoration: none' >Xiang Li</a><sup>1,2*</sup>&emsp;
+        <p align="center">
+        $^{1}$ VCIP, CS, Nankai University, $^{2}$ NKIARI, Shenzhen Futian, $^{3}$ JIIOV Technology, 
+        $^{4}$ Harvard University, $^{5}$ University of Chinese Academy of Sciences 
+        <p align='center'>
+      <div align="center">
+       <a href='https://arxiv.org/abs/2507.01467v2'><img src='https://img.shields.io/badge/arXiv-2507.01467v2-brown.svg?logo=arxiv&logoColor=white'></a>
+	<a href='https://huggingface.co/Martinser/REG/tree/main'><img src='https://img.shields.io/badge/🤗-Model-blue.svg'></a>
+		  <a href='https://zhuanlan.zhihu.com/p/1952346823168595518'><img src='https://img.shields.io/badge/Zhihu-chinese_article-blue.svg?logo=zhihu&logoColor=white'></a>
+	  </div>
+    <p align='center'>
+    </p>
+   </p>
+</p>
+
+
+## 🚩 Overview 
+
+![overview](fig/reg.png)
+
+REPA and its variants effectively mitigate training challenges in diffusion models by incorporating external visual representations from pretrained models, through alignment between the noisy hidden projections of denoising networks and foundational clean image representations. 
+We argue that the external alignment, which is absent during the entire denoising inference process, falls short of fully harnessing the potential of discriminative representations. 
+
+In this work, we propose a straightforward method called Representation Entanglement for Generation (REG), which entangles low-level image latents with a single high-level class token from pretrained foundation models for denoising. 
+REG acquires the capability to produce coherent image-class pairs directly from pure noise, 
+substantially improving both generation quality and training efficiency. 
+This is accomplished with negligible additional inference overhead, **requiring only one single additional token for denoising (<0.5\% increase in FLOPs and latency).**
+The inference process concurrently reconstructs both image latents and their corresponding global semantics, where the acquired semantic knowledge actively guides and enhances the image generation process.
+
+On ImageNet $256{\times}256$, SiT-XL/2 + REG demonstrates remarkable convergence acceleration, **achieving $\textbf{63}\times$ and $\textbf{23}\times$ faster training than SiT-XL/2 and SiT-XL/2 + REPA, respectively.** 
+More impressively, SiT-L/2 + REG trained for merely 400K iterations outperforms SiT-XL/2 + REPA trained for 4M iterations ($\textbf{10}\times$ longer).
+
+
+
+## 📰 News
+
+- **[2025.08.05]** We have released the pre-trained weights of REG + SiT-XL/2 in 4M (800 epochs).
+
+
+## 📝 Results
+
+- Performance on ImageNet $256{\times}256$ with FID=1.36 by introducing a single class token.
+- $\textbf{63}\times$ and $\textbf{23}\times$ faster training than SiT-XL/2 and SiT-XL/2 + REPA.
+
+<div align="center">
+<img src="fig/img.png" alt="Results">
+</div>
+
+
+## 📋 Plan
+- More training steps on ImageNet 256&512 and T2I.
+
+
+## 👊 Usage
+
+### 1. Environment setup
+
+```bash
+conda create -n reg python=3.10.16 -y
+conda activate reg
+pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1
+pip install -r requirements.txt
+```
+
+### 2. Dataset
+
+#### Dataset download
+
+Currently, we provide experiments for ImageNet. You can place the data that you want and can specifiy it via `--data-dir` arguments in training scripts.
+
+#### Preprocessing data
+Please refer to the preprocessing guide. And you can directly download our processed data, ImageNet data [link](https://huggingface.co/WindATree/ImageNet-256-VAE/tree/main), and ImageNet data after VAE encoder [link]( https://huggingface.co/WindATree/vae-sd/tree/main)
+
+### 3. Training
+Run train.sh
+```bash
+bash train.sh
+```
+
+train.sh contains the following content.
+```bash
+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-B/2" \
+    --enc-type="dinov2-vit-b" \
+    --proj-coeff=0.5 \
+    --encoder-depth=4 \     #SiT-L/XL use 8, SiT-B use 4
+    --output-dir="your_path" \
+    --exp-name="linear-dinov2-b-enc4" \
+    --batch-size=256 \
+    --data-dir="data_path/imagenet_vae" \
+    --cls=0.03
+```
+
+Then this script will automatically create the folder in `exps` to save logs and checkpoints. You can adjust the following options:
+
+- `--models`: `[SiT-B/2, SiT-L/2, SiT-XL/2]`
+- `--enc-type`: `[dinov2-vit-b, clip-vit-L]`
+- `--proj-coeff`: Any values larger than 0
+- `--encoder-depth`: Any values between 1 to the depth of the model
+- `--output-dir`: Any directory that you want to save checkpoints and logs
+- `--exp-name`: Any string name (the folder will be created under `output-dir`)
+- `--cls`: Weight coefficients of REG loss
+
+
+### 4. Generate images and evaluation
+You can generate images and get the final results through the following script. 
+The weight of REG can be found in this [link](https://pan.baidu.com/s/1QX2p3ybh1KfNU7wsp5McWw?pwd=khpp) or [HF](https://huggingface.co/Martinser/REG/tree/main).
+
+```bash
+bash eval.sh
+```
+
+
+## Citation
+If you find our work, this repository, or pretrained models useful, please consider giving a star and citation.
+```
+@article{wu2025representation,
+  title={Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think},
+  author={Wu, Ge and Zhang, Shen and Shi, Ruijing and Gao, Shanghua and Chen, Zhenyuan and Wang, Lei and Chen, Zhaowei and Gao, Hongcheng and Tang, Yao and Yang, Jian and others},
+  journal={arXiv preprint arXiv:2507.01467},
+  year={2025}
+}
+```
+
+## Contact
+If you have any questions, please create an issue on this repository, contact at gewu.nku@gmail.com or wechat(wg1158848).
+
+
+## Acknowledgements
+
+Our code is based on [REPA](https://github.com/sihyun-yu/REPA), along with [SiT](https://github.com/willisma/SiT), [DINOv2](https://github.com/facebookresearch/dinov2), [ADM](https://github.com/openai/guided-diffusion) and [U-ViT](https://github.com/baofff/U-ViT) repositories. We thank the authors for releasing their code. If you use our model and code, please consider citing these works as well.
+
+
+

REG copy/dataset.py

@@ -0,0 +1,80 @@
+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):
+    def __init__(self, data_dir):
+        PIL.Image.init()
+        supported_ext = PIL.Image.EXTENSION.keys() | {'.npy'}
+
+        self.images_dir = os.path.join(data_dir, 'imagenet_256_vae')
+        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):
+        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):
+        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])

REG copy/eval.sh

@@ -0,0 +1,52 @@
+
+random_number=$((RANDOM % 100 + 1200))
+NUM_GPUS=8
+STEP="4000000"
+SAVE_PATH="your_path/reg_xlarge_dinov2_base_align_8_cls/linear-dinov2-b-enc8"
+VAE_PATH="your_vae_path/"
+NUM_STEP=250
+MODEL_SIZE='XL'
+CFG_SCALE=2.3
+CLS_CFG_SCALE=2.3
+GH=0.85
+
+export NCCL_P2P_DISABLE=1
+
+python -m torch.distributed.launch --master_port=$random_number --nproc_per_node=$NUM_GPUS generate.py \
+  --model SiT-XL/2 \
+  --num-fid-samples 50000 \
+  --ckpt ${SAVE_PATH}/checkpoints/${STEP}.pt \
+  --path-type=linear \
+  --encoder-depth=8 \
+  --projector-embed-dims=768 \
+  --per-proc-batch-size=64 \
+  --mode=sde \
+  --num-steps=${NUM_STEP} \
+  --cfg-scale=${CFG_SCALE} \
+  --cls-cfg-scale=${CLS_CFG_SCALE} \
+  --guidance-high=${GH} \
+  --sample-dir ${SAVE_PATH}/checkpoints \
+  --cls=768
+
+
+python ./evaluations/evaluator.py \
+    --ref_batch your_path/VIRTUAL_imagenet256_labeled.npz \
+    --sample_batch ${SAVE_PATH}/checkpoints/SiT-${MODEL_SIZE}-2-${STEP}-size-256-vae-ema-cfg-${CFG_SCALE}-seed-0-sde-${GH}-${CLS_CFG_SCALE}.npz \
+    --save_path ${SAVE_PATH}/checkpoints \
+    --cfg_cond 1 \
+    --step ${STEP} \
+    --num_steps ${NUM_STEP} \
+    --cfg ${CFG_SCALE} \
+    --cls_cfg ${CLS_CFG_SCALE} \
+    --gh ${GH}
+
+
+
+
+
+
+
+
+
+
+

REG copy/generate.py

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

REG copy/loss.py

@@ -0,0 +1,102 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+def mean_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return torch.mean(x, dim=list(range(1, len(x.size()))))
+
+def sum_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return torch.sum(x, dim=list(range(1, len(x.size()))))
+
+class SILoss:
+    def __init__(
+            self,
+            prediction='v',
+            path_type="linear",
+            weighting="uniform",
+            encoders=[], 
+            accelerator=None, 
+            latents_scale=None, 
+            latents_bias=None,
+            ):
+        self.prediction = prediction
+        self.weighting = weighting
+        self.path_type = path_type
+        self.encoders = encoders
+        self.accelerator = accelerator
+        self.latents_scale = latents_scale
+        self.latents_bias = latents_bias
+
+    def interpolant(self, t):
+        if self.path_type == "linear":
+            alpha_t = 1 - t
+            sigma_t = t
+            d_alpha_t = -1
+            d_sigma_t =  1
+        elif self.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()
+
+        return alpha_t, sigma_t, d_alpha_t, d_sigma_t
+
+    def __call__(self, model, images, model_kwargs=None, zs=None, cls_token=None,
+                 time_input=None, noises=None,):
+        if model_kwargs == None:
+            model_kwargs = {}
+        # sample timesteps
+        if time_input is None:
+            if self.weighting == "uniform":
+                time_input = torch.rand((images.shape[0], 1, 1, 1))
+            elif self.weighting == "lognormal":
+                # sample timestep according to log-normal distribution of sigmas following EDM
+                rnd_normal = torch.randn((images.shape[0], 1 ,1, 1))
+                sigma = rnd_normal.exp()
+                if self.path_type == "linear":
+                    time_input = sigma / (1 + sigma)
+                elif self.path_type == "cosine":
+                    time_input = 2 / np.pi * torch.atan(sigma)
+                
+        time_input = time_input.to(device=images.device, dtype=images.dtype)
+
+        if noises is None:
+            noises = torch.randn_like(images)
+            noises_cls = torch.randn_like(cls_token)
+
+        alpha_t, sigma_t, d_alpha_t, d_sigma_t = self.interpolant(time_input)
+
+        model_input = alpha_t * images + sigma_t * noises
+        cls_input = alpha_t.squeeze(-1).squeeze(-1) * cls_token + sigma_t.squeeze(-1).squeeze(-1) * noises_cls
+        if self.prediction == 'v':
+            model_target = d_alpha_t * images + d_sigma_t * noises
+            cls_target = d_alpha_t * cls_token + d_sigma_t * noises_cls
+        else:
+            raise NotImplementedError()
+
+        model_output, zs_tilde, cls_output = model(model_input, time_input.flatten(), **model_kwargs,
+                                                    cls_token=cls_input)
+
+        #denoising_loss
+        denoising_loss = mean_flat((model_output - model_target) ** 2)
+        denoising_loss_cls = mean_flat((cls_output - cls_target) ** 2)
+
+        # projection loss
+        proj_loss = 0.
+        bsz = zs[0].shape[0]
+        for i, (z, z_tilde) in enumerate(zip(zs, zs_tilde)):
+            for j, (z_j, z_tilde_j) in enumerate(zip(z, z_tilde)):
+                z_tilde_j = torch.nn.functional.normalize(z_tilde_j, dim=-1) 
+                z_j = torch.nn.functional.normalize(z_j, dim=-1) 
+                proj_loss += mean_flat(-(z_j * z_tilde_j).sum(dim=-1))
+        proj_loss /= (len(zs) * bsz)
+
+        return denoising_loss, proj_loss, time_input, noises, denoising_loss_cls

REG copy/requirements.txt

@@ -0,0 +1,97 @@
+  - pip:
+  absl-py==2.2.2
+  accelerate==1.2.1
+  aiohappyeyeballs==2.6.1
+  aiohttp==3.11.16
+  aiosignal==1.3.2
+  astunparse==1.6.3
+  async-timeout==5.0.1
+  attrs==25.3.0
+  certifi==2022.12.7
+  charset-normalizer==2.1.1
+  click==8.1.8
+  datasets==2.20.0
+  diffusers==0.32.1
+  dill==0.3.8
+  docker-pycreds==0.4.0
+  einops==0.8.1
+  filelock==3.13.1
+  flatbuffers==25.2.10
+  frozenlist==1.5.0
+  fsspec==2024.5.0
+  ftfy==6.3.1
+  gast==0.6.0
+  gitdb==4.0.12
+  gitpython==3.1.44
+  google-pasta==0.2.0
+  grpcio==1.71.0
+  h5py==3.13.0
+  huggingface-hub==0.27.1
+  idna==3.4
+  importlib-metadata==8.6.1
+  jinja2==3.1.4
+  joblib==1.4.2
+  keras==3.9.2
+  libclang==18.1.1
+  markdown==3.8
+  markdown-it-py==3.0.0
+  markupsafe==2.1.5
+  mdurl==0.1.2
+  ml-dtypes==0.3.2
+  mpmath==1.3.0
+  multidict==6.4.3
+  multiprocess==0.70.16
+  namex==0.0.8
+  networkx==3.3
+  numpy==1.26.4
+  opt-einsum==3.4.0
+  optree==0.15.0
+  packaging==24.2
+  pandas==2.2.3
+  pillow==11.0.0
+  platformdirs==4.3.7
+  propcache==0.3.1
+  protobuf==4.25.6
+  psutil==7.0.0
+  pyarrow==19.0.1
+  pyarrow-hotfix==0.6
+  pygments==2.19.1
+  python-dateutil==2.9.0.post0
+  pytz==2025.2
+  pyyaml==6.0.2
+  regex==2024.11.6
+  requests==2.32.3
+  rich==14.0.0
+  safetensors==0.5.3
+  scikit-learn==1.5.1
+  scipy==1.15.2
+  sentry-sdk==2.26.1
+  setproctitle==1.3.5
+  six==1.17.0
+  smmap==5.0.2
+  sympy==1.13.1
+  tensorboard==2.16.1
+  tensorboard-data-server==0.7.2
+  tensorflow==2.16.1
+  tensorflow-io-gcs-filesystem==0.37.1
+  termcolor==3.0.1
+  tf-keras==2.16.0
+  threadpoolctl==3.6.0
+  timm==1.0.12
+  tokenizers==0.21.0
+  tqdm==4.67.1
+  transformers==4.47.0
+  triton==2.1.0
+  typing-extensions==4.12.2
+  tzdata==2025.2
+  urllib3==1.26.13
+  wandb==0.17.6
+  wcwidth==0.2.13
+  werkzeug==3.1.3
+  wrapt==1.17.2
+  xformer==1.0.1
+  xformers==0.0.23
+  xxhash==3.5.0
+  yarl==1.20.0
+  zipp==3.21.0
+

REG copy/samplers.py

@@ -0,0 +1,169 @@
+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
+        ):
+    # 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
+
+
+    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
+
+            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
+
+    # 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
+
+    return mean_x

REG copy/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

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.

REG/README.md

@@ -0,0 +1,156 @@
+<p align="center">
+  <h1 align="center">Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think (NeurIPS 2025 Oral)
+</h1>
+  <p align="center">
+      <a href='https://github.com/Martinser' style='text-decoration: none' >Ge Wu</a><sup>1</sup>&emsp;
+      <a href='https://github.com/ShenZhang-Shin' style='text-decoration: none' >Shen Zhang</a><sup>3</sup>&emsp;
+      <a href='' style='text-decoration: none' >Ruijing Shi</a><sup>1</sup>&emsp;
+      <a href='https://shgao.site/' style='text-decoration: none' >Shanghua Gao</a><sup>4</sup>&emsp;
+      <a href='https://zhenyuanchenai.github.io/' style='text-decoration: none' >Zhenyuan Chen</a><sup>1</sup>&emsp;
+      <a href='https://scholar.google.com/citations?user=6Z66DAwAAAAJ&hl=en' style='text-decoration: none' >Lei Wang</a><sup>1</sup>&emsp;     
+      <a href='https://www.zhihu.com/people/chen-zhao-wei-16-2' style='text-decoration: none' >Zhaowei Chen</a><sup>3</sup>&emsp;
+      <a href='https://gao-hongcheng.github.io/' style='text-decoration: none' >Hongcheng Gao</a><sup>5</sup>&emsp;
+      <a href='https://scholar.google.com/citations?view_op=list_works&hl=zh-CN&hl=zh-CN&user=0xP6bxcAAAAJ' style='text-decoration: none' >Yao Tang</a><sup>3</sup>&emsp;
+      <a href='https://scholar.google.com/citations?user=6CIDtZQAAAAJ&hl=en' style='text-decoration: none' >Jian Yang</a><sup>1</sup>&emsp;
+      <a href='https://mmcheng.net/cmm/' style='text-decoration: none' >Ming-Ming Cheng</a><sup>1,2</sup>&emsp;
+      <a href='https://implus.github.io/' style='text-decoration: none' >Xiang Li</a><sup>1,2*</sup>&emsp;
+        <p align="center">
+        $^{1}$ VCIP, CS, Nankai University, $^{2}$ NKIARI, Shenzhen Futian, $^{3}$ JIIOV Technology, 
+        $^{4}$ Harvard University, $^{5}$ University of Chinese Academy of Sciences 
+        <p align='center'>
+      <div align="center">
+       <a href='https://arxiv.org/abs/2507.01467v2'><img src='https://img.shields.io/badge/arXiv-2507.01467v2-brown.svg?logo=arxiv&logoColor=white'></a>
+	<a href='https://huggingface.co/Martinser/REG/tree/main'><img src='https://img.shields.io/badge/🤗-Model-blue.svg'></a>
+		  <a href='https://zhuanlan.zhihu.com/p/1952346823168595518'><img src='https://img.shields.io/badge/Zhihu-chinese_article-blue.svg?logo=zhihu&logoColor=white'></a>
+	  </div>
+    <p align='center'>
+    </p>
+   </p>
+</p>
+
+
+## 🚩 Overview 
+
+![overview](fig/reg.png)
+
+REPA and its variants effectively mitigate training challenges in diffusion models by incorporating external visual representations from pretrained models, through alignment between the noisy hidden projections of denoising networks and foundational clean image representations. 
+We argue that the external alignment, which is absent during the entire denoising inference process, falls short of fully harnessing the potential of discriminative representations. 
+
+In this work, we propose a straightforward method called Representation Entanglement for Generation (REG), which entangles low-level image latents with a single high-level class token from pretrained foundation models for denoising. 
+REG acquires the capability to produce coherent image-class pairs directly from pure noise, 
+substantially improving both generation quality and training efficiency. 
+This is accomplished with negligible additional inference overhead, **requiring only one single additional token for denoising (<0.5\% increase in FLOPs and latency).**
+The inference process concurrently reconstructs both image latents and their corresponding global semantics, where the acquired semantic knowledge actively guides and enhances the image generation process.
+
+On ImageNet $256{\times}256$, SiT-XL/2 + REG demonstrates remarkable convergence acceleration, **achieving $\textbf{63}\times$ and $\textbf{23}\times$ faster training than SiT-XL/2 and SiT-XL/2 + REPA, respectively.** 
+More impressively, SiT-L/2 + REG trained for merely 400K iterations outperforms SiT-XL/2 + REPA trained for 4M iterations ($\textbf{10}\times$ longer).
+
+
+
+## 📰 News
+
+- **[2025.08.05]** We have released the pre-trained weights of REG + SiT-XL/2 in 4M (800 epochs).
+
+
+## 📝 Results
+
+- Performance on ImageNet $256{\times}256$ with FID=1.36 by introducing a single class token.
+- $\textbf{63}\times$ and $\textbf{23}\times$ faster training than SiT-XL/2 and SiT-XL/2 + REPA.
+
+<div align="center">
+<img src="fig/img.png" alt="Results">
+</div>
+
+
+## 📋 Plan
+- More training steps on ImageNet 256&512 and T2I.
+
+
+## 👊 Usage
+
+### 1. Environment setup
+
+```bash
+conda create -n reg python=3.10.16 -y
+conda activate reg
+pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1
+pip install -r requirements.txt
+```
+
+### 2. Dataset
+
+#### Dataset download
+
+Currently, we provide experiments for ImageNet. You can place the data that you want and can specifiy it via `--data-dir` arguments in training scripts.
+
+#### Preprocessing data
+Please refer to the preprocessing guide. And you can directly download our processed data, ImageNet data [link](https://huggingface.co/WindATree/ImageNet-256-VAE/tree/main), and ImageNet data after VAE encoder [link]( https://huggingface.co/WindATree/vae-sd/tree/main)
+
+### 3. Training
+Run train.sh
+```bash
+bash train.sh
+```
+
+train.sh contains the following content.
+```bash
+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-B/2" \
+    --enc-type="dinov2-vit-b" \
+    --proj-coeff=0.5 \
+    --encoder-depth=4 \     #SiT-L/XL use 8, SiT-B use 4
+    --output-dir="your_path" \
+    --exp-name="linear-dinov2-b-enc4" \
+    --batch-size=256 \
+    --data-dir="data_path/imagenet_vae" \
+    --cls=0.03
+```
+
+Then this script will automatically create the folder in `exps` to save logs and checkpoints. You can adjust the following options:
+
+- `--models`: `[SiT-B/2, SiT-L/2, SiT-XL/2]`
+- `--enc-type`: `[dinov2-vit-b, clip-vit-L]`
+- `--proj-coeff`: Any values larger than 0
+- `--encoder-depth`: Any values between 1 to the depth of the model
+- `--output-dir`: Any directory that you want to save checkpoints and logs
+- `--exp-name`: Any string name (the folder will be created under `output-dir`)
+- `--cls`: Weight coefficients of REG loss
+
+
+### 4. Generate images and evaluation
+You can generate images and get the final results through the following script. 
+The weight of REG can be found in this [link](https://pan.baidu.com/s/1QX2p3ybh1KfNU7wsp5McWw?pwd=khpp) or [HF](https://huggingface.co/Martinser/REG/tree/main).
+
+```bash
+bash eval.sh
+```
+
+
+## Citation
+If you find our work, this repository, or pretrained models useful, please consider giving a star and citation.
+```
+@article{wu2025representation,
+  title={Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think},
+  author={Wu, Ge and Zhang, Shen and Shi, Ruijing and Gao, Shanghua and Chen, Zhenyuan and Wang, Lei and Chen, Zhaowei and Gao, Hongcheng and Tang, Yao and Yang, Jian and others},
+  journal={arXiv preprint arXiv:2507.01467},
+  year={2025}
+}
+```
+
+## Contact
+If you have any questions, please create an issue on this repository, contact at gewu.nku@gmail.com or wechat(wg1158848).
+
+
+## Acknowledgements
+
+Our code is based on [REPA](https://github.com/sihyun-yu/REPA), along with [SiT](https://github.com/willisma/SiT), [DINOv2](https://github.com/facebookresearch/dinov2), [ADM](https://github.com/openai/guided-diffusion) and [U-ViT](https://github.com/baofff/U-ViT) repositories. We thank the authors for releasing their code. If you use our model and code, please consider citing these works as well.
+
+
+

REG/dataset.py

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

REG/eval.sh

@@ -0,0 +1,52 @@
+
+random_number=$((RANDOM % 100 + 1200))
+NUM_GPUS=8
+STEP="4000000"
+SAVE_PATH="your_path/reg_xlarge_dinov2_base_align_8_cls/linear-dinov2-b-enc8"
+VAE_PATH="your_vae_path/"
+NUM_STEP=250
+MODEL_SIZE='XL'
+CFG_SCALE=2.3
+CLS_CFG_SCALE=2.3
+GH=0.85
+
+export NCCL_P2P_DISABLE=1
+
+python -m torch.distributed.launch --master_port=$random_number --nproc_per_node=$NUM_GPUS generate.py \
+  --model SiT-XL/2 \
+  --num-fid-samples 50000 \
+  --ckpt ${SAVE_PATH}/checkpoints/${STEP}.pt \
+  --path-type=linear \
+  --encoder-depth=8 \
+  --projector-embed-dims=768 \
+  --per-proc-batch-size=64 \
+  --mode=sde \
+  --num-steps=${NUM_STEP} \
+  --cfg-scale=${CFG_SCALE} \
+  --cls-cfg-scale=${CLS_CFG_SCALE} \
+  --guidance-high=${GH} \
+  --sample-dir ${SAVE_PATH}/checkpoints \
+  --cls=768
+
+
+python ./evaluations/evaluator.py \
+    --ref_batch your_path/VIRTUAL_imagenet256_labeled.npz \
+    --sample_batch ${SAVE_PATH}/checkpoints/SiT-${MODEL_SIZE}-2-${STEP}-size-256-vae-ema-cfg-${CFG_SCALE}-seed-0-sde-${GH}-${CLS_CFG_SCALE}.npz \
+    --save_path ${SAVE_PATH}/checkpoints \
+    --cfg_cond 1 \
+    --step ${STEP} \
+    --num_steps ${NUM_STEP} \
+    --cfg ${CFG_SCALE} \
+    --cls_cfg ${CLS_CFG_SCALE} \
+    --gh ${GH}
+
+
+
+
+
+
+
+
+
+
+

REG/eval_custom_0.25.log

@@ -0,0 +1 @@
+python: can't open file 'fid_custom.py': [Errno 2] No such file or directory

REG/generate.py

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

REG/loss.py

@@ -0,0 +1,193 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+try:
+    from scipy.optimize import linear_sum_assignment
+except ImportError:
+    linear_sum_assignment = None
+
+
+def ot_pair_noise_to_cls(noise_cls, cls_gt):
+    """
+    Minibatch OT（与 conditional-flow-matching / torchcfm 中 sample_plan_with_scipy 一致）：
+    在 batch 内用平方欧氏代价重排 noise，使 noise_ot[i] 与 cls_gt[i] 构成近似最优传输配对。
+    noise_cls, cls_gt: (N, D) 或任意可在最后一维展平为 D 的形状。
+    """
+    n = noise_cls.shape[0]
+    if n <= 1:
+        return noise_cls, cls_gt
+    if linear_sum_assignment is None:
+        return noise_cls, cls_gt
+    x0 = noise_cls.detach().float().reshape(n, -1)
+    x1 = cls_gt.detach().float().reshape(n, -1)
+    M = torch.cdist(x0, x1) ** 2
+    _, j = linear_sum_assignment(M.cpu().numpy())
+    j = torch.as_tensor(j, device=noise_cls.device, dtype=torch.long)
+    return noise_cls[j], cls_gt
+
+
+def mean_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return torch.mean(x, dim=list(range(1, len(x.size()))))
+
+def sum_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return torch.sum(x, dim=list(range(1, len(x.size()))))
+
+class SILoss:
+    def __init__(
+            self,
+            prediction='v',
+            path_type="linear",
+            weighting="uniform",
+            encoders=[],
+            accelerator=None,
+            latents_scale=None,
+            latents_bias=None,
+            t_c=0.5,
+            ot_cls=True,
+            ):
+        self.prediction = prediction
+        self.weighting = weighting
+        self.path_type = path_type
+        self.encoders = encoders
+        self.accelerator = accelerator
+        self.latents_scale = latents_scale
+        self.latents_bias = latents_bias
+        # t 与 train.py / JsFlow 一致：t=0 为干净 latent，t=1 为纯噪声。
+        # t ∈ (t_c, 1]：语义 cls 沿 OT 配对后的路径从噪声演化为 cls_gt（生成语义通道）；
+        # t ∈ [0, t_c]：cls 恒为真实 cls_gt，目标速度为 0（通道不再插值）。
+        tc = float(t_c)
+        self.t_c = min(max(tc, 1e-4), 1.0 - 1e-4)
+        self.ot_cls = bool(ot_cls)
+
+    def interpolant(self, t):
+        if self.path_type == "linear":
+            alpha_t = 1 - t
+            sigma_t = t
+            d_alpha_t = -1
+            d_sigma_t =  1
+        elif self.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()
+
+        return alpha_t, sigma_t, d_alpha_t, d_sigma_t
+
+    def __call__(self, model, images, model_kwargs=None, zs=None, cls_token=None,
+                 time_input=None, noises=None,):
+        if model_kwargs == None:
+            model_kwargs = {}
+        # sample timesteps
+        if time_input is None:
+            if self.weighting == "uniform":
+                time_input = torch.rand((images.shape[0], 1, 1, 1))
+            elif self.weighting == "lognormal":
+                # sample timestep according to log-normal distribution of sigmas following EDM
+                rnd_normal = torch.randn((images.shape[0], 1 ,1, 1))
+                sigma = rnd_normal.exp()
+                if self.path_type == "linear":
+                    time_input = sigma / (1 + sigma)
+                elif self.path_type == "cosine":
+                    time_input = 2 / np.pi * torch.atan(sigma)
+
+        time_input = time_input.to(device=images.device, dtype=torch.float32)
+        cls_token = cls_token.to(device=images.device, dtype=torch.float32)
+
+        if noises is None:
+            noises = torch.randn_like(images)
+            noises_cls = torch.randn_like(cls_token)
+        else:
+            if isinstance(noises, (tuple, list)) and len(noises) == 2:
+                noises, noises_cls = noises
+            else:
+                noises_cls = torch.randn_like(cls_token)
+
+        alpha_t, sigma_t, d_alpha_t, d_sigma_t = self.interpolant(time_input)
+
+        model_input = alpha_t * images + sigma_t * noises
+        if self.prediction == 'v':
+            model_target = d_alpha_t * images + d_sigma_t * noises
+        else:
+            raise NotImplementedError()
+
+        N = images.shape[0]
+        t_flat = time_input.view(-1).float()
+        high_noise_mask = (t_flat > self.t_c).float().view(N, *([1] * (cls_token.dim() - 1)))
+        low_noise_mask = 1.0 - high_noise_mask
+
+        noise_cls_raw = noises_cls
+        if self.ot_cls:
+            noise_cls_paired, cls_gt_paired = ot_pair_noise_to_cls(noise_cls_raw, cls_token)
+        else:
+            noise_cls_paired, cls_gt_paired = noise_cls_raw, cls_token
+
+        tau_shape = (N,) + (1,) * max(0, cls_token.dim() - 1)
+        tau = (time_input.reshape(tau_shape) - self.t_c) / (1.0 - self.t_c + 1e-8)
+        tau = torch.clamp(tau, 0.0, 1.0)
+        alpha_sem = 1.0 - tau
+        sigma_sem = tau
+
+        cls_t_high = alpha_sem * cls_gt_paired + sigma_sem * noise_cls_paired
+        cls_t = high_noise_mask * cls_t_high + low_noise_mask * cls_token
+        cls_t = torch.nan_to_num(cls_t, nan=0.0, posinf=1e4, neginf=-1e4)
+        cls_t = torch.clamp(cls_t, -1e4, 1e4)
+
+        cls_for_model = cls_t * high_noise_mask + cls_t.detach() * low_noise_mask
+
+        inv_scale = 1.0 / (1.0 - self.t_c + 1e-8)
+        v_cls_high = (noise_cls_paired - cls_gt_paired) * inv_scale
+        v_cls_target = high_noise_mask * v_cls_high
+
+        model_output, zs_tilde, cls_output = model(
+            model_input, time_input.flatten(), **model_kwargs, cls_token=cls_for_model
+        )
+
+        #denoising_loss
+        denoising_loss = mean_flat((model_output - model_target) ** 2)
+        denoising_loss_cls = mean_flat((cls_output - v_cls_target) ** 2)
+
+        # projection loss
+        proj_loss = 0.
+        bsz = zs[0].shape[0]
+        for i, (z, z_tilde) in enumerate(zip(zs, zs_tilde)):
+            for j, (z_j, z_tilde_j) in enumerate(zip(z, z_tilde)):
+                z_tilde_j = torch.nn.functional.normalize(z_tilde_j, dim=-1) 
+                z_j = torch.nn.functional.normalize(z_j, dim=-1) 
+                proj_loss += mean_flat(-(z_j * z_tilde_j).sum(dim=-1))
+        proj_loss /= (len(zs) * bsz)
+
+        return denoising_loss, proj_loss, time_input, noises, denoising_loss_cls
+
+    def tc_velocity_loss(self, model, images, model_kwargs=None, cls_token=None, noises=None):
+        """
+        额外约束：在 t=t_c 处直接监督图像 velocity 场，增强单步（t_c -> 0）稳定性。
+        仅作用于图像分支，不改变原有 cls/projection 主损失定义。
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if cls_token is None:
+            raise ValueError("tc_velocity_loss requires cls_token")
+        if noises is None:
+            noises = torch.randn_like(images)
+
+        bsz = images.shape[0]
+        time_input = torch.full(
+            (bsz, 1, 1, 1), float(self.t_c), device=images.device, dtype=torch.float32
+        )
+        alpha_t, sigma_t, d_alpha_t, d_sigma_t = self.interpolant(time_input)
+        model_input = alpha_t * images + sigma_t * noises
+        model_target = d_alpha_t * images + d_sigma_t * noises
+
+        model_output, _, _ = model(
+            model_input, time_input.flatten(), **model_kwargs, cls_token=cls_token
+        )
+        return mean_flat((model_output - model_target) ** 2)

REG/requirements.txt

@@ -0,0 +1,97 @@
+  - pip:
+  absl-py==2.2.2
+  accelerate==1.2.1
+  aiohappyeyeballs==2.6.1
+  aiohttp==3.11.16
+  aiosignal==1.3.2
+  astunparse==1.6.3
+  async-timeout==5.0.1
+  attrs==25.3.0
+  certifi==2022.12.7
+  charset-normalizer==2.1.1
+  click==8.1.8
+  datasets==2.20.0
+  diffusers==0.32.1
+  dill==0.3.8
+  docker-pycreds==0.4.0
+  einops==0.8.1
+  filelock==3.13.1
+  flatbuffers==25.2.10
+  frozenlist==1.5.0
+  fsspec==2024.5.0
+  ftfy==6.3.1
+  gast==0.6.0
+  gitdb==4.0.12
+  gitpython==3.1.44
+  google-pasta==0.2.0
+  grpcio==1.71.0
+  h5py==3.13.0
+  huggingface-hub==0.27.1
+  idna==3.4
+  importlib-metadata==8.6.1
+  jinja2==3.1.4
+  joblib==1.4.2
+  keras==3.9.2
+  libclang==18.1.1
+  markdown==3.8
+  markdown-it-py==3.0.0
+  markupsafe==2.1.5
+  mdurl==0.1.2
+  ml-dtypes==0.3.2
+  mpmath==1.3.0
+  multidict==6.4.3
+  multiprocess==0.70.16
+  namex==0.0.8
+  networkx==3.3
+  numpy==1.26.4
+  opt-einsum==3.4.0
+  optree==0.15.0
+  packaging==24.2
+  pandas==2.2.3
+  pillow==11.0.0
+  platformdirs==4.3.7
+  propcache==0.3.1
+  protobuf==4.25.6
+  psutil==7.0.0
+  pyarrow==19.0.1
+  pyarrow-hotfix==0.6
+  pygments==2.19.1
+  python-dateutil==2.9.0.post0
+  pytz==2025.2
+  pyyaml==6.0.2
+  regex==2024.11.6
+  requests==2.32.3
+  rich==14.0.0
+  safetensors==0.5.3
+  scikit-learn==1.5.1
+  scipy==1.15.2
+  sentry-sdk==2.26.1
+  setproctitle==1.3.5
+  six==1.17.0
+  smmap==5.0.2
+  sympy==1.13.1
+  tensorboard==2.16.1
+  tensorboard-data-server==0.7.2
+  tensorflow==2.16.1
+  tensorflow-io-gcs-filesystem==0.37.1
+  termcolor==3.0.1
+  tf-keras==2.16.0
+  threadpoolctl==3.6.0
+  timm==1.0.12
+  tokenizers==0.21.0
+  tqdm==4.67.1
+  transformers==4.47.0
+  triton==2.1.0
+  typing-extensions==4.12.2
+  tzdata==2025.2
+  urllib3==1.26.13
+  wandb==0.17.6
+  wcwidth==0.2.13
+  werkzeug==3.1.3
+  wrapt==1.17.2
+  xformer==1.0.1
+  xformers==0.0.23
+  xxhash==3.5.0
+  yarl==1.20.0
+  zipp==3.21.0
+

REG/sample_from_checkpoint.py

@@ -0,0 +1,611 @@
+#!/usr/bin/env python3
+"""
+从 REG/train.py 保存的检查点加载权重，在指定目录生成若干 PNG。
+
+示例：
+  python sample_from_checkpoint.py \\
+    --ckpt exps/jsflow-experiment/checkpoints/0050000.pt \\
+    --out-dir ./samples_gen \\
+    --num-images 64 \\
+    --batch-size 8
+
+  # 按训练 t_c 分段分配步数（t=1→t_c 与 t_c→0；--t-c 可省略若检查点含 t_c）：
+  python sample_from_checkpoint.py ... \\
+    --steps-before-tc 150 --steps-after-tc 100 --t-c 0.5
+
+  # 同一批初始噪声连跑两种 t_c 后段步数（输出到 out-dir 下子目录）：
+  python sample_from_checkpoint.py ... \\
+    --steps-before-tc 150 --steps-after-tc 5 --dual-compare-after
+  # 分段时会在 at_tc/（或 at_tc/after_input、at_tc/after_equal_before）额外保存 t≈t_c 的解码图。
+
+检查点需包含 train.py 写入的键：ema（或 model）、args（推荐，用于自动还原结构）。
+若缺少 args，需通过命令行显式传入 --model、--resolution、--enc-type 等。
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import types
+import numpy as np
+import torch
+from diffusers.models import AutoencoderKL
+from PIL import Image
+from tqdm import tqdm
+
+from models.sit import SiT_models
+from samplers import (
+    euler_maruyama_image_noise_before_tc_sampler,
+    euler_maruyama_image_noise_sampler,
+    euler_maruyama_sampler,
+    euler_ode_sampler,
+)
+
+
+def semantic_dim_from_enc_type(enc_type):
+    """与 train.py 一致：按 enc_type 推断语义/class token 维度。"""
+    if enc_type is None:
+        return 768
+    s = str(enc_type).lower()
+    if "vit-g" in s or "vitg" in s:
+        return 1536
+    if "vit-l" in s or "vitl" in s:
+        return 1024
+    if "vit-s" in s or "vits" in s:
+        return 384
+    return 768
+
+
+def load_train_args_from_ckpt(ckpt: dict) -> argparse.Namespace | None:
+    a = ckpt.get("args")
+    if a is None:
+        return None
+    if isinstance(a, argparse.Namespace):
+        return a
+    if isinstance(a, dict):
+        return argparse.Namespace(**a)
+    if isinstance(a, types.SimpleNamespace):
+        return argparse.Namespace(**vars(a))
+    return None
+
+
+def load_vae(device: torch.device):
+    """与 train.py 相同策略：优先本地 diffusers 缓存中的 sd-vae-ft-mse。"""
+    try:
+        from preprocessing import dnnlib
+
+        cache_dir = dnnlib.make_cache_dir_path("diffusers")
+        os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
+        os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+        os.environ["HF_HOME"] = cache_dir
+        try:
+            vae = AutoencoderKL.from_pretrained(
+                "stabilityai/sd-vae-ft-mse",
+                cache_dir=cache_dir,
+                local_files_only=True,
+            ).to(device)
+            vae.eval()
+            print(f"Loaded VAE from local cache: {cache_dir}")
+            return vae
+        except Exception:
+            pass
+        candidate_dir = None
+        for root_dir in [
+            cache_dir,
+            os.path.join(os.path.expanduser("~"), ".cache", "dnnlib", "diffusers"),
+            os.path.join(os.path.expanduser("~"), ".cache", "diffusers"),
+            os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub"),
+        ]:
+            if not os.path.isdir(root_dir):
+                continue
+            for root, _, files in os.walk(root_dir):
+                if "config.json" in files and "sd-vae-ft-mse" in root.replace("\\", "/"):
+                    candidate_dir = root
+                    break
+            if candidate_dir is not None:
+                break
+        if candidate_dir is not None:
+            vae = AutoencoderKL.from_pretrained(candidate_dir, local_files_only=True).to(device)
+            vae.eval()
+            print(f"Loaded VAE from {candidate_dir}")
+            return vae
+    except Exception as e:
+        print(f"VAE local cache search failed: {e}", file=sys.stderr)
+    try:
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device)
+        vae.eval()
+        print("Loaded VAE from Hub: stabilityai/sd-vae-ft-mse")
+        return vae
+    except Exception as e:
+        raise RuntimeError(
+            "无法加载 VAE stabilityai/sd-vae-ft-mse，请确认已下载或网络可用。"
+        ) from e
+
+
+def build_model_from_train_args(ta: argparse.Namespace, device: torch.device):
+    res = int(getattr(ta, "resolution", 256))
+    latent_size = res // 8
+    enc_type = getattr(ta, "enc_type", "dinov2-vit-b")
+    z_dims = [semantic_dim_from_enc_type(enc_type)]
+    block_kwargs = {
+        "fused_attn": getattr(ta, "fused_attn", True),
+        "qk_norm": getattr(ta, "qk_norm", False),
+    }
+    cfg_prob = float(getattr(ta, "cfg_prob", 0.1))
+    if ta.model not in SiT_models:
+        raise ValueError(f"未知 model={ta.model!r}，可选：{list(SiT_models.keys())}")
+    model = SiT_models[ta.model](
+        input_size=latent_size,
+        num_classes=int(getattr(ta, "num_classes", 1000)),
+        use_cfg=(cfg_prob > 0),
+        z_dims=z_dims,
+        encoder_depth=int(getattr(ta, "encoder_depth", 8)),
+        **block_kwargs,
+    ).to(device)
+    return model, z_dims[0]
+
+
+def resolve_tc_schedule(cli, ta):
+    """
+    若同时给出 --steps-before-tc 与 --steps-after-tc：在 t_c 处分段（--t-c 缺省则用检查点 args.t_c）。
+    否则使用均匀 --num-steps（与旧版一致）。
+    """
+    sb = cli.steps_before_tc
+    sa = cli.steps_after_tc
+    tc = cli.t_c
+    if sb is None and sa is None:
+        return None, None, None
+    if sb is None or sa is None:
+        print(
+            "使用分段步数时必须同时指定 --steps-before-tc 与 --steps-after-tc。",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    if tc is None:
+        tc = getattr(ta, "t_c", None) if ta is not None else None
+    if tc is None:
+        print(
+            "分段采样需要 --t-c，或检查点 args 中含 t_c。",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    return float(tc), int(sb), int(sa)
+
+
+def parse_cli():
+    p = argparse.ArgumentParser(description="REG 检查点采样出图（可选 ODE/EM/EM-图像噪声）")
+    p.add_argument("--ckpt", type=str, required=True, help="train.py 保存的 .pt 路径")
+    p.add_argument("--out-dir", type=str, required=True, help="输出 PNG 目录（会创建）")
+    p.add_argument("--num-images", type=int, required=True, help="生成图片总数")
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument(
+        "--weights",
+        type=str,
+        choices=("ema", "model"),
+        default="ema",
+        help="使用检查点中的 ema 或 model 权重",
+    )
+    p.add_argument("--device", type=str, default="cuda", help="如 cuda 或 cuda:0")
+    p.add_argument(
+        "--num-steps",
+        type=int,
+        default=50,
+        help="均匀时间网格时的欧拉步数（未使用 --steps-before-tc/--steps-after-tc 时生效）",
+    )
+    p.add_argument(
+        "--t-c",
+        type=float,
+        default=None,
+        help="分段时刻：t∈(t_c,1] 与 t∈[0,t_c] 两段；缺省可用检查点 args.t_c（需配合两段步数）",
+    )
+    p.add_argument(
+        "--steps-before-tc",
+        type=int,
+        default=None,
+        help="从 t=1 积分到 t=t_c 的步数（与 --steps-after-tc 成对使用）",
+    )
+    p.add_argument(
+        "--steps-after-tc",
+        type=int,
+        default=None,
+        help="从 t=t_c 积分到 t=0（经 t_floor=0.04）的步数",
+    )
+    p.add_argument("--cfg-scale", type=float, default=1.0)
+    p.add_argument("--cls-cfg-scale", type=float, default=0.0, help="cls 分支 CFG（>0 时需 cfg-scale>1）")
+    p.add_argument("--guidance-low", type=float, default=0.0)
+    p.add_argument("--guidance-high", type=float, default=1.0)
+    p.add_argument(
+        "--path-type",
+        type=str,
+        default=None,
+        choices=["linear", "cosine"],
+        help="默认从检查点 args 读取；可覆盖",
+    )
+    p.add_argument("--legacy", action=argparse.BooleanOptionalAction, default=False)
+    # 无 args 时的兜底
+    p.add_argument("--model", type=str, default=None, help="无检查点 args 时必填；与 SiT_models 键一致，如 SiT-XL/2")
+    p.add_argument("--resolution", type=int, default=None, choices=[256, 512])
+    p.add_argument("--num-classes", type=int, default=None)
+    p.add_argument("--encoder-depth", type=int, default=None)
+    p.add_argument("--enc-type", type=str, default=None)
+    p.add_argument("--fused-attn", action=argparse.BooleanOptionalAction, default=None)
+    p.add_argument("--qk-norm", action=argparse.BooleanOptionalAction, default=None)
+    p.add_argument("--cfg-prob", type=float, default=None)
+    p.add_argument(
+        "--sampler",
+        type=str,
+        default="em_image_noise",
+        choices=["ode", "em", "em_image_noise", "em_image_noise_before_tc"],
+        help="采样器：ode=euler_sampler 确定性漂移（linspace 1→0 或 t_c 分段直连 0，无 t_floor；与 EM 网格不同），"
+             "em=标准EM(含图像+cls噪声)，em_image_noise=仅图像噪声，"
+             "em_image_noise_before_tc=t<=t_c时图像去随机+cls全程去随机",
+    )
+    p.add_argument(
+        "--dual-compare-after",
+        action="store_true",
+        help="需配合分段步数：同批 z/y/cls 连跑两次；after_input 用 --steps-after-tc，"
+        "after_equal_before 将 after 步数设为与 --steps-before-tc 相同",
+    )
+    p.add_argument(
+        "--save-fixed-trajectory",
+        action="store_true",
+        help="保存固定步采样轨迹（npy）；仅对非 em 采样器启用，输出在 out-dir/trajectory",
+    )
+    return p.parse_args()
+
+
+def _decode_to_uint8_hwc(latents, latents_bias, latents_scale, vae):
+    imgs = vae.decode((latents - latents_bias) / latents_scale).sample
+    imgs = (imgs + 1) / 2.0
+    imgs = torch.clamp(imgs, 0, 1)
+    return (
+        (imgs * 255.0)
+        .round()
+        .to(torch.uint8)
+        .permute(0, 2, 3, 1)
+        .cpu()
+        .numpy()
+    )
+
+
+def main():
+    cli = parse_cli()
+    device = torch.device(cli.device if torch.cuda.is_available() else "cpu")
+    if device.type == "cuda":
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    try:
+        ckpt = torch.load(cli.ckpt, map_location="cpu", weights_only=False)
+    except TypeError:
+        ckpt = torch.load(cli.ckpt, map_location="cpu")
+    ta = load_train_args_from_ckpt(ckpt)
+    if ta is None:
+        if cli.model is None or cli.resolution is None or cli.enc_type is None:
+            print(
+                "检查点中无 args，请至少指定：--model --resolution --enc-type "
+                "（以及按需 --num-classes --encoder-depth）",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        ta = argparse.Namespace(
+            model=cli.model,
+            resolution=cli.resolution,
+            num_classes=cli.num_classes if cli.num_classes is not None else 1000,
+            encoder_depth=cli.encoder_depth if cli.encoder_depth is not None else 8,
+            enc_type=cli.enc_type,
+            fused_attn=cli.fused_attn if cli.fused_attn is not None else True,
+            qk_norm=cli.qk_norm if cli.qk_norm is not None else False,
+            cfg_prob=cli.cfg_prob if cli.cfg_prob is not None else 0.1,
+        )
+    else:
+        if cli.model is not None:
+            ta.model = cli.model
+        if cli.resolution is not None:
+            ta.resolution = cli.resolution
+        if cli.num_classes is not None:
+            ta.num_classes = cli.num_classes
+        if cli.encoder_depth is not None:
+            ta.encoder_depth = cli.encoder_depth
+        if cli.enc_type is not None:
+            ta.enc_type = cli.enc_type
+        if cli.fused_attn is not None:
+            ta.fused_attn = cli.fused_attn
+        if cli.qk_norm is not None:
+            ta.qk_norm = cli.qk_norm
+        if cli.cfg_prob is not None:
+            ta.cfg_prob = cli.cfg_prob
+
+    path_type = cli.path_type if cli.path_type is not None else getattr(ta, "path_type", "linear")
+
+    tc_split = resolve_tc_schedule(cli, ta)
+    if cli.dual_compare_after and tc_split[0] is None:
+        print("--dual-compare-after 必须配合 --steps-before-tc 与 --steps-after-tc（分段采样）", file=sys.stderr)
+        sys.exit(1)
+    if tc_split[0] is not None:
+        if cli.dual_compare_after:
+            print(
+                f"双次对比：t_c={tc_split[0]}, before={tc_split[1]}, "
+                f"after_input={tc_split[2]}, after_equal_before={tc_split[1]}"
+            )
+        else:
+            print(
+                f"时间网格：t_c={tc_split[0]}, 步数 (1→t_c)={tc_split[1]}, (t_c→0)={tc_split[2]} "
+                f"（总模型前向约 {tc_split[1] + tc_split[2] + 1} 次）"
+            )
+    else:
+        print(f"时间网格：均匀 num_steps={cli.num_steps}")
+
+    if cli.sampler == "ode":
+        sampler_fn = euler_ode_sampler
+    elif cli.sampler == "em":
+        sampler_fn = euler_maruyama_sampler
+    elif cli.sampler == "em_image_noise_before_tc":
+        sampler_fn = euler_maruyama_image_noise_before_tc_sampler
+    else:
+        sampler_fn = euler_maruyama_image_noise_sampler
+
+    model, cls_dim = build_model_from_train_args(ta, device)
+    wkey = cli.weights
+    if wkey not in ckpt:
+        raise KeyError(f"检查点中无 '{wkey}' 键，现有键：{list(ckpt.keys())}")
+    state = ckpt[wkey]
+    if cli.legacy:
+        from utils import load_legacy_checkpoints
+
+        state = load_legacy_checkpoints(
+            state_dict=state, encoder_depth=int(getattr(ta, "encoder_depth", 8))
+        )
+    model.load_state_dict(state, strict=True)
+    model.eval()
+
+    vae = load_vae(device)
+    latents_scale = torch.tensor([0.18215] * 4, device=device).view(1, 4, 1, 1)
+    latents_bias = torch.tensor([0.0] * 4, device=device).view(1, 4, 1, 1)
+
+    sampler_args = argparse.Namespace(cls_cfg_scale=float(cli.cls_cfg_scale))
+
+    at_tc_dir = at_tc_a = at_tc_b = None
+    pair_dir = None
+    traj_dir = traj_a = traj_b = None
+    if cli.dual_compare_after:
+        out_a = os.path.join(cli.out_dir, "after_input")
+        out_b = os.path.join(cli.out_dir, "after_equal_before")
+        pair_dir = os.path.join(cli.out_dir, "pair")
+        os.makedirs(out_a, exist_ok=True)
+        os.makedirs(out_b, exist_ok=True)
+        os.makedirs(pair_dir, exist_ok=True)
+        if tc_split[0] is not None:
+            at_tc_a = os.path.join(cli.out_dir, "at_tc", "after_input")
+            at_tc_b = os.path.join(cli.out_dir, "at_tc", "after_equal_before")
+            os.makedirs(at_tc_a, exist_ok=True)
+            os.makedirs(at_tc_b, exist_ok=True)
+        if cli.save_fixed_trajectory and cli.sampler != "em":
+            traj_a = os.path.join(cli.out_dir, "trajectory", "after_input")
+            traj_b = os.path.join(cli.out_dir, "trajectory", "after_equal_before")
+            os.makedirs(traj_a, exist_ok=True)
+            os.makedirs(traj_b, exist_ok=True)
+    else:
+        os.makedirs(cli.out_dir, exist_ok=True)
+        if tc_split[0] is not None:
+            at_tc_dir = os.path.join(cli.out_dir, "at_tc")
+            os.makedirs(at_tc_dir, exist_ok=True)
+        if cli.save_fixed_trajectory and cli.sampler != "em":
+            traj_dir = os.path.join(cli.out_dir, "trajectory")
+            os.makedirs(traj_dir, exist_ok=True)
+    latent_size = int(getattr(ta, "resolution", 256)) // 8
+    n_total = int(cli.num_images)
+    b = max(1, int(cli.batch_size))
+
+    torch.manual_seed(cli.seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(cli.seed)
+
+    written = 0
+    pbar = tqdm(total=n_total, desc="sampling")
+    while written < n_total:
+        cur = min(b, n_total - written)
+        z = torch.randn(cur, model.in_channels, latent_size, latent_size, device=device)
+        y = torch.randint(0, int(ta.num_classes), (cur,), device=device)
+        cls_z = torch.randn(cur, cls_dim, device=device)
+
+        with torch.no_grad():
+            base_kw = dict(
+                num_steps=cli.num_steps,
+                cfg_scale=cli.cfg_scale,
+                guidance_low=cli.guidance_low,
+                guidance_high=cli.guidance_high,
+                path_type=path_type,
+                cls_latents=cls_z,
+                args=sampler_args,
+            )
+            if cli.dual_compare_after:
+                tc_v, sb, sa_in = tc_split
+                # 两次完整采样会各自消耗 RNG；不重置则第二条的 1→t_c 噪声与第一条不同，z_tc/at_tc 会对不齐。
+                # 在固定 z/y/cls_z 之后打快照，第二条运行前恢复，使 t_c 中间态一致（仅后段步数不同）。
+                _rng_cpu_dual = torch.random.get_rng_state()
+                _rng_cuda_dual = (
+                    torch.cuda.get_rng_state_all()
+                    if device.type == "cuda"
+                    else None
+                )
+                batch_imgs = {}
+                for _run_i, (subdir, sa, tc_save_dir) in enumerate(
+                    (
+                        (out_a, sa_in, at_tc_a),
+                        (out_b, sb, at_tc_b),
+                    )
+                ):
+                    if _run_i > 0:
+                        torch.random.set_rng_state(_rng_cpu_dual)
+                        if _rng_cuda_dual is not None:
+                            torch.cuda.set_rng_state_all(_rng_cuda_dual)
+                    em_kw = dict(base_kw)
+                    em_kw["t_c"] = tc_v
+                    em_kw["num_steps_before_tc"] = sb
+                    em_kw["num_steps_after_tc"] = sa
+                    if cli.sampler == "em_image_noise_before_tc":
+                        if cli.save_fixed_trajectory and cli.sampler != "em":
+                            latents, z_tc, cls_tc, cls_t0, traj = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_v),
+                                return_cls_final=True,
+                                return_trajectory=True,
+                            )
+                        else:
+                            latents, z_tc, cls_tc, cls_t0 = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_v),
+                                return_cls_final=True,
+                            )
+                            traj = None
+                    else:
+                        if cli.save_fixed_trajectory and cli.sampler != "em":
+                            latents, z_tc, cls_tc, traj = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_v),
+                                return_trajectory=True,
+                            )
+                        else:
+                            latents, z_tc, cls_tc = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_v),
+                            )
+                            traj = None
+                        cls_t0 = None
+                    latents = latents.to(torch.float32)
+                    imgs = _decode_to_uint8_hwc(latents, latents_bias, latents_scale, vae)
+                    batch_imgs[subdir] = imgs
+                    for i in range(cur):
+                        Image.fromarray(imgs[i]).save(
+                            os.path.join(subdir, f"{written + i:06d}.png")
+                        )
+                    if tc_save_dir is not None and z_tc is not None:
+                        imgs_tc = _decode_to_uint8_hwc(
+                            z_tc.to(torch.float32), latents_bias, latents_scale, vae
+                        )
+                        for i in range(cur):
+                            Image.fromarray(imgs_tc[i]).save(
+                                os.path.join(tc_save_dir, f"{written + i:06d}.png")
+                            )
+                    if traj is not None:
+                        traj_np = torch.stack(traj, dim=0).to(torch.float32).cpu().numpy()
+                        save_traj_dir = traj_a if subdir == out_a else traj_b
+                        np.save(os.path.join(save_traj_dir, f"{written:06d}_traj.npy"), traj_np)
+                imgs_a = batch_imgs.get(out_a)
+                imgs_b = batch_imgs.get(out_b)
+                if pair_dir is not None and imgs_a is not None and imgs_b is not None:
+                    for i in range(cur):
+                        pair_img = np.concatenate([imgs_a[i], imgs_b[i]], axis=1)
+                        Image.fromarray(pair_img).save(
+                            os.path.join(pair_dir, f"{written + i:06d}.png")
+                        )
+            else:
+                em_kw = dict(base_kw)
+                if tc_split[0] is not None:
+                    em_kw["t_c"] = tc_split[0]
+                    em_kw["num_steps_before_tc"] = tc_split[1]
+                    em_kw["num_steps_after_tc"] = tc_split[2]
+                    if cli.sampler == "em_image_noise_before_tc":
+                        if cli.save_fixed_trajectory and cli.sampler != "em":
+                            latents, z_tc, cls_tc, cls_t0, traj = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_split[0]),
+                                return_cls_final=True,
+                                return_trajectory=True,
+                            )
+                        else:
+                            latents, z_tc, cls_tc, cls_t0 = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_split[0]),
+                                return_cls_final=True,
+                            )
+                            traj = None
+                    else:
+                        if cli.save_fixed_trajectory and cli.sampler != "em":
+                            latents, z_tc, cls_tc, traj = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_split[0]),
+                                return_trajectory=True,
+                            )
+                        else:
+                            latents, z_tc, cls_tc = sampler_fn(
+                                model,
+                                z,
+                                y,
+                                **em_kw,
+                                return_mid_state=True,
+                                t_mid=float(tc_split[0]),
+                            )
+                            traj = None
+                        cls_t0 = None
+                    latents = latents.to(torch.float32)
+                    if z_tc is not None and at_tc_dir is not None:
+                        imgs_tc = _decode_to_uint8_hwc(
+                            z_tc.to(torch.float32), latents_bias, latents_scale, vae
+                        )
+                        for i in range(cur):
+                            Image.fromarray(imgs_tc[i]).save(
+                                os.path.join(at_tc_dir, f"{written + i:06d}.png")
+                            )
+                    if traj is not None and traj_dir is not None:
+                        traj_np = torch.stack(traj, dim=0).to(torch.float32).cpu().numpy()
+                        np.save(os.path.join(traj_dir, f"{written:06d}_traj.npy"), traj_np)
+                else:
+                    latents = sampler_fn(model, z, y, **em_kw).to(torch.float32)
+                imgs = _decode_to_uint8_hwc(latents, latents_bias, latents_scale, vae)
+                for i in range(cur):
+                    Image.fromarray(imgs[i]).save(
+                        os.path.join(cli.out_dir, f"{written + i:06d}.png")
+                    )
+        written += cur
+        pbar.update(cur)
+    pbar.close()
+    if cli.dual_compare_after:
+        msg = (
+            f"Done. Saved {written} images per run under {out_a} and {out_b} "
+            f"(parent: {cli.out_dir})"
+        )
+        if pair_dir is not None:
+            msg += f"; paired comparisons under {pair_dir}"
+        if tc_split[0] is not None and at_tc_a is not None:
+            msg += f"; t≈t_c decoded under {at_tc_a} and {at_tc_b}"
+        print(msg)
+    else:
+        msg = f"Done. Saved {written} images under {cli.out_dir}"
+        if tc_split[0] is not None and at_tc_dir is not None:
+            msg += f"; t≈t_c decoded under {at_tc_dir}"
+        print(msg)
+
+
+if __name__ == "__main__":
+    main()

REG/sample_from_checkpoint_ddp.py

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

REG/samplers.py

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

REG/samples.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+# 双次对比步数请用 --dual-compare-after（见 sample_from_checkpoint.py），输出在 out-dir 子目录。
+
+CUDA_VISIBLE_DEVICES=1 python sample_from_checkpoint.py \
+  --ckpt /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/exps/jsflow-experiment-0.75-0.01-one-step/checkpoints/1970000.pt \
+  --out-dir ./my_samples_test \
+  --num-images 100 \
+  --batch-size 4 \
+  --seed 0 \
+  --t-c 0.75 \
+  --steps-before-tc 100\
+  --steps-after-tc 2 \
+  --sampler ode \
+  --cfg-scale 1.0 \
+  --dual-compare-after \

REG/samples_0.25_new.log

@@ -0,0 +1,43 @@
+W0408 11:05:05.680000 100112 site-packages/torch/distributed/run.py:793] 
+W0408 11:05:05.680000 100112 site-packages/torch/distributed/run.py:793] *****************************************
+W0408 11:05:05.680000 100112 site-packages/torch/distributed/run.py:793] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0408 11:05:05.680000 100112 site-packages/torch/distributed/run.py:793] *****************************************
+时间网格：t_c=0.25, 步数 (1→t_c)=100, (t_c→0)=2
+Total number of images that will be sampled: 5120
+
+[rank0]:[W408 11:06:31.621528015 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+[rank2]:[W408 11:06:31.627182760 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 2]  using GPU 2 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+[rank1]:[W408 11:06:32.966218444 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 1]  using GPU 1 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+
+W0408 11:11:04.746000 100112 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100158 closing signal SIGTERM
+W0408 11:11:04.748000 100112 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100159 closing signal SIGTERM
+W0408 11:11:04.749000 100112 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100160 closing signal SIGTERM
+W0408 11:11:04.749000 100112 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100161 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/bin/torchrun", line 33, in <module>
+    sys.exit(load_entry_point('torch==2.5.1', 'console_scripts', 'torchrun')())
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/run.py", line 919, in main
+    run(args)
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/run.py", line 910, in run
+    elastic_launch(
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 138, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 260, in launch_agent
+    result = agent.run()
+             ^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/elastic/metrics/api.py", line 137, in wrapper
+    result = f(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/elastic/agent/server/api.py", line 696, in run
+    result = self._invoke_run(role)
+             ^^^^^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/elastic/agent/server/api.py", line 855, in _invoke_run
+    time.sleep(monitor_interval)
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 84, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 100112 got signal: 15

REG/samples_0.75_new.log

@@ -0,0 +1,46 @@
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] 
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] *****************************************
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0325 16:55:16.395000 538513 site-packages/torch/distributed/run.py:793] *****************************************
+时间网格：t_c=0.75, 步数 (1→t_c)=100, (t_c→0)=50
+Total number of images that will be sampled: 40192
+
+[rank3]:[W325 16:57:00.799344818 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 3]  using GPU 3 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+[rank1]:[W325 16:57:00.847229448 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 1]  using GPU 1 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+[rank0]:[W325 16:57:01.326116049 ProcessGroupNCCL.cpp:4115] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device,or call init_process_group() with a device_id.
+
+W0325 18:03:18.212000 538513 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 538677 closing signal SIGTERM
+W0325 18:03:18.212000 538513 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 538678 closing signal SIGTERM
+E0325 18:03:18.554000 538513 site-packages/torch/distributed/elastic/multiprocessing/api.py:869] failed (exitcode: -9) local_rank: 1 (pid: 538676) of binary: /gemini/space/zhaozy/guzhenyu/envs/envs/SiT/bin/python
+Traceback (most recent call last):
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/bin/torchrun", line 33, in <module>
+    sys.exit(load_entry_point('torch==2.5.1', 'console_scripts', 'torchrun')())
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/run.py", line 919, in main
+    run(args)
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/run.py", line 910, in run
+    elastic_launch(
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 138, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+==========================================================
+sample_from_checkpoint_ddp.py FAILED
+----------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+----------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-03-25_18:03:18
+  host      : 24c964746905d416ce09d045f9a06f23-taskrole1-0
+  rank      : 1 (local_rank: 1)
+  exitcode  : -9 (pid: 538676)
+  error_file: <N/A>
+  traceback : Signal 9 (SIGKILL) received by PID 538676
+==========================================================

REG/samples_ddp.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+
+# 4 卡 DDP 单路径采样（不做 dual-compare，不保存 at_tc 中间图）
+CUDA_VISIBLE_DEVICES=0,1,2,3 nohup nohup torchrun \
+  --nnodes=1 \
+  --nproc_per_node=4 \
+  --rdzv_endpoint=localhost:29110 \
+  sample_from_checkpoint_ddp.py \
+  --ckpt /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/exps/jsflow-experiment-0.25/checkpoints/0230000.pt \
+  --out-dir ./my_samples_25 \
+  --num-images 5000 \
+  --batch-size 64 \
+  --seed 0 \
+  --t-c 0.25 \
+  --steps-before-tc 100 \
+  --steps-after-tc 2 \
+  --sampler em_image_noise_before_tc \
+  --cfg-scale 1.0 \
+  > samples_0.25_new.log 2>&1 & 
+
+# nohup python sample_from_checkpoint_ddp.py \
+#   --ckpt /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/exps/jsflow-experiment-0.5/checkpoints/0250000.pt \
+#   --out-dir ./my_samples_5 \
+#   --num-images 20000 \
+#   --batch-size 16 \
+#   --seed 0 \
+#   --t-c 0.5 \
+#   --steps-before-tc 100 \
+#   --steps-after-tc 50 \
+#   --sampler em_image_noise_before_tc \
+#   --cfg-scale 1.0 \
+#   > samples_0.5.log 2>&1 & 

REG/train.py

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

REG/train.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+# REG/train.py：与主仓库类似，可单独指定数据根目录与预处理 cls 特征目录。
+# 数据布局：${DATA_DIR}/imagenet_256_vae/ 下 VAE latent；
+#          ${SEMANTIC_FEATURES_DIR}/ 下 img-feature-*.npy + dataset.json（与 parallel_encode 一致）。
+
+NUM_GPUS=4
+
+# ------------ 按本机路径修改 ------------
+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"
+
+# 后台示例（与主实验脚本风格一致）：
+# nohup bash train.sh > jsflow-experiment.log 2>&1 &
+
+nohup accelerate launch --multi_gpu --num_processes "${NUM_GPUS}" --mixed_precision bf16 train.py \
+    --report-to wandb \
+    --allow-tf32 \
+    --mixed-precision bf16 \
+    --seed 0 \
+    --path-type linear \
+    --prediction v \
+    --weighting uniform \
+    --model SiT-XL/2 \
+    --enc-type dinov2-vit-b \
+    --encoder-depth 8 \
+    --proj-coeff 0.5 \
+    --output-dir exps \
+    --exp-name jsflow-experiment-0.75-0.01 \
+    --batch-size 256 \
+    --data-dir "${DATA_DIR}" \
+    --semantic-features-dir "${SEMANTIC_FEATURES_DIR}" \
+    --learning-rate 0.00005 \
+    --t-c 0.75 \
+    --cls 0.01 \
+    --ot-cls \
+    > jsflow-experiment.log 2>&1 & 
+
+# 说明：
+# - 不使用预处理特征、改在线抽 DINO 时：去掉 --semantic-features-dir，并保证 data-dir 为 REG 原布局
+#  （imagenet_256_vae + vae-sd）。
+# - 关闭 minibatch OT：追加 --no-ot-cls。
+# - 主仓库 train.py 中的 --weight-ratio / --semantic-reg-coeff / --repa-* 等为本 REG 脚本未实现项；
+#   投影强度请用 --proj-coeff，cls 流损失权重用 --cls。

REG/train_resume_tc_velocity.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# 基于指定 checkpoint 续训，并额外加入 t_c 处速度场监督项（用于增强单步性）
+set -euo pipefail
+
+NUM_GPUS=4
+
+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"
+
+# 用户指定的续训起点 checkpoint
+RESUME_CKPT="/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/exps/jsflow-experiment-0.75-0.01-one-step/checkpoints/1920000.pt"
+
+# 新增的 t_c 速度场损失权重（可按需调大/调小）
+TC_VEL_COEFF=2
+
+nohup accelerate launch --multi_gpu --num_processes "${NUM_GPUS}" --mixed_precision bf16 train.py \
+    --report-to wandb \
+    --allow-tf32 \
+    --mixed-precision bf16 \
+    --seed 0 \
+    --path-type linear \
+    --prediction v \
+    --weighting uniform \
+    --model SiT-XL/2 \
+    --enc-type dinov2-vit-b \
+    --encoder-depth 8 \
+    --proj-coeff 0.5 \
+    --output-dir exps \
+    --exp-name jsflow-experiment-0.75-0.01-one-step \
+    --batch-size 256 \
+    --data-dir "${DATA_DIR}" \
+    --semantic-features-dir "${SEMANTIC_FEATURES_DIR}" \
+    --learning-rate 0.00005 \
+    --t-c 0.75 \
+    --cls 0.005 \
+    --ot-cls \
+    --resume-from-ckpt "${RESUME_CKPT}" \
+    --tc-velocity-loss-coeff "${TC_VEL_COEFF}" \
+    > jsflow-experiment-0.75-0.01-tcvel.log 2>&1 &
+
+echo "Launched resume training with tc velocity loss."

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

REG/wandb/run-20260322_150022-yhxc5cgu/files/config.yaml

@@ -0,0 +1,202 @@
+_wandb:
+    value:
+        cli_version: 0.25.0
+        e:
+            ucanic8s891x6sl28vnbha78lzoecw66:
+                args:
+                    - --report-to
+                    - wandb
+                    - --allow-tf32
+                    - --mixed-precision
+                    - bf16
+                    - --seed
+                    - "0"
+                    - --path-type
+                    - linear
+                    - --prediction
+                    - v
+                    - --weighting
+                    - uniform
+                    - --model
+                    - SiT-XL/2
+                    - --enc-type
+                    - dinov2-vit-b
+                    - --encoder-depth
+                    - "8"
+                    - --proj-coeff
+                    - "0.5"
+                    - --output-dir
+                    - exps
+                    - --exp-name
+                    - jsflow-experiment
+                    - --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
+                    - --learning-rate
+                    - "0.00005"
+                    - --t-c
+                    - "0.5"
+                    - --cls
+                    - "0.2"
+                    - --ot-cls
+                codePath: train.py
+                codePathLocal: train.py
+                cpu_count: 96
+                cpu_count_logical: 192
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "3838880616448"
+                        used: "357557354496"
+                email: 2365972933@qq.com
+                executable: /gemini/space/zhaozy/guzhenyu/envs/envs/SiT/bin/python
+                git:
+                    commit: 021ea2e50c38c5803bd9afff16316958a01fbd1d
+                    remote: https://github.com/Martinser/REG.git
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 4
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-757303bb-4ec2-808b-a17f-95f6f5bad6dc
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-a09f2421-99e6-a72e-63bd-fd7452510758
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-9c670cc7-60a8-17f8-9b39-7ced3744976d
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-e6b1d8da-68d7-ed83-90d0-a4dedf33120e
+                host: 24c964746905d416ce09d045f9a06f23-taskrole1-0
+                memory:
+                    total: "2164115296256"
+                os: Linux-5.15.0-94-generic-x86_64-with-glibc2.35
+                program: /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py
+                python: CPython 3.12.9
+                root: /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG
+                startedAt: "2026-03-22T07:00:22.092510Z"
+                writerId: ucanic8s891x6sl28vnbha78lzoecw66
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 63
+                - 71
+                - 83
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 63
+                - 71
+                - 83
+                - 98
+            "3":
+                - 13
+            "4": 3.12.9
+            "5": 0.25.0
+            "6": 4.53.2
+            "12": 0.25.0
+            "13": linux-x86_64
+adam_beta1:
+    value: 0.9
+adam_beta2:
+    value: 0.999
+adam_epsilon:
+    value: 1e-08
+adam_weight_decay:
+    value: 0
+allow_tf32:
+    value: true
+batch_size:
+    value: 256
+cfg_prob:
+    value: 0.1
+checkpointing_steps:
+    value: 10000
+cls:
+    value: 0.2
+data_dir:
+    value: /gemini/space/zhaozy/dataset/Imagenet/imagenet_256
+enc_type:
+    value: dinov2-vit-b
+encoder_depth:
+    value: 8
+epochs:
+    value: 1400
+exp_name:
+    value: jsflow-experiment
+fused_attn:
+    value: true
+gradient_accumulation_steps:
+    value: 1
+learning_rate:
+    value: 5e-05
+legacy:
+    value: false
+logging_dir:
+    value: logs
+max_grad_norm:
+    value: 1
+max_train_steps:
+    value: 1000000
+mixed_precision:
+    value: bf16
+model:
+    value: SiT-XL/2
+num_classes:
+    value: 1000
+num_workers:
+    value: 4
+ops_head:
+    value: 16
+ot_cls:
+    value: true
+output_dir:
+    value: exps
+path_type:
+    value: linear
+prediction:
+    value: v
+proj_coeff:
+    value: 0.5
+qk_norm:
+    value: false
+report_to:
+    value: wandb
+resolution:
+    value: 256
+resume_step:
+    value: 0
+sampling_steps:
+    value: 2000
+seed:
+    value: 0
+semantic_features_dir:
+    value: /gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0
+t_c:
+    value: 0.5
+weighting:
+    value: uniform

REG/wandb/run-20260322_150022-yhxc5cgu/files/wandb-summary.json

@@ -0,0 +1 @@
+{"_wandb":{"runtime":2},"_runtime":2}

REG/wandb/run-20260322_150443-e3yw9ii4/files/config.yaml

@@ -0,0 +1,202 @@
+_wandb:
+    value:
+        cli_version: 0.25.0
+        e:
+            q63x26q8nhayytv8q2rrmj9j9uy9kvub:
+                args:
+                    - --report-to
+                    - wandb
+                    - --allow-tf32
+                    - --mixed-precision
+                    - bf16
+                    - --seed
+                    - "0"
+                    - --path-type
+                    - linear
+                    - --prediction
+                    - v
+                    - --weighting
+                    - uniform
+                    - --model
+                    - SiT-XL/2
+                    - --enc-type
+                    - dinov2-vit-b
+                    - --encoder-depth
+                    - "8"
+                    - --proj-coeff
+                    - "0.5"
+                    - --output-dir
+                    - exps
+                    - --exp-name
+                    - jsflow-experiment
+                    - --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
+                    - --learning-rate
+                    - "0.00005"
+                    - --t-c
+                    - "0.5"
+                    - --cls
+                    - "0.2"
+                    - --ot-cls
+                codePath: train.py
+                codePathLocal: train.py
+                cpu_count: 96
+                cpu_count_logical: 192
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "3838880616448"
+                        used: "357557714944"
+                email: 2365972933@qq.com
+                executable: /gemini/space/zhaozy/guzhenyu/envs/envs/SiT/bin/python
+                git:
+                    commit: 021ea2e50c38c5803bd9afff16316958a01fbd1d
+                    remote: https://github.com/Martinser/REG.git
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 4
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-757303bb-4ec2-808b-a17f-95f6f5bad6dc
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-a09f2421-99e6-a72e-63bd-fd7452510758
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-9c670cc7-60a8-17f8-9b39-7ced3744976d
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-e6b1d8da-68d7-ed83-90d0-a4dedf33120e
+                host: 24c964746905d416ce09d045f9a06f23-taskrole1-0
+                memory:
+                    total: "2164115296256"
+                os: Linux-5.15.0-94-generic-x86_64-with-glibc2.35
+                program: /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py
+                python: CPython 3.12.9
+                root: /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG
+                startedAt: "2026-03-22T07:04:43.133739Z"
+                writerId: q63x26q8nhayytv8q2rrmj9j9uy9kvub
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 63
+                - 71
+                - 83
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 63
+                - 71
+                - 83
+                - 98
+            "3":
+                - 13
+            "4": 3.12.9
+            "5": 0.25.0
+            "6": 4.53.2
+            "12": 0.25.0
+            "13": linux-x86_64
+adam_beta1:
+    value: 0.9
+adam_beta2:
+    value: 0.999
+adam_epsilon:
+    value: 1e-08
+adam_weight_decay:
+    value: 0
+allow_tf32:
+    value: true
+batch_size:
+    value: 256
+cfg_prob:
+    value: 0.1
+checkpointing_steps:
+    value: 10000
+cls:
+    value: 0.2
+data_dir:
+    value: /gemini/space/zhaozy/dataset/Imagenet/imagenet_256
+enc_type:
+    value: dinov2-vit-b
+encoder_depth:
+    value: 8
+epochs:
+    value: 1400
+exp_name:
+    value: jsflow-experiment
+fused_attn:
+    value: true
+gradient_accumulation_steps:
+    value: 1
+learning_rate:
+    value: 5e-05
+legacy:
+    value: false
+logging_dir:
+    value: logs
+max_grad_norm:
+    value: 1
+max_train_steps:
+    value: 1000000
+mixed_precision:
+    value: bf16
+model:
+    value: SiT-XL/2
+num_classes:
+    value: 1000
+num_workers:
+    value: 4
+ops_head:
+    value: 16
+ot_cls:
+    value: true
+output_dir:
+    value: exps
+path_type:
+    value: linear
+prediction:
+    value: v
+proj_coeff:
+    value: 0.5
+qk_norm:
+    value: false
+report_to:
+    value: wandb
+resolution:
+    value: 256
+resume_step:
+    value: 0
+sampling_steps:
+    value: 2000
+seed:
+    value: 0
+semantic_features_dir:
+    value: /gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0
+t_c:
+    value: 0.5
+weighting:
+    value: uniform

REG/wandb/run-20260322_150443-e3yw9ii4/files/output.log

@@ -0,0 +1,15 @@
+Steps:   0%|          | 1/1000000 [00:02<588:29:38,  2.12s/it][2026-03-22 15:04:48] Generating EMA samples for evaluation (SDE → t=0)...
+Traceback (most recent call last):
+  File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 572, in <module>
+    main(args)
+  File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 444, in main
+    if vae is not None:
+       ^^^
+NameError: name 'vae' is not defined
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 572, in <module>
+[rank0]:     main(args)
+[rank0]:   File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 444, in main
+[rank0]:     if vae is not None:
+[rank0]:        ^^^
+[rank0]: NameError: name 'vae' is not defined

REG/wandb/run-20260322_150443-e3yw9ii4/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

REG/wandb/run-20260322_150443-e3yw9ii4/files/wandb-metadata.json

@@ -0,0 +1,101 @@
+{
+  "os":  "Linux-5.15.0-94-generic-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-03-22T07:04:43.133739Z",
+  "args":  [
+    "--report-to",
+    "wandb",
+    "--allow-tf32",
+    "--mixed-precision",
+    "bf16",
+    "--seed",
+    "0",
+    "--path-type",
+    "linear",
+    "--prediction",
+    "v",
+    "--weighting",
+    "uniform",
+    "--model",
+    "SiT-XL/2",
+    "--enc-type",
+    "dinov2-vit-b",
+    "--encoder-depth",
+    "8",
+    "--proj-coeff",
+    "0.5",
+    "--output-dir",
+    "exps",
+    "--exp-name",
+    "jsflow-experiment",
+    "--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",
+    "--learning-rate",
+    "0.00005",
+    "--t-c",
+    "0.5",
+    "--cls",
+    "0.2",
+    "--ot-cls"
+  ],
+  "program":  "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/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/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":  "357557714944"
+    }
+  },
+  "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":  "q63x26q8nhayytv8q2rrmj9j9uy9kvub"
+}

REG/wandb/run-20260322_150443-e3yw9ii4/files/wandb-summary.json

@@ -0,0 +1 @@
+{"_wandb":{"runtime":3},"_runtime":3}

REG/wandb/run-20260322_150443-e3yw9ii4/logs/debug-internal.log

@@ -0,0 +1,7 @@
+{"time":"2026-03-22T15:04:43.390486873+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-22T15:04:44.970687851+08:00","level":"INFO","msg":"stream: created new stream","id":"e3yw9ii4"}
+{"time":"2026-03-22T15:04:44.970802178+08:00","level":"INFO","msg":"handler: started","stream_id":"e3yw9ii4"}
+{"time":"2026-03-22T15:04:44.971744065+08:00","level":"INFO","msg":"stream: started","id":"e3yw9ii4"}
+{"time":"2026-03-22T15:04:44.97174913+08:00","level":"INFO","msg":"writer: started","stream_id":"e3yw9ii4"}
+{"time":"2026-03-22T15:04:44.971758857+08:00","level":"INFO","msg":"sender: started","stream_id":"e3yw9ii4"}
+{"time":"2026-03-22T15:04:50.286711145+08:00","level":"INFO","msg":"stream: closing","id":"e3yw9ii4"}

REG/wandb/run-20260322_150443-e3yw9ii4/logs/debug.log

@@ -0,0 +1,22 @@
+2026-03-22 15:04:43,155 INFO    MainThread:326012 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-22 15:04:43,156 INFO    MainThread:326012 [wandb_setup.py:_flush():81] Configure stats pid to 326012
+2026-03-22 15:04:43,156 INFO    MainThread:326012 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-22 15:04:43,156 INFO    MainThread:326012 [wandb_init.py:setup_run_log_directory():717] Logging user logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260322_150443-e3yw9ii4/logs/debug.log
+2026-03-22 15:04:43,156 INFO    MainThread:326012 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260322_150443-e3yw9ii4/logs/debug-internal.log
+2026-03-22 15:04:43,156 INFO    MainThread:326012 [wandb_init.py:init():844] calling init triggers
+2026-03-22 15:04:43,156 INFO    MainThread:326012 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-22 15:04:43,156 INFO    MainThread:326012 [wandb_init.py:init():892] starting backend
+2026-03-22 15:04:43,378 INFO    MainThread:326012 [wandb_init.py:init():895] sending inform_init request
+2026-03-22 15:04:43,388 INFO    MainThread:326012 [wandb_init.py:init():903] backend started and connected
+2026-03-22 15:04:43,389 INFO    MainThread:326012 [wandb_init.py:init():973] updated telemetry
+2026-03-22 15:04:43,402 INFO    MainThread:326012 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-22 15:04:46,450 INFO    MainThread:326012 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-22 15:04:46,541 INFO    MainThread:326012 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-22 15:04:46,541 INFO    MainThread:326012 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-22 15:04:46,541 INFO    MainThread:326012 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-22 15:04:46,541 INFO    MainThread:326012 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-22 15:04:46,545 INFO    MainThread:326012 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-22 15:04:46,545 INFO    MainThread:326012 [wandb_run.py:_config_callback():1403] config_cb None None {'output_dir': 'exps', 'exp_name': 'jsflow-experiment', 'logging_dir': 'logs', 'report_to': 'wandb', 'sampling_steps': 2000, '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': 'bf16', 'epochs': 1400, 'max_train_steps': 1000000, 'checkpointing_steps': 10000, 'gradient_accumulation_steps': 1, 'learning_rate': 5e-05, '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.2, 't_c': 0.5, 'ot_cls': True}
+2026-03-22 15:04:50,286 INFO    wandb-AsyncioManager-main:326012 [service_client.py:_forward_responses():134] Reached EOF.
+2026-03-22 15:04:50,286 INFO    wandb-AsyncioManager-main:326012 [mailbox.py:close():155] Closing mailbox, abandoning 1 handles.

REG/wandb/run-20260322_150635-o2w3z8rq/logs/debug-internal.log

@@ -0,0 +1,6 @@
+{"time":"2026-03-22T15:06:35.605542246+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-22T15:06:37.585674571+08:00","level":"INFO","msg":"stream: created new stream","id":"o2w3z8rq"}
+{"time":"2026-03-22T15:06:37.585934805+08:00","level":"INFO","msg":"handler: started","stream_id":"o2w3z8rq"}
+{"time":"2026-03-22T15:06:37.586954142+08:00","level":"INFO","msg":"stream: started","id":"o2w3z8rq"}
+{"time":"2026-03-22T15:06:37.587002572+08:00","level":"INFO","msg":"sender: started","stream_id":"o2w3z8rq"}
+{"time":"2026-03-22T15:06:37.58696296+08:00","level":"INFO","msg":"writer: started","stream_id":"o2w3z8rq"}

REG/wandb/run-20260322_150635-o2w3z8rq/logs/debug.log

@@ -0,0 +1,20 @@
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_setup.py:_flush():81] Configure stats pid to 328110
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_init.py:setup_run_log_directory():717] Logging user logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260322_150635-o2w3z8rq/logs/debug.log
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260322_150635-o2w3z8rq/logs/debug-internal.log
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_init.py:init():844] calling init triggers
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-22 15:06:35,364 INFO    MainThread:328110 [wandb_init.py:init():892] starting backend
+2026-03-22 15:06:35,590 INFO    MainThread:328110 [wandb_init.py:init():895] sending inform_init request
+2026-03-22 15:06:35,601 INFO    MainThread:328110 [wandb_init.py:init():903] backend started and connected
+2026-03-22 15:06:35,604 INFO    MainThread:328110 [wandb_init.py:init():973] updated telemetry
+2026-03-22 15:06:35,618 INFO    MainThread:328110 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-22 15:06:38,166 INFO    MainThread:328110 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-22 15:06:38,258 INFO    MainThread:328110 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-22 15:06:38,258 INFO    MainThread:328110 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-22 15:06:38,259 INFO    MainThread:328110 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-22 15:06:38,259 INFO    MainThread:328110 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-22 15:06:38,262 INFO    MainThread:328110 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-22 15:06:38,263 INFO    MainThread:328110 [wandb_run.py:_config_callback():1403] config_cb None None {'output_dir': 'exps', 'exp_name': 'jsflow-experiment', 'logging_dir': 'logs', 'report_to': 'wandb', 'sampling_steps': 2000, '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': 'bf16', 'epochs': 1400, 'max_train_steps': 1000000, 'checkpointing_steps': 10000, 'gradient_accumulation_steps': 1, 'learning_rate': 5e-05, '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.2, 't_c': 0.5, 'ot_cls': True}

REG/wandb/run-20260323_135607-zue1y2ba/files/output.log

@@ -0,0 +1,289 @@
+Steps:   0%|          | 1/1000000 [00:02<603:07:37,  2.17s/it][2026-03-23 13:56:13] Generating EMA samples (Euler-Maruyama; t≈0.75 → t=0)...
+[2026-03-23 13:56:16] Step: 1, Training Logs: loss_final: 4.886707, loss_mean: 1.706308, proj_loss: 0.001541, loss_mean_cls: 3.178859, grad_norm: 1.484174
+Steps:   0%|          | 2/1000000 [00:04<669:29:28,  2.41s/it, grad_norm=1.48, loss_final=4.89, loss_mean=1.71, loss_mean_cls=3.18, proj_loss=0.00154][2026-03-23 13:56:16] Step: 2, Training Logs: loss_final: 4.294325, loss_mean: 1.689698, proj_loss: -0.010266, loss_mean_cls: 2.614893, grad_norm: 1.062257
+Steps:   0%|          | 3/1000000 [00:04<394:40:41,  1.42s/it, grad_norm=1.06, loss_final=4.29, loss_mean=1.69, loss_mean_cls=2.61, proj_loss=-0.0103][2026-03-23 13:56:16] Step: 3, Training Logs: loss_final: 4.829489, loss_mean: 1.666703, proj_loss: -0.019215, loss_mean_cls: 3.182001, grad_norm: 1.115940
+Steps:   0%|          | 4/1000000 [00:05<265:36:26,  1.05it/s, grad_norm=1.12, loss_final=4.83, loss_mean=1.67, loss_mean_cls=3.18, proj_loss=-0.0192][2026-03-23 13:56:16] Step: 4, Training Logs: loss_final: 4.838729, loss_mean: 1.683388, proj_loss: -0.026348, loss_mean_cls: 3.181689, grad_norm: 0.753645
+Steps:   0%|          | 5/1000000 [00:05<194:13:56,  1.43it/s, grad_norm=0.754, loss_final=4.84, loss_mean=1.68, loss_mean_cls=3.18, proj_loss=-0.0263][2026-03-23 13:56:17] Step: 5, Training Logs: loss_final: 4.434019, loss_mean: 1.678989, proj_loss: -0.034618, loss_mean_cls: 2.789647, grad_norm: 0.828005
+Steps:   0%|          | 6/1000000 [00:05<151:13:33,  1.84it/s, grad_norm=0.828, loss_final=4.43, loss_mean=1.68, loss_mean_cls=2.79, proj_loss=-0.0346][2026-03-23 13:56:17] Step: 6, Training Logs: loss_final: 4.717834, loss_mean: 1.685299, proj_loss: -0.039478, loss_mean_cls: 3.072013, grad_norm: 0.944329
+Steps:   0%|          | 7/1000000 [00:05<123:56:37,  2.24it/s, grad_norm=0.944, loss_final=4.72, loss_mean=1.69, loss_mean_cls=3.07, proj_loss=-0.0395][2026-03-23 13:56:17] Step: 7, Training Logs: loss_final: 4.768410, loss_mean: 1.687585, proj_loss: -0.042706, loss_mean_cls: 3.123530, grad_norm: 0.827338
+Steps:   0%|          | 8/1000000 [00:06<106:03:40,  2.62it/s, grad_norm=0.827, loss_final=4.77, loss_mean=1.69, loss_mean_cls=3.12, proj_loss=-0.0427][2026-03-23 13:56:17] Step: 8, Training Logs: loss_final: 4.815378, loss_mean: 1.655840, proj_loss: -0.044987, loss_mean_cls: 3.204525, grad_norm: 0.854108
+Steps:   0%|          | 9/1000000 [00:06<94:05:27,  2.95it/s, grad_norm=0.854, loss_final=4.82, loss_mean=1.66, loss_mean_cls=3.2, proj_loss=-0.045] [2026-03-23 13:56:18] Step: 9, Training Logs: loss_final: 4.777181, loss_mean: 1.657609, proj_loss: -0.046041, loss_mean_cls: 3.165613, grad_norm: 0.939216
+Steps:   0%|          | 10/1000000 [00:06<86:16:44,  3.22it/s, grad_norm=0.939, loss_final=4.78, loss_mean=1.66, loss_mean_cls=3.17, proj_loss=-0.046][2026-03-23 13:56:18] Step: 10, Training Logs: loss_final: 4.966327, loss_mean: 1.662955, proj_loss: -0.047787, loss_mean_cls: 3.351158, grad_norm: 0.991139
+Steps:   0%|          | 11/1000000 [00:06<80:53:45,  3.43it/s, grad_norm=0.991, loss_final=4.97, loss_mean=1.66, loss_mean_cls=3.35, proj_loss=-0.0478][2026-03-23 13:56:18] Step: 11, Training Logs: loss_final: 5.292466, loss_mean: 1.650309, proj_loss: -0.049361, loss_mean_cls: 3.691518, grad_norm: 0.968609
+Steps:   0%|          | 12/1000000 [00:07<76:53:52,  3.61it/s, grad_norm=0.969, loss_final=5.29, loss_mean=1.65, loss_mean_cls=3.69, proj_loss=-0.0494][2026-03-23 13:56:18] Step: 12, Training Logs: loss_final: 5.088713, loss_mean: 1.626424, proj_loss: -0.049838, loss_mean_cls: 3.512127, grad_norm: 1.162980
+Steps:   0%|          | 13/1000000 [00:07<74:06:27,  3.75it/s, grad_norm=1.16, loss_final=5.09, loss_mean=1.63, loss_mean_cls=3.51, proj_loss=-0.0498][2026-03-23 13:56:19] Step: 13, Training Logs: loss_final: 4.928039, loss_mean: 1.623318, proj_loss: -0.051428, loss_mean_cls: 3.356148, grad_norm: 1.370081
+Steps:   0%|          | 14/1000000 [00:07<72:19:32,  3.84it/s, grad_norm=1.37, loss_final=4.93, loss_mean=1.62, loss_mean_cls=3.36, proj_loss=-0.0514][2026-03-23 13:56:19] Step: 14, Training Logs: loss_final: 4.342262, loss_mean: 1.599151, proj_loss: -0.051557, loss_mean_cls: 2.794668, grad_norm: 1.247963
+Steps:   0%|          | 15/1000000 [00:07<70:56:55,  3.92it/s, grad_norm=1.25, loss_final=4.34, loss_mean=1.6, loss_mean_cls=2.79, proj_loss=-0.0516][2026-03-23 13:56:19] Step: 15, Training Logs: loss_final: 5.220107, loss_mean: 1.589133, proj_loss: -0.054820, loss_mean_cls: 3.685793, grad_norm: 1.207489
+Steps:   0%|          | 16/1000000 [00:08<69:57:27,  3.97it/s, grad_norm=1.21, loss_final=5.22, loss_mean=1.59, loss_mean_cls=3.69, proj_loss=-0.0548][2026-03-23 13:56:19] Step: 16, Training Logs: loss_final: 4.653113, loss_mean: 1.599458, proj_loss: -0.052329, loss_mean_cls: 3.105984, grad_norm: 1.168053
+Steps:   0%|          | 17/1000000 [00:08<69:18:40,  4.01it/s, grad_norm=1.17, loss_final=4.65, loss_mean=1.6, loss_mean_cls=3.11, proj_loss=-0.0523][2026-03-23 13:56:20] Step: 17, Training Logs: loss_final: 4.984623, loss_mean: 1.639365, proj_loss: -0.051690, loss_mean_cls: 3.396948, grad_norm: 2.396022
+Steps:   0%|          | 18/1000000 [00:08<69:00:10,  4.03it/s, grad_norm=2.4, loss_final=4.98, loss_mean=1.64, loss_mean_cls=3.4, proj_loss=-0.0517][2026-03-23 13:56:20] Step: 18, Training Logs: loss_final: 4.356441, loss_mean: 1.560717, proj_loss: -0.054883, loss_mean_cls: 2.850607, grad_norm: 0.931938
+Steps:   0%|          | 19/1000000 [00:08<68:37:45,  4.05it/s, grad_norm=0.932, loss_final=4.36, loss_mean=1.56, loss_mean_cls=2.85, proj_loss=-0.0549][2026-03-23 13:56:20] Step: 19, Training Logs: loss_final: 4.922860, loss_mean: 1.625848, proj_loss: -0.055209, loss_mean_cls: 3.352221, grad_norm: 2.649956
+Steps:   0%|          | 20/1000000 [00:09<68:31:08,  4.05it/s, grad_norm=2.65, loss_final=4.92, loss_mean=1.63, loss_mean_cls=3.35, proj_loss=-0.0552][2026-03-23 13:56:20] Step: 20, Training Logs: loss_final: 4.693020, loss_mean: 1.635111, proj_loss: -0.053463, loss_mean_cls: 3.111372, grad_norm: 2.886806
+Steps:   0%|          | 21/1000000 [00:09<68:17:30,  4.07it/s, grad_norm=2.89, loss_final=4.69, loss_mean=1.64, loss_mean_cls=3.11, proj_loss=-0.0535][2026-03-23 13:56:21] Step: 21, Training Logs: loss_final: 4.847960, loss_mean: 1.608342, proj_loss: -0.053926, loss_mean_cls: 3.293545, grad_norm: 1.804076
+Steps:   0%|          | 22/1000000 [00:09<68:07:59,  4.08it/s, grad_norm=1.8, loss_final=4.85, loss_mean=1.61, loss_mean_cls=3.29, proj_loss=-0.0539][2026-03-23 13:56:21] Step: 22, Training Logs: loss_final: 4.531857, loss_mean: 1.556872, proj_loss: -0.053500, loss_mean_cls: 3.028484, grad_norm: 1.136153
+Steps:   0%|          | 23/1000000 [00:09<68:01:00,  4.08it/s, grad_norm=1.14, loss_final=4.53, loss_mean=1.56, loss_mean_cls=3.03, proj_loss=-0.0535][2026-03-23 13:56:21] Step: 23, Training Logs: loss_final: 4.347858, loss_mean: 1.571953, proj_loss: -0.051053, loss_mean_cls: 2.826958, grad_norm: 1.379415
+Steps:   0%|          | 24/1000000 [00:10<67:55:02,  4.09it/s, grad_norm=1.38, loss_final=4.35, loss_mean=1.57, loss_mean_cls=2.83, proj_loss=-0.0511][2026-03-23 13:56:21] Step: 24, Training Logs: loss_final: 4.812301, loss_mean: 1.573597, proj_loss: -0.054251, loss_mean_cls: 3.292954, grad_norm: 1.536997
+Steps:   0%|          | 25/1000000 [00:10<67:50:56,  4.09it/s, grad_norm=1.54, loss_final=4.81, loss_mean=1.57, loss_mean_cls=3.29, proj_loss=-0.0543][2026-03-23 13:56:22] Step: 25, Training Logs: loss_final: 5.140490, loss_mean: 1.572134, proj_loss: -0.055976, loss_mean_cls: 3.624332, grad_norm: 1.501310
+Steps:   0%|          | 26/1000000 [00:10<67:49:32,  4.10it/s, grad_norm=1.5, loss_final=5.14, loss_mean=1.57, loss_mean_cls=3.62, proj_loss=-0.056][2026-03-23 13:56:22] Step: 26, Training Logs: loss_final: 4.614564, loss_mean: 1.575298, proj_loss: -0.055522, loss_mean_cls: 3.094787, grad_norm: 1.374379
+Steps:   0%|          | 27/1000000 [00:10<67:48:51,  4.10it/s, grad_norm=1.37, loss_final=4.61, loss_mean=1.58, loss_mean_cls=3.09, proj_loss=-0.0555][2026-03-23 13:56:22] Step: 27, Training Logs: loss_final: 4.428196, loss_mean: 1.559373, proj_loss: -0.054907, loss_mean_cls: 2.923730, grad_norm: 1.247505
+Steps:   0%|          | 28/1000000 [00:11<67:51:05,  4.09it/s, grad_norm=1.25, loss_final=4.43, loss_mean=1.56, loss_mean_cls=2.92, proj_loss=-0.0549][2026-03-23 13:56:22] Step: 28, Training Logs: loss_final: 4.514633, loss_mean: 1.561718, proj_loss: -0.054182, loss_mean_cls: 3.007097, grad_norm: 1.272959
+Steps:   0%|          | 29/1000000 [00:11<67:50:57,  4.09it/s, grad_norm=1.27, loss_final=4.51, loss_mean=1.56, loss_mean_cls=3.01, proj_loss=-0.0542][2026-03-23 13:56:23] Step: 29, Training Logs: loss_final: 4.064789, loss_mean: 1.520746, proj_loss: -0.055004, loss_mean_cls: 2.599047, grad_norm: 1.213601
+Steps:   0%|          | 30/1000000 [00:11<67:49:21,  4.10it/s, grad_norm=1.21, loss_final=4.06, loss_mean=1.52, loss_mean_cls=2.6, proj_loss=-0.055][2026-03-23 13:56:23] Step: 30, Training Logs: loss_final: 4.343926, loss_mean: 1.523063, proj_loss: -0.055674, loss_mean_cls: 2.876538, grad_norm: 1.148791
+Steps:   0%|          | 31/1000000 [00:11<67:50:28,  4.09it/s, grad_norm=1.15, loss_final=4.34, loss_mean=1.52, loss_mean_cls=2.88, proj_loss=-0.0557][2026-03-23 13:56:23] Step: 31, Training Logs: loss_final: 4.943371, loss_mean: 1.500860, proj_loss: -0.056089, loss_mean_cls: 3.498600, grad_norm: 1.126629
+Steps:   0%|          | 32/1000000 [00:12<67:51:49,  4.09it/s, grad_norm=1.13, loss_final=4.94, loss_mean=1.5, loss_mean_cls=3.5, proj_loss=-0.0561][2026-03-23 13:56:23] Step: 32, Training Logs: loss_final: 4.913333, loss_mean: 1.487980, proj_loss: -0.053908, loss_mean_cls: 3.479261, grad_norm: 1.111562
+Steps:   0%|          | 33/1000000 [00:12<67:49:26,  4.10it/s, grad_norm=1.11, loss_final=4.91, loss_mean=1.49, loss_mean_cls=3.48, proj_loss=-0.0539][2026-03-23 13:56:24] Step: 33, Training Logs: loss_final: 5.135122, loss_mean: 1.485391, proj_loss: -0.058625, loss_mean_cls: 3.708356, grad_norm: 1.166992
+Steps:   0%|          | 34/1000000 [00:12<67:49:31,  4.10it/s, grad_norm=1.17, loss_final=5.14, loss_mean=1.49, loss_mean_cls=3.71, proj_loss=-0.0586][2026-03-23 13:56:24] Step: 34, Training Logs: loss_final: 4.094241, loss_mean: 1.475285, proj_loss: -0.056888, loss_mean_cls: 2.675844, grad_norm: 1.037110
+Steps:   0%|          | 35/1000000 [00:12<67:47:41,  4.10it/s, grad_norm=1.04, loss_final=4.09, loss_mean=1.48, loss_mean_cls=2.68, proj_loss=-0.0569][2026-03-23 13:56:24] Step: 35, Training Logs: loss_final: 4.305937, loss_mean: 1.460218, proj_loss: -0.052705, loss_mean_cls: 2.898424, grad_norm: 1.091828
+Steps:   0%|          | 36/1000000 [00:13<67:50:00,  4.09it/s, grad_norm=1.09, loss_final=4.31, loss_mean=1.46, loss_mean_cls=2.9, proj_loss=-0.0527][2026-03-23 13:56:24] Step: 36, Training Logs: loss_final: 4.382219, loss_mean: 1.466251, proj_loss: -0.054344, loss_mean_cls: 2.970313, grad_norm: 1.093016
+Steps:   0%|          | 37/1000000 [00:13<68:14:38,  4.07it/s, grad_norm=1.09, loss_final=4.38, loss_mean=1.47, loss_mean_cls=2.97, proj_loss=-0.0543][2026-03-23 13:56:25] Step: 37, Training Logs: loss_final: 4.428336, loss_mean: 1.469206, proj_loss: -0.057364, loss_mean_cls: 3.016495, grad_norm: 1.468433
+Steps:   0%|          | 38/1000000 [00:13<68:07:38,  4.08it/s, grad_norm=1.47, loss_final=4.43, loss_mean=1.47, loss_mean_cls=3.02, proj_loss=-0.0574][2026-03-23 13:56:25] Step: 38, Training Logs: loss_final: 4.412411, loss_mean: 1.453988, proj_loss: -0.056726, loss_mean_cls: 3.015149, grad_norm: 0.825492
+Steps:   0%|          | 39/1000000 [00:13<68:01:05,  4.08it/s, grad_norm=0.825, loss_final=4.41, loss_mean=1.45, loss_mean_cls=3.02, proj_loss=-0.0567][2026-03-23 13:56:25] Step: 39, Training Logs: loss_final: 4.554020, loss_mean: 1.450588, proj_loss: -0.055717, loss_mean_cls: 3.159149, grad_norm: 1.281297
+Steps:   0%|          | 40/1000000 [00:14<68:02:46,  4.08it/s, grad_norm=1.28, loss_final=4.55, loss_mean=1.45, loss_mean_cls=3.16, proj_loss=-0.0557][2026-03-23 13:56:25] Step: 40, Training Logs: loss_final: 4.900630, loss_mean: 1.421538, proj_loss: -0.056336, loss_mean_cls: 3.535427, grad_norm: 1.096009
+Steps:   0%|          | 41/1000000 [00:14<67:56:56,  4.09it/s, grad_norm=1.1, loss_final=4.9, loss_mean=1.42, loss_mean_cls=3.54, proj_loss=-0.0563][2026-03-23 13:56:26] Step: 41, Training Logs: loss_final: 5.000130, loss_mean: 1.418576, proj_loss: -0.056068, loss_mean_cls: 3.637622, grad_norm: 1.111790
+Steps:   0%|          | 42/1000000 [00:14<67:53:46,  4.09it/s, grad_norm=1.11, loss_final=5, loss_mean=1.42, loss_mean_cls=3.64, proj_loss=-0.0561][2026-03-23 13:56:26] Step: 42, Training Logs: loss_final: 4.223974, loss_mean: 1.442425, proj_loss: -0.054436, loss_mean_cls: 2.835985, grad_norm: 1.223004
+Steps:   0%|          | 43/1000000 [00:14<67:51:37,  4.09it/s, grad_norm=1.22, loss_final=4.22, loss_mean=1.44, loss_mean_cls=2.84, proj_loss=-0.0544][2026-03-23 13:56:26] Step: 43, Training Logs: loss_final: 5.023710, loss_mean: 1.441718, proj_loss: -0.055548, loss_mean_cls: 3.637539, grad_norm: 1.391159
+Steps:   0%|          | 44/1000000 [00:15<67:53:09,  4.09it/s, grad_norm=1.39, loss_final=5.02, loss_mean=1.44, loss_mean_cls=3.64, proj_loss=-0.0555][2026-03-23 13:56:26] Step: 44, Training Logs: loss_final: 5.000957, loss_mean: 1.414230, proj_loss: -0.055362, loss_mean_cls: 3.642089, grad_norm: 1.162115
+Steps:   0%|          | 45/1000000 [00:15<67:51:51,  4.09it/s, grad_norm=1.16, loss_final=5, loss_mean=1.41, loss_mean_cls=3.64, proj_loss=-0.0554][2026-03-23 13:56:27] Step: 45, Training Logs: loss_final: 4.689414, loss_mean: 1.389372, proj_loss: -0.054793, loss_mean_cls: 3.354835, grad_norm: 0.814496
+Steps:   0%|          | 46/1000000 [00:15<67:51:06,  4.09it/s, grad_norm=0.814, loss_final=4.69, loss_mean=1.39, loss_mean_cls=3.35, proj_loss=-0.0548][2026-03-23 13:56:27] Step: 46, Training Logs: loss_final: 4.452005, loss_mean: 1.403370, proj_loss: -0.056346, loss_mean_cls: 3.104981, grad_norm: 1.062373
+Steps:   0%|          | 47/1000000 [00:15<67:50:59,  4.09it/s, grad_norm=1.06, loss_final=4.45, loss_mean=1.4, loss_mean_cls=3.1, proj_loss=-0.0563][2026-03-23 13:56:27] Step: 47, Training Logs: loss_final: 4.638161, loss_mean: 1.413495, proj_loss: -0.055944, loss_mean_cls: 3.280609, grad_norm: 0.873316
+Steps:   0%|          | 48/1000000 [00:15<67:51:08,  4.09it/s, grad_norm=0.873, loss_final=4.64, loss_mean=1.41, loss_mean_cls=3.28, proj_loss=-0.0559][2026-03-23 13:56:27] Step: 48, Training Logs: loss_final: 4.429680, loss_mean: 1.411911, proj_loss: -0.054931, loss_mean_cls: 3.072701, grad_norm: 0.759066
+Steps:   0%|          | 49/1000000 [00:16<67:51:48,  4.09it/s, grad_norm=0.759, loss_final=4.43, loss_mean=1.41, loss_mean_cls=3.07, proj_loss=-0.0549][2026-03-23 13:56:27] Step: 49, Training Logs: loss_final: 4.901843, loss_mean: 1.388557, proj_loss: -0.055650, loss_mean_cls: 3.568936, grad_norm: 0.750354
+Steps:   0%|          | 50/1000000 [00:16<67:49:39,  4.10it/s, grad_norm=0.75, loss_final=4.9, loss_mean=1.39, loss_mean_cls=3.57, proj_loss=-0.0557][2026-03-23 13:56:28] Step: 50, Training Logs: loss_final: 4.100002, loss_mean: 1.412004, proj_loss: -0.052795, loss_mean_cls: 2.740793, grad_norm: 0.767011
+Steps:   0%|          | 51/1000000 [00:16<67:51:25,  4.09it/s, grad_norm=0.767, loss_final=4.1, loss_mean=1.41, loss_mean_cls=2.74, proj_loss=-0.0528][2026-03-23 13:56:28] Step: 51, Training Logs: loss_final: 4.759223, loss_mean: 1.391587, proj_loss: -0.054960, loss_mean_cls: 3.422596, grad_norm: 0.719157
+Steps:   0%|          | 52/1000000 [00:16<67:51:11,  4.09it/s, grad_norm=0.719, loss_final=4.76, loss_mean=1.39, loss_mean_cls=3.42, proj_loss=-0.055][2026-03-23 13:56:28] Step: 52, Training Logs: loss_final: 4.482382, loss_mean: 1.395736, proj_loss: -0.056075, loss_mean_cls: 3.142721, grad_norm: 0.726095
+Steps:   0%|          | 53/1000000 [00:17<67:50:55,  4.09it/s, grad_norm=0.726, loss_final=4.48, loss_mean=1.4, loss_mean_cls=3.14, proj_loss=-0.0561][2026-03-23 13:56:28] Step: 53, Training Logs: loss_final: 4.466825, loss_mean: 1.373211, proj_loss: -0.055033, loss_mean_cls: 3.148648, grad_norm: 0.663936
+Steps:   0%|          | 54/1000000 [00:17<67:49:43,  4.10it/s, grad_norm=0.664, loss_final=4.47, loss_mean=1.37, loss_mean_cls=3.15, proj_loss=-0.055][2026-03-23 13:56:29] Step: 54, Training Logs: loss_final: 4.687413, loss_mean: 1.367031, proj_loss: -0.056856, loss_mean_cls: 3.377239, grad_norm: 0.706146
+Steps:   0%|          | 55/1000000 [00:17<67:47:50,  4.10it/s, grad_norm=0.706, loss_final=4.69, loss_mean=1.37, loss_mean_cls=3.38, proj_loss=-0.0569][2026-03-23 13:56:29] Step: 55, Training Logs: loss_final: 4.724741, loss_mean: 1.355043, proj_loss: -0.056556, loss_mean_cls: 3.426254, grad_norm: 0.891681
+Steps:   0%|          | 56/1000000 [00:17<67:49:45,  4.10it/s, grad_norm=0.892, loss_final=4.72, loss_mean=1.36, loss_mean_cls=3.43, proj_loss=-0.0566][2026-03-23 13:56:29] Step: 56, Training Logs: loss_final: 4.598913, loss_mean: 1.357820, proj_loss: -0.053707, loss_mean_cls: 3.294799, grad_norm: 0.688579
+Steps:   0%|          | 57/1000000 [00:18<67:48:24,  4.10it/s, grad_norm=0.689, loss_final=4.6, loss_mean=1.36, loss_mean_cls=3.29, proj_loss=-0.0537][2026-03-23 13:56:29] Step: 57, Training Logs: loss_final: 4.647234, loss_mean: 1.346388, proj_loss: -0.055346, loss_mean_cls: 3.356192, grad_norm: 0.643795
+Steps:   0%|          | 58/1000000 [00:18<67:48:01,  4.10it/s, grad_norm=0.644, loss_final=4.65, loss_mean=1.35, loss_mean_cls=3.36, proj_loss=-0.0553][2026-03-23 13:56:30] Step: 58, Training Logs: loss_final: 4.629024, loss_mean: 1.361648, proj_loss: -0.056354, loss_mean_cls: 3.323730, grad_norm: 0.640788
+Steps:   0%|          | 59/1000000 [00:18<67:48:29,  4.10it/s, grad_norm=0.641, loss_final=4.63, loss_mean=1.36, loss_mean_cls=3.32, proj_loss=-0.0564][2026-03-23 13:56:30] Step: 59, Training Logs: loss_final: 4.318052, loss_mean: 1.352944, proj_loss: -0.054974, loss_mean_cls: 3.020082, grad_norm: 1.103365
+Steps:   0%|          | 60/1000000 [00:18<67:52:13,  4.09it/s, grad_norm=1.1, loss_final=4.32, loss_mean=1.35, loss_mean_cls=3.02, proj_loss=-0.055][2026-03-23 13:56:30] Step: 60, Training Logs: loss_final: 4.558999, loss_mean: 1.356349, proj_loss: -0.056304, loss_mean_cls: 3.258954, grad_norm: 2.178408
+Steps:   0%|          | 61/1000000 [00:19<67:50:06,  4.09it/s, grad_norm=2.18, loss_final=4.56, loss_mean=1.36, loss_mean_cls=3.26, proj_loss=-0.0563][2026-03-23 13:56:30] Step: 61, Training Logs: loss_final: 4.405947, loss_mean: 1.345089, proj_loss: -0.054088, loss_mean_cls: 3.114946, grad_norm: 0.833812
+Steps:   0%|          | 62/1000000 [00:19<67:49:48,  4.09it/s, grad_norm=0.834, loss_final=4.41, loss_mean=1.35, loss_mean_cls=3.11, proj_loss=-0.0541][2026-03-23 13:56:31] Step: 62, Training Logs: loss_final: 4.248630, loss_mean: 1.369897, proj_loss: -0.055107, loss_mean_cls: 2.933840, grad_norm: 2.604588
+Steps:   0%|          | 63/1000000 [00:19<67:49:04,  4.10it/s, grad_norm=2.6, loss_final=4.25, loss_mean=1.37, loss_mean_cls=2.93, proj_loss=-0.0551][2026-03-23 13:56:31] Step: 63, Training Logs: loss_final: 4.370800, loss_mean: 1.356826, proj_loss: -0.054216, loss_mean_cls: 3.068191, grad_norm: 1.995899
+Steps:   0%|          | 64/1000000 [00:19<67:52:12,  4.09it/s, grad_norm=2, loss_final=4.37, loss_mean=1.36, loss_mean_cls=3.07, proj_loss=-0.0542][2026-03-23 13:56:31] Step: 64, Training Logs: loss_final: 3.620857, loss_mean: 1.371637, proj_loss: -0.056188, loss_mean_cls: 2.305408, grad_norm: 1.270011
+Steps:   0%|          | 65/1000000 [00:20<67:52:26,  4.09it/s, grad_norm=1.27, loss_final=3.62, loss_mean=1.37, loss_mean_cls=2.31, proj_loss=-0.0562][2026-03-23 13:56:31] Step: 65, Training Logs: loss_final: 3.841930, loss_mean: 1.345907, proj_loss: -0.057290, loss_mean_cls: 2.553313, grad_norm: 1.103519
+Steps:   0%|          | 66/1000000 [00:20<67:50:32,  4.09it/s, grad_norm=1.1, loss_final=3.84, loss_mean=1.35, loss_mean_cls=2.55, proj_loss=-0.0573][2026-03-23 13:56:32] Step: 66, Training Logs: loss_final: 4.823415, loss_mean: 1.309009, proj_loss: -0.057261, loss_mean_cls: 3.571667, grad_norm: 1.146726
+Steps:   0%|          | 67/1000000 [00:20<67:51:11,  4.09it/s, grad_norm=1.15, loss_final=4.82, loss_mean=1.31, loss_mean_cls=3.57, proj_loss=-0.0573][2026-03-23 13:56:32] Step: 67, Training Logs: loss_final: 4.654735, loss_mean: 1.315030, proj_loss: -0.055812, loss_mean_cls: 3.395517, grad_norm: 0.900543
+Steps:   0%|          | 68/1000000 [00:20<67:53:43,  4.09it/s, grad_norm=0.901, loss_final=4.65, loss_mean=1.32, loss_mean_cls=3.4, proj_loss=-0.0558][2026-03-23 13:56:32] Step: 68, Training Logs: loss_final: 3.865313, loss_mean: 1.311920, proj_loss: -0.055171, loss_mean_cls: 2.608564, grad_norm: 0.968363
+Steps:   0%|          | 69/1000000 [00:21<67:51:01,  4.09it/s, grad_norm=0.968, loss_final=3.87, loss_mean=1.31, loss_mean_cls=2.61, proj_loss=-0.0552][2026-03-23 13:56:32] Step: 69, Training Logs: loss_final: 4.332185, loss_mean: 1.287292, proj_loss: -0.055524, loss_mean_cls: 3.100417, grad_norm: 0.992981
+Steps:   0%|          | 70/1000000 [00:21<67:50:34,  4.09it/s, grad_norm=0.993, loss_final=4.33, loss_mean=1.29, loss_mean_cls=3.1, proj_loss=-0.0555][2026-03-23 13:56:33] Step: 70, Training Logs: loss_final: 5.214293, loss_mean: 1.271327, proj_loss: -0.056000, loss_mean_cls: 3.998966, grad_norm: 0.950980
+Steps:   0%|          | 71/1000000 [00:21<67:51:17,  4.09it/s, grad_norm=0.951, loss_final=5.21, loss_mean=1.27, loss_mean_cls=4, proj_loss=-0.056][2026-03-23 13:56:33] Step: 71, Training Logs: loss_final: 4.271416, loss_mean: 1.298077, proj_loss: -0.058767, loss_mean_cls: 3.032106, grad_norm: 0.897326
+Steps:   0%|          | 72/1000000 [00:21<67:53:03,  4.09it/s, grad_norm=0.897, loss_final=4.27, loss_mean=1.3, loss_mean_cls=3.03, proj_loss=-0.0588][2026-03-23 13:56:33] Step: 72, Training Logs: loss_final: 4.944882, loss_mean: 1.275258, proj_loss: -0.054986, loss_mean_cls: 3.724611, grad_norm: 1.406698
+Steps:   0%|          | 73/1000000 [00:22<67:51:07,  4.09it/s, grad_norm=1.41, loss_final=4.94, loss_mean=1.28, loss_mean_cls=3.72, proj_loss=-0.055][2026-03-23 13:56:33] Step: 73, Training Logs: loss_final: 4.262454, loss_mean: 1.391114, proj_loss: -0.054645, loss_mean_cls: 2.925985, grad_norm: 4.082608
+Steps:   0%|          | 74/1000000 [00:22<67:49:45,  4.09it/s, grad_norm=4.08, loss_final=4.26, loss_mean=1.39, loss_mean_cls=2.93, proj_loss=-0.0546][2026-03-23 13:56:34] Step: 74, Training Logs: loss_final: 4.441927, loss_mean: 1.368322, proj_loss: -0.056059, loss_mean_cls: 3.129665, grad_norm: 3.740271
+Steps:   0%|          | 75/1000000 [00:22<67:48:41,  4.10it/s, grad_norm=3.74, loss_final=4.44, loss_mean=1.37, loss_mean_cls=3.13, proj_loss=-0.0561][2026-03-23 13:56:34] Step: 75, Training Logs: loss_final: 4.786633, loss_mean: 1.323001, proj_loss: -0.054923, loss_mean_cls: 3.518555, grad_norm: 1.360202
+Steps:   0%|          | 76/1000000 [00:22<67:47:44,  4.10it/s, grad_norm=1.36, loss_final=4.79, loss_mean=1.32, loss_mean_cls=3.52, proj_loss=-0.0549][2026-03-23 13:56:34] Step: 76, Training Logs: loss_final: 4.880804, loss_mean: 1.312154, proj_loss: -0.058812, loss_mean_cls: 3.627462, grad_norm: 2.462873
+Steps:   0%|          | 77/1000000 [00:23<67:48:38,  4.10it/s, grad_norm=2.46, loss_final=4.88, loss_mean=1.31, loss_mean_cls=3.63, proj_loss=-0.0588][2026-03-23 13:56:34] Step: 77, Training Logs: loss_final: 4.347323, loss_mean: 1.277282, proj_loss: -0.056339, loss_mean_cls: 3.126381, grad_norm: 1.089468
+Steps:   0%|          | 78/1000000 [00:23<67:47:33,  4.10it/s, grad_norm=1.09, loss_final=4.35, loss_mean=1.28, loss_mean_cls=3.13, proj_loss=-0.0563][2026-03-23 13:56:35] Step: 78, Training Logs: loss_final: 4.452549, loss_mean: 1.308045, proj_loss: -0.057751, loss_mean_cls: 3.202255, grad_norm: 1.409237
+Steps:   0%|          | 79/1000000 [00:23<67:48:01,  4.10it/s, grad_norm=1.41, loss_final=4.45, loss_mean=1.31, loss_mean_cls=3.2, proj_loss=-0.0578][2026-03-23 13:56:35] Step: 79, Training Logs: loss_final: 4.300065, loss_mean: 1.256680, proj_loss: -0.057729, loss_mean_cls: 3.101114, grad_norm: 1.063736
+Steps:   0%|          | 80/1000000 [00:23<67:46:24,  4.10it/s, grad_norm=1.06, loss_final=4.3, loss_mean=1.26, loss_mean_cls=3.1, proj_loss=-0.0577][2026-03-23 13:56:35] Step: 80, Training Logs: loss_final: 5.255699, loss_mean: 1.230699, proj_loss: -0.055231, loss_mean_cls: 4.080230, grad_norm: 1.159766
+Steps:   0%|          | 81/1000000 [00:24<67:45:22,  4.10it/s, grad_norm=1.16, loss_final=5.26, loss_mean=1.23, loss_mean_cls=4.08, proj_loss=-0.0552][2026-03-23 13:56:35] Step: 81, Training Logs: loss_final: 4.184104, loss_mean: 1.256977, proj_loss: -0.057558, loss_mean_cls: 2.984685, grad_norm: 1.074197
+Steps:   0%|          | 82/1000000 [00:24<67:46:48,  4.10it/s, grad_norm=1.07, loss_final=4.18, loss_mean=1.26, loss_mean_cls=2.98, proj_loss=-0.0576][2026-03-23 13:56:36] Step: 82, Training Logs: loss_final: 4.368578, loss_mean: 1.246449, proj_loss: -0.055964, loss_mean_cls: 3.178093, grad_norm: 1.212082
+Steps:   0%|          | 83/1000000 [00:24<67:46:42,  4.10it/s, grad_norm=1.21, loss_final=4.37, loss_mean=1.25, loss_mean_cls=3.18, proj_loss=-0.056][2026-03-23 13:56:36] Step: 83, Training Logs: loss_final: 4.066692, loss_mean: 1.279943, proj_loss: -0.057100, loss_mean_cls: 2.843850, grad_norm: 1.081535
+Steps:   0%|          | 84/1000000 [00:24<67:45:25,  4.10it/s, grad_norm=1.08, loss_final=4.07, loss_mean=1.28, loss_mean_cls=2.84, proj_loss=-0.0571][2026-03-23 13:56:36] Step: 84, Training Logs: loss_final: 4.419612, loss_mean: 1.222443, proj_loss: -0.056123, loss_mean_cls: 3.253291, grad_norm: 1.017672
+Steps:   0%|          | 85/1000000 [00:25<67:46:48,  4.10it/s, grad_norm=1.02, loss_final=4.42, loss_mean=1.22, loss_mean_cls=3.25, proj_loss=-0.0561][2026-03-23 13:56:36] Step: 85, Training Logs: loss_final: 4.330603, loss_mean: 1.217940, proj_loss: -0.055811, loss_mean_cls: 3.168474, grad_norm: 1.048292
+Steps:   0%|          | 86/1000000 [00:25<67:46:18,  4.10it/s, grad_norm=1.05, loss_final=4.33, loss_mean=1.22, loss_mean_cls=3.17, proj_loss=-0.0558][2026-03-23 13:56:37] Step: 86, Training Logs: loss_final: 4.073160, loss_mean: 1.261520, proj_loss: -0.057447, loss_mean_cls: 2.869087, grad_norm: 2.760909
+Steps:   0%|          | 87/1000000 [00:25<67:46:56,  4.10it/s, grad_norm=2.76, loss_final=4.07, loss_mean=1.26, loss_mean_cls=2.87, proj_loss=-0.0574][2026-03-23 13:56:37] Step: 87, Training Logs: loss_final: 4.349278, loss_mean: 1.225316, proj_loss: -0.055978, loss_mean_cls: 3.179941, grad_norm: 1.415633
+Steps:   0%|          | 88/1000000 [00:25<67:48:19,  4.10it/s, grad_norm=1.42, loss_final=4.35, loss_mean=1.23, loss_mean_cls=3.18, proj_loss=-0.056][2026-03-23 13:56:37] Step: 88, Training Logs: loss_final: 4.191838, loss_mean: 1.244159, proj_loss: -0.053158, loss_mean_cls: 3.000837, grad_norm: 2.142123
+Steps:   0%|          | 89/1000000 [00:26<67:48:47,  4.10it/s, grad_norm=2.14, loss_final=4.19, loss_mean=1.24, loss_mean_cls=3, proj_loss=-0.0532][2026-03-23 13:56:37] Step: 89, Training Logs: loss_final: 3.810838, loss_mean: 1.241034, proj_loss: -0.056567, loss_mean_cls: 2.626369, grad_norm: 1.977536
+Steps:   0%|          | 90/1000000 [00:26<67:49:26,  4.10it/s, grad_norm=1.98, loss_final=3.81, loss_mean=1.24, loss_mean_cls=2.63, proj_loss=-0.0566][2026-03-23 13:56:37] Step: 90, Training Logs: loss_final: 4.302361, loss_mean: 1.235736, proj_loss: -0.055757, loss_mean_cls: 3.122383, grad_norm: 1.465364
+Steps:   0%|          | 91/1000000 [00:26<67:49:17,  4.10it/s, grad_norm=1.47, loss_final=4.3, loss_mean=1.24, loss_mean_cls=3.12, proj_loss=-0.0558][2026-03-23 13:56:38] Step: 91, Training Logs: loss_final: 4.793566, loss_mean: 1.295077, proj_loss: -0.054678, loss_mean_cls: 3.553167, grad_norm: 4.777657
+Steps:   0%|          | 92/1000000 [00:26<67:50:23,  4.09it/s, grad_norm=4.78, loss_final=4.79, loss_mean=1.3, loss_mean_cls=3.55, proj_loss=-0.0547][2026-03-23 13:56:38] Step: 92, Training Logs: loss_final: 4.318158, loss_mean: 1.268673, proj_loss: -0.056761, loss_mean_cls: 3.106246, grad_norm: 3.593853
+Steps:   0%|          | 93/1000000 [00:26<67:49:16,  4.10it/s, grad_norm=3.59, loss_final=4.32, loss_mean=1.27, loss_mean_cls=3.11, proj_loss=-0.0568][2026-03-23 13:56:38] Step: 93, Training Logs: loss_final: 4.038731, loss_mean: 1.237636, proj_loss: -0.058346, loss_mean_cls: 2.859441, grad_norm: 2.444315
+Steps:   0%|          | 94/1000000 [00:27<67:49:30,  4.10it/s, grad_norm=2.44, loss_final=4.04, loss_mean=1.24, loss_mean_cls=2.86, proj_loss=-0.0583][2026-03-23 13:56:38] Step: 94, Training Logs: loss_final: 5.108729, loss_mean: 1.212993, proj_loss: -0.058394, loss_mean_cls: 3.954130, grad_norm: 1.895388
+Steps:   0%|          | 95/1000000 [00:27<67:48:34,  4.10it/s, grad_norm=1.9, loss_final=5.11, loss_mean=1.21, loss_mean_cls=3.95, proj_loss=-0.0584][2026-03-23 13:56:39] Step: 95, Training Logs: loss_final: 4.209163, loss_mean: 1.207963, proj_loss: -0.055074, loss_mean_cls: 3.056275, grad_norm: 1.442985
+Steps:   0%|          | 96/1000000 [00:27<67:50:56,  4.09it/s, grad_norm=1.44, loss_final=4.21, loss_mean=1.21, loss_mean_cls=3.06, proj_loss=-0.0551][2026-03-23 13:56:39] Step: 96, Training Logs: loss_final: 4.601864, loss_mean: 1.191701, proj_loss: -0.057576, loss_mean_cls: 3.467739, grad_norm: 1.299760
+Steps:   0%|          | 97/1000000 [00:27<67:51:29,  4.09it/s, grad_norm=1.3, loss_final=4.6, loss_mean=1.19, loss_mean_cls=3.47, proj_loss=-0.0576][2026-03-23 13:56:39] Step: 97, Training Logs: loss_final: 4.614489, loss_mean: 1.185281, proj_loss: -0.056098, loss_mean_cls: 3.485306, grad_norm: 1.128512
+Steps:   0%|          | 98/1000000 [00:28<67:50:38,  4.09it/s, grad_norm=1.13, loss_final=4.61, loss_mean=1.19, loss_mean_cls=3.49, proj_loss=-0.0561][2026-03-23 13:56:39] Step: 98, Training Logs: loss_final: 4.294365, loss_mean: 1.196555, proj_loss: -0.058301, loss_mean_cls: 3.156111, grad_norm: 1.399289
+Steps:   0%|          | 99/1000000 [00:28<67:49:41,  4.09it/s, grad_norm=1.4, loss_final=4.29, loss_mean=1.2, loss_mean_cls=3.16, proj_loss=-0.0583][2026-03-23 13:56:40] Step: 99, Training Logs: loss_final: 4.215539, loss_mean: 1.199367, proj_loss: -0.057493, loss_mean_cls: 3.073665, grad_norm: 1.343760
+Steps:   0%|          | 100/1000000 [00:28<67:51:40,  4.09it/s, grad_norm=1.34, loss_final=4.22, loss_mean=1.2, loss_mean_cls=3.07, proj_loss=-0.0575][2026-03-23 13:56:40] Step: 100, Training Logs: loss_final: 4.488984, loss_mean: 1.149765, proj_loss: -0.056815, loss_mean_cls: 3.396034, grad_norm: 1.045582
+Steps:   0%|          | 101/1000000 [00:28<67:50:50,  4.09it/s, grad_norm=1.05, loss_final=4.49, loss_mean=1.15, loss_mean_cls=3.4, proj_loss=-0.0568][2026-03-23 13:56:40] Step: 101, Training Logs: loss_final: 4.368436, loss_mean: 1.178788, proj_loss: -0.058116, loss_mean_cls: 3.247765, grad_norm: 1.342256
+Steps:   0%|          | 102/1000000 [00:29<67:50:14,  4.09it/s, grad_norm=1.34, loss_final=4.37, loss_mean=1.18, loss_mean_cls=3.25, proj_loss=-0.0581][2026-03-23 13:56:40] Step: 102, Training Logs: loss_final: 4.096940, loss_mean: 1.195665, proj_loss: -0.053437, loss_mean_cls: 2.954712, grad_norm: 1.167338
+Steps:   0%|          | 103/1000000 [00:29<67:52:23,  4.09it/s, grad_norm=1.17, loss_final=4.1, loss_mean=1.2, loss_mean_cls=2.95, proj_loss=-0.0534][2026-03-23 13:56:41] Step: 103, Training Logs: loss_final: 4.406911, loss_mean: 1.123473, proj_loss: -0.055880, loss_mean_cls: 3.339318, grad_norm: 0.934016
+Steps:   0%|          | 104/1000000 [00:29<67:52:00,  4.09it/s, grad_norm=0.934, loss_final=4.41, loss_mean=1.12, loss_mean_cls=3.34, proj_loss=-0.0559][2026-03-23 13:56:41] Step: 104, Training Logs: loss_final: 3.831389, loss_mean: 1.127624, proj_loss: -0.056442, loss_mean_cls: 2.760207, grad_norm: 1.079511
+Steps:   0%|          | 105/1000000 [00:29<67:50:29,  4.09it/s, grad_norm=1.08, loss_final=3.83, loss_mean=1.13, loss_mean_cls=2.76, proj_loss=-0.0564][2026-03-23 13:56:41] Step: 105, Training Logs: loss_final: 3.949926, loss_mean: 1.169529, proj_loss: -0.056276, loss_mean_cls: 2.836674, grad_norm: 0.823243
+Steps:   0%|          | 106/1000000 [00:30<67:49:55,  4.09it/s, grad_norm=0.823, loss_final=3.95, loss_mean=1.17, loss_mean_cls=2.84, proj_loss=-0.0563][2026-03-23 13:56:41] Step: 106, Training Logs: loss_final: 4.447761, loss_mean: 1.131930, proj_loss: -0.057379, loss_mean_cls: 3.373209, grad_norm: 0.780265
+Steps:   0%|          | 107/1000000 [00:30<67:50:26,  4.09it/s, grad_norm=0.78, loss_final=4.45, loss_mean=1.13, loss_mean_cls=3.37, proj_loss=-0.0574][2026-03-23 13:56:42] Step: 107, Training Logs: loss_final: 4.185682, loss_mean: 1.137062, proj_loss: -0.056792, loss_mean_cls: 3.105412, grad_norm: 0.896073
+Steps:   0%|          | 108/1000000 [00:30<67:51:47,  4.09it/s, grad_norm=0.896, loss_final=4.19, loss_mean=1.14, loss_mean_cls=3.11, proj_loss=-0.0568][2026-03-23 13:56:42] Step: 108, Training Logs: loss_final: 4.381330, loss_mean: 1.127375, proj_loss: -0.057231, loss_mean_cls: 3.311185, grad_norm: 0.791624
+Steps:   0%|          | 109/1000000 [00:30<67:50:27,  4.09it/s, grad_norm=0.792, loss_final=4.38, loss_mean=1.13, loss_mean_cls=3.31, proj_loss=-0.0572][2026-03-23 13:56:42] Step: 109, Training Logs: loss_final: 4.101658, loss_mean: 1.133106, proj_loss: -0.058366, loss_mean_cls: 3.026918, grad_norm: 1.389539
+Steps:   0%|          | 110/1000000 [00:31<67:50:56,  4.09it/s, grad_norm=1.39, loss_final=4.1, loss_mean=1.13, loss_mean_cls=3.03, proj_loss=-0.0584][2026-03-23 13:56:42] Step: 110, Training Logs: loss_final: 4.222694, loss_mean: 1.125139, proj_loss: -0.056927, loss_mean_cls: 3.154482, grad_norm: 1.616820
+Steps:   0%|          | 111/1000000 [00:31<67:52:18,  4.09it/s, grad_norm=1.62, loss_final=4.22, loss_mean=1.13, loss_mean_cls=3.15, proj_loss=-0.0569][2026-03-23 13:56:43] Step: 111, Training Logs: loss_final: 3.836849, loss_mean: 1.123032, proj_loss: -0.057576, loss_mean_cls: 2.771394, grad_norm: 1.516640
+Steps:   0%|          | 112/1000000 [00:31<67:50:27,  4.09it/s, grad_norm=1.52, loss_final=3.84, loss_mean=1.12, loss_mean_cls=2.77, proj_loss=-0.0576][2026-03-23 13:56:43] Step: 112, Training Logs: loss_final: 4.385750, loss_mean: 1.170239, proj_loss: -0.057827, loss_mean_cls: 3.273338, grad_norm: 3.234420
+Steps:   0%|          | 113/1000000 [00:31<67:50:21,  4.09it/s, grad_norm=3.23, loss_final=4.39, loss_mean=1.17, loss_mean_cls=3.27, proj_loss=-0.0578][2026-03-23 13:56:43] Step: 113, Training Logs: loss_final: 4.791656, loss_mean: 1.139730, proj_loss: -0.057919, loss_mean_cls: 3.709845, grad_norm: 2.920728
+Steps:   0%|          | 114/1000000 [00:32<67:52:04,  4.09it/s, grad_norm=2.92, loss_final=4.79, loss_mean=1.14, loss_mean_cls=3.71, proj_loss=-0.0579][2026-03-23 13:56:43] Step: 114, Training Logs: loss_final: 4.198164, loss_mean: 1.112274, proj_loss: -0.054186, loss_mean_cls: 3.140076, grad_norm: 1.197350
+Steps:   0%|          | 115/1000000 [00:32<67:52:07,  4.09it/s, grad_norm=1.2, loss_final=4.2, loss_mean=1.11, loss_mean_cls=3.14, proj_loss=-0.0542][2026-03-23 13:56:44] Step: 115, Training Logs: loss_final: 4.495673, loss_mean: 1.149152, proj_loss: -0.056316, loss_mean_cls: 3.402837, grad_norm: 1.762717
+Steps:   0%|          | 116/1000000 [00:32<67:52:00,  4.09it/s, grad_norm=1.76, loss_final=4.5, loss_mean=1.15, loss_mean_cls=3.4, proj_loss=-0.0563][2026-03-23 13:56:44] Step: 116, Training Logs: loss_final: 3.863493, loss_mean: 1.173006, proj_loss: -0.057397, loss_mean_cls: 2.747884, grad_norm: 2.976036
+Steps:   0%|          | 117/1000000 [00:32<67:52:53,  4.09it/s, grad_norm=2.98, loss_final=3.86, loss_mean=1.17, loss_mean_cls=2.75, proj_loss=-0.0574][2026-03-23 13:56:44] Step: 117, Training Logs: loss_final: 4.637027, loss_mean: 1.119735, proj_loss: -0.056042, loss_mean_cls: 3.573334, grad_norm: 2.094532
+Steps:   0%|          | 118/1000000 [00:33<67:52:15,  4.09it/s, grad_norm=2.09, loss_final=4.64, loss_mean=1.12, loss_mean_cls=3.57, proj_loss=-0.056][2026-03-23 13:56:44] Step: 118, Training Logs: loss_final: 3.896062, loss_mean: 1.133374, proj_loss: -0.057744, loss_mean_cls: 2.820432, grad_norm: 2.091462
+Steps:   0%|          | 119/1000000 [00:33<67:52:19,  4.09it/s, grad_norm=2.09, loss_final=3.9, loss_mean=1.13, loss_mean_cls=2.82, proj_loss=-0.0577][2026-03-23 13:56:45] Step: 119, Training Logs: loss_final: 4.494811, loss_mean: 1.127229, proj_loss: -0.056165, loss_mean_cls: 3.423747, grad_norm: 2.573430
+Steps:   0%|          | 120/1000000 [00:33<67:52:07,  4.09it/s, grad_norm=2.57, loss_final=4.49, loss_mean=1.13, loss_mean_cls=3.42, proj_loss=-0.0562][2026-03-23 13:56:45] Step: 120, Training Logs: loss_final: 4.395417, loss_mean: 1.094057, proj_loss: -0.057758, loss_mean_cls: 3.359118, grad_norm: 1.652974
+Steps:   0%|          | 121/1000000 [00:33<67:49:02,  4.10it/s, grad_norm=1.65, loss_final=4.4, loss_mean=1.09, loss_mean_cls=3.36, proj_loss=-0.0578][2026-03-23 13:56:45] Step: 121, Training Logs: loss_final: 3.780295, loss_mean: 1.143866, proj_loss: -0.057413, loss_mean_cls: 2.693842, grad_norm: 2.330424
+Steps:   0%|          | 122/1000000 [00:34<67:51:29,  4.09it/s, grad_norm=2.33, loss_final=3.78, loss_mean=1.14, loss_mean_cls=2.69, proj_loss=-0.0574][2026-03-23 13:56:45] Step: 122, Training Logs: loss_final: 5.011144, loss_mean: 1.122681, proj_loss: -0.058837, loss_mean_cls: 3.947300, grad_norm: 1.509129
+Steps:   0%|          | 123/1000000 [00:34<67:51:03,  4.09it/s, grad_norm=1.51, loss_final=5.01, loss_mean=1.12, loss_mean_cls=3.95, proj_loss=-0.0588][2026-03-23 13:56:46] Step: 123, Training Logs: loss_final: 4.172338, loss_mean: 1.115492, proj_loss: -0.056144, loss_mean_cls: 3.112991, grad_norm: 1.528705
+Steps:   0%|          | 124/1000000 [00:34<67:50:03,  4.09it/s, grad_norm=1.53, loss_final=4.17, loss_mean=1.12, loss_mean_cls=3.11, proj_loss=-0.0561][2026-03-23 13:56:46] Step: 124, Training Logs: loss_final: 4.010628, loss_mean: 1.117229, proj_loss: -0.054621, loss_mean_cls: 2.948020, grad_norm: 1.329769
+Steps:   0%|          | 125/1000000 [00:34<67:48:41,  4.10it/s, grad_norm=1.33, loss_final=4.01, loss_mean=1.12, loss_mean_cls=2.95, proj_loss=-0.0546][2026-03-23 13:56:46] Step: 125, Training Logs: loss_final: 3.764182, loss_mean: 1.127925, proj_loss: -0.057976, loss_mean_cls: 2.694233, grad_norm: 1.674507
+Steps:   0%|          | 126/1000000 [00:35<67:47:01,  4.10it/s, grad_norm=1.67, loss_final=3.76, loss_mean=1.13, loss_mean_cls=2.69, proj_loss=-0.058][2026-03-23 13:56:46] Step: 126, Training Logs: loss_final: 4.371668, loss_mean: 1.097472, proj_loss: -0.055551, loss_mean_cls: 3.329747, grad_norm: 1.962917
+Steps:   0%|          | 127/1000000 [00:35<67:47:01,  4.10it/s, grad_norm=1.96, loss_final=4.37, loss_mean=1.1, loss_mean_cls=3.33, proj_loss=-0.0556][2026-03-23 13:56:47] Step: 127, Training Logs: loss_final: 4.417885, loss_mean: 1.097842, proj_loss: -0.057305, loss_mean_cls: 3.377348, grad_norm: 1.860904
+Steps:   0%|          | 128/1000000 [00:35<67:46:39,  4.10it/s, grad_norm=1.86, loss_final=4.42, loss_mean=1.1, loss_mean_cls=3.38, proj_loss=-0.0573][2026-03-23 13:56:47] Step: 128, Training Logs: loss_final: 4.017467, loss_mean: 1.103394, proj_loss: -0.057005, loss_mean_cls: 2.971077, grad_norm: 1.122871
+Steps:   0%|          | 129/1000000 [00:35<67:55:50,  4.09it/s, grad_norm=1.12, loss_final=4.02, loss_mean=1.1, loss_mean_cls=2.97, proj_loss=-0.057][2026-03-23 13:56:47] Step: 129, Training Logs: loss_final: 4.208836, loss_mean: 1.102054, proj_loss: -0.056320, loss_mean_cls: 3.163102, grad_norm: 1.851275
+Steps:   0%|          | 130/1000000 [00:36<67:52:57,  4.09it/s, grad_norm=1.85, loss_final=4.21, loss_mean=1.1, loss_mean_cls=3.16, proj_loss=-0.0563][2026-03-23 13:56:47] Step: 130, Training Logs: loss_final: 4.207572, loss_mean: 1.127353, proj_loss: -0.057998, loss_mean_cls: 3.138218, grad_norm: 1.432804
+Steps:   0%|          | 131/1000000 [00:36<67:51:54,  4.09it/s, grad_norm=1.43, loss_final=4.21, loss_mean=1.13, loss_mean_cls=3.14, proj_loss=-0.058][2026-03-23 13:56:48] Step: 131, Training Logs: loss_final: 4.999227, loss_mean: 1.065102, proj_loss: -0.057811, loss_mean_cls: 3.991935, grad_norm: 1.260865
+Steps:   0%|          | 132/1000000 [00:36<67:50:15,  4.09it/s, grad_norm=1.26, loss_final=5, loss_mean=1.07, loss_mean_cls=3.99, proj_loss=-0.0578][2026-03-23 13:56:48] Step: 132, Training Logs: loss_final: 4.958624, loss_mean: 1.080354, proj_loss: -0.055375, loss_mean_cls: 3.933645, grad_norm: 1.087305
+Steps:   0%|          | 133/1000000 [00:36<67:49:24,  4.10it/s, grad_norm=1.09, loss_final=4.96, loss_mean=1.08, loss_mean_cls=3.93, proj_loss=-0.0554][2026-03-23 13:56:48] Step: 133, Training Logs: loss_final: 4.336567, loss_mean: 1.074137, proj_loss: -0.057906, loss_mean_cls: 3.320336, grad_norm: 0.845532
+Steps:   0%|          | 134/1000000 [00:36<67:49:27,  4.09it/s, grad_norm=0.846, loss_final=4.34, loss_mean=1.07, loss_mean_cls=3.32, proj_loss=-0.0579][2026-03-23 13:56:48] Step: 134, Training Logs: loss_final: 4.113417, loss_mean: 1.101781, proj_loss: -0.055828, loss_mean_cls: 3.067464, grad_norm: 0.853130
+Steps:   0%|          | 135/1000000 [00:37<67:48:26,  4.10it/s, grad_norm=0.853, loss_final=4.11, loss_mean=1.1, loss_mean_cls=3.07, proj_loss=-0.0558][2026-03-23 13:56:48] Step: 135, Training Logs: loss_final: 4.451187, loss_mean: 1.071875, proj_loss: -0.057522, loss_mean_cls: 3.436834, grad_norm: 1.211643
+Steps:   0%|          | 136/1000000 [00:37<67:48:00,  4.10it/s, grad_norm=1.21, loss_final=4.45, loss_mean=1.07, loss_mean_cls=3.44, proj_loss=-0.0575][2026-03-23 13:56:49] Step: 136, Training Logs: loss_final: 3.470972, loss_mean: 1.105708, proj_loss: -0.057636, loss_mean_cls: 2.422900, grad_norm: 1.028398
+Steps:   0%|          | 137/1000000 [00:37<67:48:46,  4.10it/s, grad_norm=1.03, loss_final=3.47, loss_mean=1.11, loss_mean_cls=2.42, proj_loss=-0.0576][2026-03-23 13:56:49] Step: 137, Training Logs: loss_final: 4.167720, loss_mean: 1.116012, proj_loss: -0.056297, loss_mean_cls: 3.108005, grad_norm: 1.331365
+Steps:   0%|          | 138/1000000 [00:37<67:49:54,  4.09it/s, grad_norm=1.33, loss_final=4.17, loss_mean=1.12, loss_mean_cls=3.11, proj_loss=-0.0563][2026-03-23 13:56:49] Step: 138, Training Logs: loss_final: 3.969088, loss_mean: 1.080545, proj_loss: -0.058032, loss_mean_cls: 2.946575, grad_norm: 1.159432
+Steps:   0%|          | 139/1000000 [00:38<67:49:15,  4.10it/s, grad_norm=1.16, loss_final=3.97, loss_mean=1.08, loss_mean_cls=2.95, proj_loss=-0.058][2026-03-23 13:56:49] Step: 139, Training Logs: loss_final: 4.553603, loss_mean: 1.100756, proj_loss: -0.057699, loss_mean_cls: 3.510547, grad_norm: 1.379237
+Steps:   0%|          | 140/1000000 [00:38<67:48:29,  4.10it/s, grad_norm=1.38, loss_final=4.55, loss_mean=1.1, loss_mean_cls=3.51, proj_loss=-0.0577][2026-03-23 13:56:50] Step: 140, Training Logs: loss_final: 3.823996, loss_mean: 1.091382, proj_loss: -0.055826, loss_mean_cls: 2.788440, grad_norm: 1.269147
+Steps:   0%|          | 141/1000000 [00:38<67:49:20,  4.10it/s, grad_norm=1.27, loss_final=3.82, loss_mean=1.09, loss_mean_cls=2.79, proj_loss=-0.0558][2026-03-23 13:56:50] Step: 141, Training Logs: loss_final: 4.244388, loss_mean: 1.086618, proj_loss: -0.060178, loss_mean_cls: 3.217947, grad_norm: 1.224737
+Steps:   0%|          | 142/1000000 [00:38<67:48:47,  4.10it/s, grad_norm=1.22, loss_final=4.24, loss_mean=1.09, loss_mean_cls=3.22, proj_loss=-0.0602][2026-03-23 13:56:50] Step: 142, Training Logs: loss_final: 4.555122, loss_mean: 1.074355, proj_loss: -0.056447, loss_mean_cls: 3.537215, grad_norm: 1.288044
+Steps:   0%|          | 143/1000000 [00:39<67:49:15,  4.10it/s, grad_norm=1.29, loss_final=4.56, loss_mean=1.07, loss_mean_cls=3.54, proj_loss=-0.0564][2026-03-23 13:56:50] Step: 143, Training Logs: loss_final: 4.631051, loss_mean: 1.031162, proj_loss: -0.056307, loss_mean_cls: 3.656196, grad_norm: 0.875154
+Steps:   0%|          | 144/1000000 [00:39<67:47:56,  4.10it/s, grad_norm=0.875, loss_final=4.63, loss_mean=1.03, loss_mean_cls=3.66, proj_loss=-0.0563][2026-03-23 13:56:51] Step: 144, Training Logs: loss_final: 4.039268, loss_mean: 1.089356, proj_loss: -0.056316, loss_mean_cls: 3.006228, grad_norm: 1.155283
+Steps:   0%|          | 145/1000000 [00:39<67:55:59,  4.09it/s, grad_norm=1.16, loss_final=4.04, loss_mean=1.09, loss_mean_cls=3.01, proj_loss=-0.0563][2026-03-23 13:56:51] Step: 145, Training Logs: loss_final: 3.760608, loss_mean: 1.070431, proj_loss: -0.059299, loss_mean_cls: 2.749476, grad_norm: 1.032837
+Steps:   0%|          | 146/1000000 [00:39<67:53:40,  4.09it/s, grad_norm=1.03, loss_final=3.76, loss_mean=1.07, loss_mean_cls=2.75, proj_loss=-0.0593][2026-03-23 13:56:51] Step: 146, Training Logs: loss_final: 4.181468, loss_mean: 1.069318, proj_loss: -0.057231, loss_mean_cls: 3.169381, grad_norm: 1.739599
+Steps:   0%|          | 147/1000000 [00:40<87:04:35,  3.19it/s, grad_norm=1.74, loss_final=4.18, loss_mean=1.07, loss_mean_cls=3.17, proj_loss=-0.0572][2026-03-23 13:56:52] Step: 147, Training Logs: loss_final: 4.397898, loss_mean: 1.077711, proj_loss: -0.057317, loss_mean_cls: 3.377505, grad_norm: 1.324051
+Steps:   0%|          | 148/1000000 [00:40<81:18:29,  3.42it/s, grad_norm=1.32, loss_final=4.4, loss_mean=1.08, loss_mean_cls=3.38, proj_loss=-0.0573][2026-03-23 13:56:52] Step: 148, Training Logs: loss_final: 3.678877, loss_mean: 1.095331, proj_loss: -0.056643, loss_mean_cls: 2.640189, grad_norm: 0.957576
+Steps:   0%|          | 149/1000000 [00:40<77:21:07,  3.59it/s, grad_norm=0.958, loss_final=3.68, loss_mean=1.1, loss_mean_cls=2.64, proj_loss=-0.0566][2026-03-23 13:56:52] Step: 149, Training Logs: loss_final: 3.970488, loss_mean: 1.096237, proj_loss: -0.057982, loss_mean_cls: 2.932232, grad_norm: 0.751983
+Steps:   0%|          | 150/1000000 [00:41<74:28:13,  3.73it/s, grad_norm=0.752, loss_final=3.97, loss_mean=1.1, loss_mean_cls=2.93, proj_loss=-0.058][2026-03-23 13:56:52] Step: 150, Training Logs: loss_final: 3.589296, loss_mean: 1.085876, proj_loss: -0.059221, loss_mean_cls: 2.562641, grad_norm: 1.001571
+Steps:   0%|          | 151/1000000 [00:41<72:29:17,  3.83it/s, grad_norm=1, loss_final=3.59, loss_mean=1.09, loss_mean_cls=2.56, proj_loss=-0.0592][2026-03-23 13:56:53] Step: 151, Training Logs: loss_final: 3.809782, loss_mean: 1.041376, proj_loss: -0.057593, loss_mean_cls: 2.825998, grad_norm: 0.772958
+Steps:   0%|          | 152/1000000 [00:41<71:06:45,  3.91it/s, grad_norm=0.773, loss_final=3.81, loss_mean=1.04, loss_mean_cls=2.83, proj_loss=-0.0576][2026-03-23 13:56:53] Step: 152, Training Logs: loss_final: 3.756772, loss_mean: 1.077754, proj_loss: -0.055612, loss_mean_cls: 2.734629, grad_norm: 0.942414
+Steps:   0%|          | 153/1000000 [00:41<70:15:13,  3.95it/s, grad_norm=0.942, loss_final=3.76, loss_mean=1.08, loss_mean_cls=2.73, proj_loss=-0.0556][2026-03-23 13:56:53] Step: 153, Training Logs: loss_final: 3.815321, loss_mean: 1.091242, proj_loss: -0.055186, loss_mean_cls: 2.779264, grad_norm: 1.427716
+Steps:   0%|          | 154/1000000 [00:42<69:30:56,  4.00it/s, grad_norm=1.43, loss_final=3.82, loss_mean=1.09, loss_mean_cls=2.78, proj_loss=-0.0552][2026-03-23 13:56:53] Step: 154, Training Logs: loss_final: 4.484177, loss_mean: 1.070384, proj_loss: -0.056569, loss_mean_cls: 3.470362, grad_norm: 0.890748
+Steps:   0%|          | 155/1000000 [00:42<68:58:59,  4.03it/s, grad_norm=0.891, loss_final=4.48, loss_mean=1.07, loss_mean_cls=3.47, proj_loss=-0.0566][2026-03-23 13:56:54] Step: 155, Training Logs: loss_final: 3.992099, loss_mean: 1.094668, proj_loss: -0.056406, loss_mean_cls: 2.953837, grad_norm: 1.103738
+Steps:   0%|          | 156/1000000 [00:42<68:38:04,  4.05it/s, grad_norm=1.1, loss_final=3.99, loss_mean=1.09, loss_mean_cls=2.95, proj_loss=-0.0564][2026-03-23 13:56:54] Step: 156, Training Logs: loss_final: 4.835541, loss_mean: 1.066255, proj_loss: -0.057052, loss_mean_cls: 3.826337, grad_norm: 1.572935
+Steps:   0%|          | 157/1000000 [00:42<68:32:39,  4.05it/s, grad_norm=1.57, loss_final=4.84, loss_mean=1.07, loss_mean_cls=3.83, proj_loss=-0.0571][2026-03-23 13:56:54] Step: 157, Training Logs: loss_final: 4.422526, loss_mean: 1.037669, proj_loss: -0.057450, loss_mean_cls: 3.442307, grad_norm: 0.881763
+Steps:   0%|          | 158/1000000 [00:43<68:18:41,  4.07it/s, grad_norm=0.882, loss_final=4.42, loss_mean=1.04, loss_mean_cls=3.44, proj_loss=-0.0574][2026-03-23 13:56:54] Step: 158, Training Logs: loss_final: 3.933815, loss_mean: 1.094053, proj_loss: -0.058011, loss_mean_cls: 2.897773, grad_norm: 1.524001
+Steps:   0%|          | 159/1000000 [00:43<68:08:39,  4.08it/s, grad_norm=1.52, loss_final=3.93, loss_mean=1.09, loss_mean_cls=2.9, proj_loss=-0.058][2026-03-23 13:56:55] Step: 159, Training Logs: loss_final: 3.294493, loss_mean: 1.101539, proj_loss: -0.057312, loss_mean_cls: 2.250265, grad_norm: 0.862417
+Steps:   0%|          | 160/1000000 [00:43<68:00:56,  4.08it/s, grad_norm=0.862, loss_final=3.29, loss_mean=1.1, loss_mean_cls=2.25, proj_loss=-0.0573][2026-03-23 13:56:55] Step: 160, Training Logs: loss_final: 4.003653, loss_mean: 1.049939, proj_loss: -0.056768, loss_mean_cls: 3.010482, grad_norm: 1.406133
+Steps:   0%|          | 161/1000000 [00:43<68:05:23,  4.08it/s, grad_norm=1.41, loss_final=4, loss_mean=1.05, loss_mean_cls=3.01, proj_loss=-0.0568][2026-03-23 13:56:55] Step: 161, Training Logs: loss_final: 3.735500, loss_mean: 1.059312, proj_loss: -0.055975, loss_mean_cls: 2.732163, grad_norm: 1.257030
+Steps:   0%|          | 162/1000000 [00:44<67:59:39,  4.08it/s, grad_norm=1.26, loss_final=3.74, loss_mean=1.06, loss_mean_cls=2.73, proj_loss=-0.056][2026-03-23 13:56:55] Step: 162, Training Logs: loss_final: 4.134730, loss_mean: 1.051557, proj_loss: -0.055925, loss_mean_cls: 3.139099, grad_norm: 1.383831
+Steps:   0%|          | 163/1000000 [00:44<67:54:57,  4.09it/s, grad_norm=1.38, loss_final=4.13, loss_mean=1.05, loss_mean_cls=3.14, proj_loss=-0.0559][2026-03-23 13:56:56] Step: 163, Training Logs: loss_final: 3.976990, loss_mean: 1.080228, proj_loss: -0.055570, loss_mean_cls: 2.952332, grad_norm: 2.321836
+Steps:   0%|          | 164/1000000 [00:44<67:52:28,  4.09it/s, grad_norm=2.32, loss_final=3.98, loss_mean=1.08, loss_mean_cls=2.95, proj_loss=-0.0556][2026-03-23 13:56:56] Step: 164, Training Logs: loss_final: 4.920438, loss_mean: 1.053734, proj_loss: -0.057831, loss_mean_cls: 3.924535, grad_norm: 1.427572
+Steps:   0%|          | 165/1000000 [00:44<67:51:20,  4.09it/s, grad_norm=1.43, loss_final=4.92, loss_mean=1.05, loss_mean_cls=3.92, proj_loss=-0.0578][2026-03-23 13:56:56] Step: 165, Training Logs: loss_final: 4.075354, loss_mean: 1.102409, proj_loss: -0.055738, loss_mean_cls: 3.028682, grad_norm: 1.554521
+Steps:   0%|          | 166/1000000 [00:45<67:50:22,  4.09it/s, grad_norm=1.55, loss_final=4.08, loss_mean=1.1, loss_mean_cls=3.03, proj_loss=-0.0557][2026-03-23 13:56:56] Step: 166, Training Logs: loss_final: 3.970260, loss_mean: 1.052696, proj_loss: -0.054281, loss_mean_cls: 2.971844, grad_norm: 1.398542
+Steps:   0%|          | 167/1000000 [00:45<67:49:03,  4.10it/s, grad_norm=1.4, loss_final=3.97, loss_mean=1.05, loss_mean_cls=2.97, proj_loss=-0.0543][2026-03-23 13:56:57] Step: 167, Training Logs: loss_final: 4.506516, loss_mean: 1.058694, proj_loss: -0.055691, loss_mean_cls: 3.503513, grad_norm: 1.465548
+Steps:   0%|          | 168/1000000 [00:45<67:48:15,  4.10it/s, grad_norm=1.47, loss_final=4.51, loss_mean=1.06, loss_mean_cls=3.5, proj_loss=-0.0557][2026-03-23 13:56:57] Step: 168, Training Logs: loss_final: 3.774779, loss_mean: 1.070169, proj_loss: -0.056798, loss_mean_cls: 2.761408, grad_norm: 1.063631
+Steps:   0%|          | 169/1000000 [00:45<67:55:41,  4.09it/s, grad_norm=1.06, loss_final=3.77, loss_mean=1.07, loss_mean_cls=2.76, proj_loss=-0.0568][2026-03-23 13:56:57] Step: 169, Training Logs: loss_final: 3.799341, loss_mean: 1.072764, proj_loss: -0.057054, loss_mean_cls: 2.783631, grad_norm: 1.432508
+Steps:   0%|          | 170/1000000 [00:46<67:53:00,  4.09it/s, grad_norm=1.43, loss_final=3.8, loss_mean=1.07, loss_mean_cls=2.78, proj_loss=-0.0571][2026-03-23 13:56:57] Step: 170, Training Logs: loss_final: 3.950215, loss_mean: 1.050979, proj_loss: -0.058352, loss_mean_cls: 2.957588, grad_norm: 1.403183
+Steps:   0%|          | 171/1000000 [00:46<67:51:12,  4.09it/s, grad_norm=1.4, loss_final=3.95, loss_mean=1.05, loss_mean_cls=2.96, proj_loss=-0.0584][2026-03-23 13:56:58] Step: 171, Training Logs: loss_final: 4.198855, loss_mean: 1.049812, proj_loss: -0.056677, loss_mean_cls: 3.205721, grad_norm: 0.921758
+Steps:   0%|          | 172/1000000 [00:46<67:49:22,  4.09it/s, grad_norm=0.922, loss_final=4.2, loss_mean=1.05, loss_mean_cls=3.21, proj_loss=-0.0567][2026-03-23 13:56:58] Step: 172, Training Logs: loss_final: 3.441348, loss_mean: 1.076280, proj_loss: -0.057867, loss_mean_cls: 2.422936, grad_norm: 0.854744
+Steps:   0%|          | 173/1000000 [00:46<67:50:30,  4.09it/s, grad_norm=0.855, loss_final=3.44, loss_mean=1.08, loss_mean_cls=2.42, proj_loss=-0.0579][2026-03-23 13:56:58] Step: 173, Training Logs: loss_final: 3.798677, loss_mean: 1.085783, proj_loss: -0.057882, loss_mean_cls: 2.770777, grad_norm: 1.453923
+Steps:   0%|          | 174/1000000 [00:46<67:50:14,  4.09it/s, grad_norm=1.45, loss_final=3.8, loss_mean=1.09, loss_mean_cls=2.77, proj_loss=-0.0579][2026-03-23 13:56:58] Step: 174, Training Logs: loss_final: 3.112827, loss_mean: 1.064432, proj_loss: -0.058472, loss_mean_cls: 2.106868, grad_norm: 1.250231
+Steps:   0%|          | 175/1000000 [00:47<67:49:59,  4.09it/s, grad_norm=1.25, loss_final=3.11, loss_mean=1.06, loss_mean_cls=2.11, proj_loss=-0.0585][2026-03-23 13:56:58] Step: 175, Training Logs: loss_final: 3.708531, loss_mean: 1.080565, proj_loss: -0.055104, loss_mean_cls: 2.683070, grad_norm: 1.818153
+Steps:   0%|          | 176/1000000 [00:47<67:49:26,  4.09it/s, grad_norm=1.82, loss_final=3.71, loss_mean=1.08, loss_mean_cls=2.68, proj_loss=-0.0551][2026-03-23 13:56:59] Step: 176, Training Logs: loss_final: 4.057882, loss_mean: 1.041325, proj_loss: -0.057044, loss_mean_cls: 3.073601, grad_norm: 1.056386
+Steps:   0%|          | 177/1000000 [00:47<67:50:14,  4.09it/s, grad_norm=1.06, loss_final=4.06, loss_mean=1.04, loss_mean_cls=3.07, proj_loss=-0.057][2026-03-23 13:56:59] Step: 177, Training Logs: loss_final: 3.843677, loss_mean: 1.082743, proj_loss: -0.056343, loss_mean_cls: 2.817278, grad_norm: 1.454900
+Steps:   0%|          | 178/1000000 [00:47<67:50:45,  4.09it/s, grad_norm=1.45, loss_final=3.84, loss_mean=1.08, loss_mean_cls=2.82, proj_loss=-0.0563][2026-03-23 13:56:59] Step: 178, Training Logs: loss_final: 4.592722, loss_mean: 1.051988, proj_loss: -0.056379, loss_mean_cls: 3.597114, grad_norm: 1.553295
+Steps:   0%|          | 179/1000000 [00:48<67:51:26,  4.09it/s, grad_norm=1.55, loss_final=4.59, loss_mean=1.05, loss_mean_cls=3.6, proj_loss=-0.0564][2026-03-23 13:56:59] Step: 179, Training Logs: loss_final: 4.083874, loss_mean: 1.044562, proj_loss: -0.057760, loss_mean_cls: 3.097073, grad_norm: 1.207282
+Steps:   0%|          | 180/1000000 [00:48<67:50:22,  4.09it/s, grad_norm=1.21, loss_final=4.08, loss_mean=1.04, loss_mean_cls=3.1, proj_loss=-0.0578][2026-03-23 13:57:00] Step: 180, Training Logs: loss_final: 3.533881, loss_mean: 1.084617, proj_loss: -0.055471, loss_mean_cls: 2.504735, grad_norm: 1.384633
+Steps:   0%|          | 181/1000000 [00:48<67:49:31,  4.09it/s, grad_norm=1.38, loss_final=3.53, loss_mean=1.08, loss_mean_cls=2.5, proj_loss=-0.0555][2026-03-23 13:57:00] Step: 181, Training Logs: loss_final: 4.409495, loss_mean: 1.036472, proj_loss: -0.057444, loss_mean_cls: 3.430467, grad_norm: 1.297928
+Steps:   0%|          | 182/1000000 [00:48<67:48:57,  4.10it/s, grad_norm=1.3, loss_final=4.41, loss_mean=1.04, loss_mean_cls=3.43, proj_loss=-0.0574][2026-03-23 13:57:00] Step: 182, Training Logs: loss_final: 3.176256, loss_mean: 1.091783, proj_loss: -0.057242, loss_mean_cls: 2.141715, grad_norm: 1.965931
+Steps:   0%|          | 183/1000000 [00:49<67:49:28,  4.09it/s, grad_norm=1.97, loss_final=3.18, loss_mean=1.09, loss_mean_cls=2.14, proj_loss=-0.0572][2026-03-23 13:57:00] Step: 183, Training Logs: loss_final: 3.651900, loss_mean: 1.063701, proj_loss: -0.057272, loss_mean_cls: 2.645471, grad_norm: 1.260781
+Steps:   0%|          | 184/1000000 [00:49<67:49:34,  4.09it/s, grad_norm=1.26, loss_final=3.65, loss_mean=1.06, loss_mean_cls=2.65, proj_loss=-0.0573][2026-03-23 13:57:01] Step: 184, Training Logs: loss_final: 4.134811, loss_mean: 1.052191, proj_loss: -0.056823, loss_mean_cls: 3.139443, grad_norm: 1.083165
+Steps:   0%|          | 185/1000000 [00:49<67:54:27,  4.09it/s, grad_norm=1.08, loss_final=4.13, loss_mean=1.05, loss_mean_cls=3.14, proj_loss=-0.0568][2026-03-23 13:57:01] Step: 185, Training Logs: loss_final: 3.952762, loss_mean: 1.034234, proj_loss: -0.056918, loss_mean_cls: 2.975446, grad_norm: 1.370537
+Steps:   0%|          | 186/1000000 [00:49<67:53:04,  4.09it/s, grad_norm=1.37, loss_final=3.95, loss_mean=1.03, loss_mean_cls=2.98, proj_loss=-0.0569][2026-03-23 13:57:01] Step: 186, Training Logs: loss_final: 4.377052, loss_mean: 1.068180, proj_loss: -0.055129, loss_mean_cls: 3.364000, grad_norm: 2.047776
+Steps:   0%|          | 187/1000000 [00:50<67:51:19,  4.09it/s, grad_norm=2.05, loss_final=4.38, loss_mean=1.07, loss_mean_cls=3.36, proj_loss=-0.0551][2026-03-23 13:57:01] Step: 187, Training Logs: loss_final: 3.465180, loss_mean: 1.067420, proj_loss: -0.055022, loss_mean_cls: 2.452782, grad_norm: 1.094218
+Steps:   0%|          | 188/1000000 [00:50<67:50:38,  4.09it/s, grad_norm=1.09, loss_final=3.47, loss_mean=1.07, loss_mean_cls=2.45, proj_loss=-0.055][2026-03-23 13:57:02] Step: 188, Training Logs: loss_final: 4.428637, loss_mean: 1.075879, proj_loss: -0.056132, loss_mean_cls: 3.408890, grad_norm: 2.243081
+Steps:   0%|          | 189/1000000 [00:50<67:50:39,  4.09it/s, grad_norm=2.24, loss_final=4.43, loss_mean=1.08, loss_mean_cls=3.41, proj_loss=-0.0561][2026-03-23 13:57:02] Step: 189, Training Logs: loss_final: 3.874612, loss_mean: 1.059214, proj_loss: -0.055212, loss_mean_cls: 2.870610, grad_norm: 1.509933
+Steps:   0%|          | 190/1000000 [00:50<67:53:49,  4.09it/s, grad_norm=1.51, loss_final=3.87, loss_mean=1.06, loss_mean_cls=2.87, proj_loss=-0.0552][2026-03-23 13:57:02] Step: 190, Training Logs: loss_final: 3.696738, loss_mean: 1.066510, proj_loss: -0.058307, loss_mean_cls: 2.688535, grad_norm: 1.849224
+Steps:   0%|          | 191/1000000 [00:51<67:54:49,  4.09it/s, grad_norm=1.85, loss_final=3.7, loss_mean=1.07, loss_mean_cls=2.69, proj_loss=-0.0583][2026-03-23 13:57:02] Step: 191, Training Logs: loss_final: 4.354342, loss_mean: 1.020227, proj_loss: -0.055745, loss_mean_cls: 3.389860, grad_norm: 1.847939
+Steps:   0%|          | 192/1000000 [00:51<67:54:25,  4.09it/s, grad_norm=1.85, loss_final=4.35, loss_mean=1.02, loss_mean_cls=3.39, proj_loss=-0.0557][2026-03-23 13:57:03] Step: 192, Training Logs: loss_final: 4.409212, loss_mean: 1.048026, proj_loss: -0.056594, loss_mean_cls: 3.417780, grad_norm: 1.524102
+Steps:   0%|          | 193/1000000 [00:51<67:53:07,  4.09it/s, grad_norm=1.52, loss_final=4.41, loss_mean=1.05, loss_mean_cls=3.42, proj_loss=-0.0566][2026-03-23 13:57:03] Step: 193, Training Logs: loss_final: 3.909914, loss_mean: 1.034270, proj_loss: -0.058143, loss_mean_cls: 2.933788, grad_norm: 1.197792
+Steps:   0%|          | 194/1000000 [00:51<67:53:28,  4.09it/s, grad_norm=1.2, loss_final=3.91, loss_mean=1.03, loss_mean_cls=2.93, proj_loss=-0.0581][2026-03-23 13:57:03] Step: 194, Training Logs: loss_final: 3.518339, loss_mean: 1.061185, proj_loss: -0.053464, loss_mean_cls: 2.510617, grad_norm: 1.685875
+Steps:   0%|          | 195/1000000 [00:52<67:52:30,  4.09it/s, grad_norm=1.69, loss_final=3.52, loss_mean=1.06, loss_mean_cls=2.51, proj_loss=-0.0535][2026-03-23 13:57:03] Step: 195, Training Logs: loss_final: 4.144682, loss_mean: 1.069985, proj_loss: -0.056024, loss_mean_cls: 3.130722, grad_norm: 1.617648
+Steps:   0%|          | 196/1000000 [00:52<67:53:25,  4.09it/s, grad_norm=1.62, loss_final=4.14, loss_mean=1.07, loss_mean_cls=3.13, proj_loss=-0.056][2026-03-23 13:57:04] Step: 196, Training Logs: loss_final: 4.209203, loss_mean: 1.043093, proj_loss: -0.056968, loss_mean_cls: 3.223077, grad_norm: 1.561610
+Steps:   0%|          | 197/1000000 [00:52<67:51:41,  4.09it/s, grad_norm=1.56, loss_final=4.21, loss_mean=1.04, loss_mean_cls=3.22, proj_loss=-0.057][2026-03-23 13:57:04] Step: 197, Training Logs: loss_final: 3.986165, loss_mean: 1.048607, proj_loss: -0.058094, loss_mean_cls: 2.995652, grad_norm: 1.465502
+Steps:   0%|          | 198/1000000 [00:52<67:51:10,  4.09it/s, grad_norm=1.47, loss_final=3.99, loss_mean=1.05, loss_mean_cls=3, proj_loss=-0.0581][2026-03-23 13:57:04] Step: 198, Training Logs: loss_final: 4.109353, loss_mean: 1.057425, proj_loss: -0.057203, loss_mean_cls: 3.109131, grad_norm: 1.602210
+Steps:   0%|          | 199/1000000 [00:53<67:51:00,  4.09it/s, grad_norm=1.6, loss_final=4.11, loss_mean=1.06, loss_mean_cls=3.11, proj_loss=-0.0572][2026-03-23 13:57:04] Step: 199, Training Logs: loss_final: 3.827809, loss_mean: 1.060561, proj_loss: -0.056529, loss_mean_cls: 2.823776, grad_norm: 1.281930
+Steps:   0%|          | 200/1000000 [00:53<67:52:45,  4.09it/s, grad_norm=1.28, loss_final=3.83, loss_mean=1.06, loss_mean_cls=2.82, proj_loss=-0.0565][2026-03-23 13:57:05] Step: 200, Training Logs: loss_final: 4.859810, loss_mean: 1.018786, proj_loss: -0.056772, loss_mean_cls: 3.897796, grad_norm: 1.346983
+Steps:   0%|          | 201/1000000 [00:53<67:50:08,  4.09it/s, grad_norm=1.35, loss_final=4.86, loss_mean=1.02, loss_mean_cls=3.9, proj_loss=-0.0568][2026-03-23 13:57:05] Step: 201, Training Logs: loss_final: 3.945022, loss_mean: 1.039917, proj_loss: -0.058203, loss_mean_cls: 2.963308, grad_norm: 1.138654
+Steps:   0%|          | 202/1000000 [00:53<67:49:17,  4.09it/s, grad_norm=1.14, loss_final=3.95, loss_mean=1.04, loss_mean_cls=2.96, proj_loss=-0.0582][2026-03-23 13:57:05] Step: 202, Training Logs: loss_final: 3.795771, loss_mean: 1.045783, proj_loss: -0.058287, loss_mean_cls: 2.808275, grad_norm: 1.378415
+Steps:   0%|          | 203/1000000 [00:54<67:47:45,  4.10it/s, grad_norm=1.38, loss_final=3.8, loss_mean=1.05, loss_mean_cls=2.81, proj_loss=-0.0583][2026-03-23 13:57:05] Step: 203, Training Logs: loss_final: 3.873578, loss_mean: 1.057585, proj_loss: -0.057143, loss_mean_cls: 2.873136, grad_norm: 1.039413
+Steps:   0%|          | 204/1000000 [00:54<67:47:23,  4.10it/s, grad_norm=1.04, loss_final=3.87, loss_mean=1.06, loss_mean_cls=2.87, proj_loss=-0.0571][2026-03-23 13:57:06] Step: 204, Training Logs: loss_final: 3.869634, loss_mean: 1.025371, proj_loss: -0.059635, loss_mean_cls: 2.903898, grad_norm: 1.113446
+Steps:   0%|          | 205/1000000 [00:54<67:46:28,  4.10it/s, grad_norm=1.11, loss_final=3.87, loss_mean=1.03, loss_mean_cls=2.9, proj_loss=-0.0596][2026-03-23 13:57:06] Step: 205, Training Logs: loss_final: 4.516129, loss_mean: 1.026917, proj_loss: -0.059113, loss_mean_cls: 3.548324, grad_norm: 1.246919
+Steps:   0%|          | 206/1000000 [00:54<67:46:58,  4.10it/s, grad_norm=1.25, loss_final=4.52, loss_mean=1.03, loss_mean_cls=3.55, proj_loss=-0.0591][2026-03-23 13:57:06] Step: 206, Training Logs: loss_final: 3.870474, loss_mean: 1.033348, proj_loss: -0.057558, loss_mean_cls: 2.894683, grad_norm: 1.104687
+Steps:   0%|          | 207/1000000 [00:55<67:49:28,  4.09it/s, grad_norm=1.1, loss_final=3.87, loss_mean=1.03, loss_mean_cls=2.89, proj_loss=-0.0576][2026-03-23 13:57:06] Step: 207, Training Logs: loss_final: 3.657545, loss_mean: 1.036661, proj_loss: -0.054751, loss_mean_cls: 2.675634, grad_norm: 1.384344
+Steps:   0%|          | 208/1000000 [00:55<67:48:34,  4.10it/s, grad_norm=1.38, loss_final=3.66, loss_mean=1.04, loss_mean_cls=2.68, proj_loss=-0.0548][2026-03-23 13:57:07] Step: 208, Training Logs: loss_final: 3.776420, loss_mean: 1.025812, proj_loss: -0.057821, loss_mean_cls: 2.808429, grad_norm: 1.096370
+Steps:   0%|          | 209/1000000 [00:55<67:48:19,  4.10it/s, grad_norm=1.1, loss_final=3.78, loss_mean=1.03, loss_mean_cls=2.81, proj_loss=-0.0578][2026-03-23 13:57:07] Step: 209, Training Logs: loss_final: 3.761090, loss_mean: 1.062359, proj_loss: -0.056922, loss_mean_cls: 2.755653, grad_norm: 1.903223
+Steps:   0%|          | 210/1000000 [00:55<67:48:31,  4.10it/s, grad_norm=1.9, loss_final=3.76, loss_mean=1.06, loss_mean_cls=2.76, proj_loss=-0.0569][2026-03-23 13:57:07] Step: 210, Training Logs: loss_final: 3.906418, loss_mean: 1.041829, proj_loss: -0.057181, loss_mean_cls: 2.921770, grad_norm: 1.735096
+Steps:   0%|          | 211/1000000 [00:56<67:48:32,  4.10it/s, grad_norm=1.74, loss_final=3.91, loss_mean=1.04, loss_mean_cls=2.92, proj_loss=-0.0572][2026-03-23 13:57:07] Step: 211, Training Logs: loss_final: 3.858996, loss_mean: 1.023147, proj_loss: -0.059874, loss_mean_cls: 2.895722, grad_norm: 1.423884
+Steps:   0%|          | 212/1000000 [00:56<67:46:51,  4.10it/s, grad_norm=1.42, loss_final=3.86, loss_mean=1.02, loss_mean_cls=2.9, proj_loss=-0.0599][2026-03-23 13:57:08] Step: 212, Training Logs: loss_final: 4.049764, loss_mean: 1.020845, proj_loss: -0.057857, loss_mean_cls: 3.086776, grad_norm: 1.539841
+Steps:   0%|          | 213/1000000 [00:56<67:48:03,  4.10it/s, grad_norm=1.54, loss_final=4.05, loss_mean=1.02, loss_mean_cls=3.09, proj_loss=-0.0579][2026-03-23 13:57:08] Step: 213, Training Logs: loss_final: 3.666535, loss_mean: 1.064638, proj_loss: -0.058414, loss_mean_cls: 2.660310, grad_norm: 1.364981
+Steps:   0%|          | 214/1000000 [00:56<67:48:11,  4.10it/s, grad_norm=1.36, loss_final=3.67, loss_mean=1.06, loss_mean_cls=2.66, proj_loss=-0.0584][2026-03-23 13:57:08] Step: 214, Training Logs: loss_final: 4.115675, loss_mean: 1.022230, proj_loss: -0.055882, loss_mean_cls: 3.149326, grad_norm: 1.587428
+Steps:   0%|          | 215/1000000 [00:57<67:46:32,  4.10it/s, grad_norm=1.59, loss_final=4.12, loss_mean=1.02, loss_mean_cls=3.15, proj_loss=-0.0559][2026-03-23 13:57:08] Step: 215, Training Logs: loss_final: 4.290797, loss_mean: 1.006664, proj_loss: -0.057540, loss_mean_cls: 3.341673, grad_norm: 1.364294
+Steps:   0%|          | 216/1000000 [00:57<67:48:39,  4.10it/s, grad_norm=1.36, loss_final=4.29, loss_mean=1.01, loss_mean_cls=3.34, proj_loss=-0.0575][2026-03-23 13:57:09] Step: 216, Training Logs: loss_final: 4.299802, loss_mean: 1.033919, proj_loss: -0.057627, loss_mean_cls: 3.323509, grad_norm: 2.076446
+Steps:   0%|          | 217/1000000 [00:57<67:48:50,  4.10it/s, grad_norm=2.08, loss_final=4.3, loss_mean=1.03, loss_mean_cls=3.32, proj_loss=-0.0576][2026-03-23 13:57:09] Step: 217, Training Logs: loss_final: 4.032390, loss_mean: 1.038509, proj_loss: -0.059780, loss_mean_cls: 3.053661, grad_norm: 2.095842
+Steps:   0%|          | 218/1000000 [00:57<67:48:13,  4.10it/s, grad_norm=2.1, loss_final=4.03, loss_mean=1.04, loss_mean_cls=3.05, proj_loss=-0.0598][2026-03-23 13:57:09] Step: 218, Training Logs: loss_final: 4.852066, loss_mean: 0.991051, proj_loss: -0.055897, loss_mean_cls: 3.916911, grad_norm: 1.589446
+Steps:   0%|          | 219/1000000 [00:57<67:54:34,  4.09it/s, grad_norm=1.59, loss_final=4.85, loss_mean=0.991, loss_mean_cls=3.92, proj_loss=-0.0559][2026-03-23 13:57:09] Step: 219, Training Logs: loss_final: 3.452114, loss_mean: 1.032779, proj_loss: -0.057448, loss_mean_cls: 2.476782, grad_norm: 2.759307
+Steps:   0%|          | 220/1000000 [00:58<67:53:33,  4.09it/s, grad_norm=2.76, loss_final=3.45, loss_mean=1.03, loss_mean_cls=2.48, proj_loss=-0.0574][2026-03-23 13:57:09] Step: 220, Training Logs: loss_final: 4.474224, loss_mean: 1.012484, proj_loss: -0.056469, loss_mean_cls: 3.518209, grad_norm: 2.037073
+Steps:   0%|          | 221/1000000 [00:58<67:58:08,  4.09it/s, grad_norm=2.04, loss_final=4.47, loss_mean=1.01, loss_mean_cls=3.52, proj_loss=-0.0565][2026-03-23 13:57:10] Step: 221, Training Logs: loss_final: 3.619599, loss_mean: 1.027089, proj_loss: -0.055216, loss_mean_cls: 2.647726, grad_norm: 1.876838
+Steps:   0%|          | 222/1000000 [00:58<67:55:14,  4.09it/s, grad_norm=1.88, loss_final=3.62, loss_mean=1.03, loss_mean_cls=2.65, proj_loss=-0.0552][2026-03-23 13:57:10] Step: 222, Training Logs: loss_final: 3.956278, loss_mean: 0.994196, proj_loss: -0.058793, loss_mean_cls: 3.020875, grad_norm: 1.756858
+Steps:   0%|          | 223/1000000 [00:58<67:57:51,  4.09it/s, grad_norm=1.76, loss_final=3.96, loss_mean=0.994, loss_mean_cls=3.02, proj_loss=-0.0588][2026-03-23 13:57:10] Step: 223, Training Logs: loss_final: 3.396563, loss_mean: 1.048192, proj_loss: -0.056162, loss_mean_cls: 2.404532, grad_norm: 1.566004
+Steps:   0%|          | 224/1000000 [00:59<67:54:18,  4.09it/s, grad_norm=1.57, loss_final=3.4, loss_mean=1.05, loss_mean_cls=2.4, proj_loss=-0.0562][2026-03-23 13:57:10] Step: 224, Training Logs: loss_final: 4.041858, loss_mean: 1.026601, proj_loss: -0.059143, loss_mean_cls: 3.074400, grad_norm: 1.464198
+Steps:   0%|          | 225/1000000 [00:59<67:52:38,  4.09it/s, grad_norm=1.46, loss_final=4.04, loss_mean=1.03, loss_mean_cls=3.07, proj_loss=-0.0591][2026-03-23 13:57:11] Step: 225, Training Logs: loss_final: 3.787225, loss_mean: 1.008417, proj_loss: -0.057760, loss_mean_cls: 2.836567, grad_norm: 1.825020
+Steps:   0%|          | 226/1000000 [00:59<67:52:45,  4.09it/s, grad_norm=1.83, loss_final=3.79, loss_mean=1.01, loss_mean_cls=2.84, proj_loss=-0.0578][2026-03-23 13:57:11] Step: 226, Training Logs: loss_final: 3.892933, loss_mean: 1.035562, proj_loss: -0.058638, loss_mean_cls: 2.916008, grad_norm: 1.450337
+Steps:   0%|          | 227/1000000 [00:59<67:53:30,  4.09it/s, grad_norm=1.45, loss_final=3.89, loss_mean=1.04, loss_mean_cls=2.92, proj_loss=-0.0586][2026-03-23 13:57:11] Step: 227, Training Logs: loss_final: 4.605734, loss_mean: 1.015767, proj_loss: -0.058328, loss_mean_cls: 3.648294, grad_norm: 1.576961
+Steps:   0%|          | 228/1000000 [01:00<67:52:30,  4.09it/s, grad_norm=1.58, loss_final=4.61, loss_mean=1.02, loss_mean_cls=3.65, proj_loss=-0.0583][2026-03-23 13:57:11] Step: 228, Training Logs: loss_final: 3.988146, loss_mean: 1.020183, proj_loss: -0.057191, loss_mean_cls: 3.025154, grad_norm: 1.494521
+Steps:   0%|          | 229/1000000 [01:00<67:49:43,  4.09it/s, grad_norm=1.49, loss_final=3.99, loss_mean=1.02, loss_mean_cls=3.03, proj_loss=-0.0572][2026-03-23 13:57:12] Step: 229, Training Logs: loss_final: 4.138639, loss_mean: 1.001511, proj_loss: -0.058552, loss_mean_cls: 3.195680, grad_norm: 1.869578
+Steps:   0%|          | 230/1000000 [01:00<67:48:23,  4.10it/s, grad_norm=1.87, loss_final=4.14, loss_mean=1, loss_mean_cls=3.2, proj_loss=-0.0586][2026-03-23 13:57:12] Step: 230, Training Logs: loss_final: 3.587549, loss_mean: 1.026086, proj_loss: -0.057275, loss_mean_cls: 2.618739, grad_norm: 1.361713
+Steps:   0%|          | 231/1000000 [01:00<67:50:24,  4.09it/s, grad_norm=1.36, loss_final=3.59, loss_mean=1.03, loss_mean_cls=2.62, proj_loss=-0.0573][2026-03-23 13:57:12] Step: 231, Training Logs: loss_final: 4.299013, loss_mean: 1.019667, proj_loss: -0.057679, loss_mean_cls: 3.337025, grad_norm: 2.507548
+Steps:   0%|          | 232/1000000 [01:01<67:49:48,  4.09it/s, grad_norm=2.51, loss_final=4.3, loss_mean=1.02, loss_mean_cls=3.34, proj_loss=-0.0577][2026-03-23 13:57:12] Step: 232, Training Logs: loss_final: 3.918818, loss_mean: 1.041992, proj_loss: -0.057748, loss_mean_cls: 2.934574, grad_norm: 1.768478
+Steps:   0%|          | 233/1000000 [01:01<67:48:26,  4.10it/s, grad_norm=1.77, loss_final=3.92, loss_mean=1.04, loss_mean_cls=2.93, proj_loss=-0.0577][2026-03-23 13:57:13] Step: 233, Training Logs: loss_final: 4.135336, loss_mean: 1.012130, proj_loss: -0.055806, loss_mean_cls: 3.179013, grad_norm: 2.010638
+Steps:   0%|          | 234/1000000 [01:01<67:48:16,  4.10it/s, grad_norm=2.01, loss_final=4.14, loss_mean=1.01, loss_mean_cls=3.18, proj_loss=-0.0558][2026-03-23 13:57:13] Step: 234, Training Logs: loss_final: 4.622794, loss_mean: 1.008856, proj_loss: -0.059102, loss_mean_cls: 3.673040, grad_norm: 1.824655
+Steps:   0%|          | 235/1000000 [01:01<67:48:42,  4.10it/s, grad_norm=1.82, loss_final=4.62, loss_mean=1.01, loss_mean_cls=3.67, proj_loss=-0.0591][2026-03-23 13:57:13] Step: 235, Training Logs: loss_final: 4.330180, loss_mean: 1.019101, proj_loss: -0.054455, loss_mean_cls: 3.365534, grad_norm: 1.865192
+Steps:   0%|          | 236/1000000 [01:02<67:47:59,  4.10it/s, grad_norm=1.87, loss_final=4.33, loss_mean=1.02, loss_mean_cls=3.37, proj_loss=-0.0545][2026-03-23 13:57:13] Step: 236, Training Logs: loss_final: 4.471676, loss_mean: 1.005463, proj_loss: -0.057791, loss_mean_cls: 3.524004, grad_norm: 1.992857
+Steps:   0%|          | 237/1000000 [01:02<67:48:33,  4.10it/s, grad_norm=1.99, loss_final=4.47, loss_mean=1.01, loss_mean_cls=3.52, proj_loss=-0.0578][2026-03-23 13:57:14] Step: 237, Training Logs: loss_final: 4.708742, loss_mean: 1.007163, proj_loss: -0.057294, loss_mean_cls: 3.758873, grad_norm: 1.696959
+Steps:   0%|          | 238/1000000 [01:02<67:48:46,  4.10it/s, grad_norm=1.7, loss_final=4.71, loss_mean=1.01, loss_mean_cls=3.76, proj_loss=-0.0573][2026-03-23 13:57:14] Step: 238, Training Logs: loss_final: 4.700453, loss_mean: 1.011082, proj_loss: -0.057899, loss_mean_cls: 3.747270, grad_norm: 2.577658
+Steps:   0%|          | 239/1000000 [01:02<67:48:18,  4.10it/s, grad_norm=2.58, loss_final=4.7, loss_mean=1.01, loss_mean_cls=3.75, proj_loss=-0.0579][2026-03-23 13:57:14] Step: 239, Training Logs: loss_final: 4.519000, loss_mean: 0.999633, proj_loss: -0.058458, loss_mean_cls: 3.577825, grad_norm: 2.147466
+Steps:   0%|          | 240/1000000 [01:03<67:48:55,  4.10it/s, grad_norm=2.15, loss_final=4.52, loss_mean=1, loss_mean_cls=3.58, proj_loss=-0.0585][2026-03-23 13:57:14] Step: 240, Training Logs: loss_final: 3.874477, loss_mean: 1.005100, proj_loss: -0.059289, loss_mean_cls: 2.928666, grad_norm: 2.779819
+Steps:   0%|          | 241/1000000 [01:03<67:50:12,  4.09it/s, grad_norm=2.78, loss_final=3.87, loss_mean=1.01, loss_mean_cls=2.93, proj_loss=-0.0593][2026-03-23 13:57:15] Step: 241, Training Logs: loss_final: 4.244301, loss_mean: 0.978814, proj_loss: -0.057319, loss_mean_cls: 3.322806, grad_norm: 1.934856
+Steps:   0%|          | 242/1000000 [01:03<67:50:20,  4.09it/s, grad_norm=1.93, loss_final=4.24, loss_mean=0.979, loss_mean_cls=3.32, proj_loss=-0.0573][2026-03-23 13:57:15] Step: 242, Training Logs: loss_final: 4.099132, loss_mean: 1.017777, proj_loss: -0.056814, loss_mean_cls: 3.138168, grad_norm: 2.223382
+Steps:   0%|          | 243/1000000 [01:03<67:49:22,  4.09it/s, grad_norm=2.22, loss_final=4.1, loss_mean=1.02, loss_mean_cls=3.14, proj_loss=-0.0568][2026-03-23 13:57:15] Step: 243, Training Logs: loss_final: 3.643645, loss_mean: 1.043657, proj_loss: -0.057582, loss_mean_cls: 2.657569, grad_norm: 1.876094
+Steps:   0%|          | 244/1000000 [01:04<67:48:22,  4.10it/s, grad_norm=1.88, loss_final=3.64, loss_mean=1.04, loss_mean_cls=2.66, proj_loss=-0.0576][2026-03-23 13:57:15] Step: 244, Training Logs: loss_final: 3.791383, loss_mean: 1.039013, proj_loss: -0.057821, loss_mean_cls: 2.810191, grad_norm: 2.101933
+Steps:   0%|          | 245/1000000 [01:04<67:48:03,  4.10it/s, grad_norm=2.1, loss_final=3.79, loss_mean=1.04, loss_mean_cls=2.81, proj_loss=-0.0578][2026-03-23 13:57:16] Step: 245, Training Logs: loss_final: 4.122998, loss_mean: 1.003735, proj_loss: -0.056978, loss_mean_cls: 3.176242, grad_norm: 1.748520
+Steps:   0%|          | 246/1000000 [01:04<67:47:44,  4.10it/s, grad_norm=1.75, loss_final=4.12, loss_mean=1, loss_mean_cls=3.18, proj_loss=-0.057][2026-03-23 13:57:16] Step: 246, Training Logs: loss_final: 3.975110, loss_mean: 1.031366, proj_loss: -0.058318, loss_mean_cls: 3.002061, grad_norm: 2.088799
+Steps:   0%|          | 247/1000000 [01:04<67:49:45,  4.09it/s, grad_norm=2.09, loss_final=3.98, loss_mean=1.03, loss_mean_cls=3, proj_loss=-0.0583][2026-03-23 13:57:16] Step: 247, Training Logs: loss_final: 3.795434, loss_mean: 1.016577, proj_loss: -0.059947, loss_mean_cls: 2.838804, grad_norm: 1.548398
+Steps:   0%|          | 248/1000000 [01:05<67:48:52,  4.10it/s, grad_norm=1.55, loss_final=3.8, loss_mean=1.02, loss_mean_cls=2.84, proj_loss=-0.0599][2026-03-23 13:57:16] Step: 248, Training Logs: loss_final: 3.547227, loss_mean: 1.007156, proj_loss: -0.057843, loss_mean_cls: 2.597914, grad_norm: 2.097709
+Steps:   0%|          | 249/1000000 [01:05<67:47:51,  4.10it/s, grad_norm=2.1, loss_final=3.55, loss_mean=1.01, loss_mean_cls=2.6, proj_loss=-0.0578][2026-03-23 13:57:17] Step: 249, Training Logs: loss_final: 3.670953, loss_mean: 1.005734, proj_loss: -0.056730, loss_mean_cls: 2.721948, grad_norm: 1.993682
+Steps:   0%|          | 250/1000000 [01:05<67:48:32,  4.10it/s, grad_norm=1.99, loss_final=3.67, loss_mean=1.01, loss_mean_cls=2.72, proj_loss=-0.0567][2026-03-23 13:57:17] Step: 250, Training Logs: loss_final: 3.783008, loss_mean: 1.009454, proj_loss: -0.057429, loss_mean_cls: 2.830983, grad_norm: 2.015732
+Steps:   0%|          | 251/1000000 [01:05<67:49:52,  4.09it/s, grad_norm=2.02, loss_final=3.78, loss_mean=1.01, loss_mean_cls=2.83, proj_loss=-0.0574][2026-03-23 13:57:17] Step: 251, Training Logs: loss_final: 3.917508, loss_mean: 1.015374, proj_loss: -0.057006, loss_mean_cls: 2.959140, grad_norm: 1.911892
+Steps:   0%|          | 252/1000000 [01:06<67:48:52,  4.10it/s, grad_norm=1.91, loss_final=3.92, loss_mean=1.02, loss_mean_cls=2.96, proj_loss=-0.057][2026-03-23 13:57:17] Step: 252, Training Logs: loss_final: 4.345321, loss_mean: 0.999146, proj_loss: -0.059442, loss_mean_cls: 3.405616, grad_norm: 2.544859
+Steps:   0%|          | 253/1000000 [01:06<67:48:39,  4.10it/s, grad_norm=2.54, loss_final=4.35, loss_mean=0.999, loss_mean_cls=3.41, proj_loss=-0.0594][2026-03-23 13:57:18] Step: 253, Training Logs: loss_final: 3.264407, loss_mean: 1.012033, proj_loss: -0.060338, loss_mean_cls: 2.312712, grad_norm: 2.002759
+Steps:   0%|          | 254/1000000 [01:06<67:47:32,  4.10it/s, grad_norm=2, loss_final=3.26, loss_mean=1.01, loss_mean_cls=2.31, proj_loss=-0.0603][2026-03-23 13:57:18] Step: 254, Training Logs: loss_final: 3.680009, loss_mean: 1.009429, proj_loss: -0.057669, loss_mean_cls: 2.728250, grad_norm: 2.470939
+Steps:   0%|          | 255/1000000 [01:06<67:46:55,  4.10it/s, grad_norm=2.47, loss_final=3.68, loss_mean=1.01, loss_mean_cls=2.73, proj_loss=-0.0577][2026-03-23 13:57:18] Step: 255, Training Logs: loss_final: 3.617963, loss_mean: 1.040573, proj_loss: -0.058244, loss_mean_cls: 2.635633, grad_norm: 2.220286
+Steps:   0%|          | 256/1000000 [01:07<67:48:18,  4.10it/s, grad_norm=2.22, loss_final=3.62, loss_mean=1.04, loss_mean_cls=2.64, proj_loss=-0.0582][2026-03-23 13:57:18] Step: 256, Training Logs: loss_final: 4.080931, loss_mean: 0.997066, proj_loss: -0.058532, loss_mean_cls: 3.142397, grad_norm: 2.140220
+Steps:   0%|          | 257/1000000 [01:07<67:47:49,  4.10it/s, grad_norm=2.14, loss_final=4.08, loss_mean=0.997, loss_mean_cls=3.14, proj_loss=-0.0585][2026-03-23 13:57:19] Step: 257, Training Logs: loss_final: 3.699068, loss_mean: 1.000059, proj_loss: -0.058102, loss_mean_cls: 2.757111, grad_norm: 1.981827
+Steps:   0%|          | 258/1000000 [01:07<67:46:24,  4.10it/s, grad_norm=1.98, loss_final=3.7, loss_mean=1, loss_mean_cls=2.76, proj_loss=-0.0581][2026-03-23 13:57:19] Step: 258, Training Logs: loss_final: 3.992317, loss_mean: 1.029925, proj_loss: -0.058207, loss_mean_cls: 3.020598, grad_norm: 2.269195
+Steps:   0%|          | 259/1000000 [01:07<67:46:06,  4.10it/s, grad_norm=2.27, loss_final=3.99, loss_mean=1.03, loss_mean_cls=3.02, proj_loss=-0.0582][2026-03-23 13:57:19] Step: 259, Training Logs: loss_final: 3.969361, loss_mean: 1.005176, proj_loss: -0.055455, loss_mean_cls: 3.019640, grad_norm: 2.430135
+Steps:   0%|          | 260/1000000 [01:08<67:47:12,  4.10it/s, grad_norm=2.43, loss_final=3.97, loss_mean=1.01, loss_mean_cls=3.02, proj_loss=-0.0555][2026-03-23 13:57:19] Step: 260, Training Logs: loss_final: 3.946629, loss_mean: 0.986893, proj_loss: -0.055241, loss_mean_cls: 3.014977, grad_norm: 2.162997
+Steps:   0%|          | 261/1000000 [01:08<79:04:37,  3.51it/s, grad_norm=2.16, loss_final=3.95, loss_mean=0.987, loss_mean_cls=3.01, proj_loss=-0.0552][2026-03-23 13:57:20] Step: 261, Training Logs: loss_final: 3.920058, loss_mean: 0.984585, proj_loss: -0.057787, loss_mean_cls: 2.993260, grad_norm: 2.114045
+Steps:   0%|          | 262/1000000 [01:08<78:26:04,  3.54it/s, grad_norm=2.11, loss_final=3.92, loss_mean=0.985, loss_mean_cls=2.99, proj_loss=-0.0578][2026-03-23 13:57:20] Step: 262, Training Logs: loss_final: 4.358081, loss_mean: 0.977391, proj_loss: -0.057046, loss_mean_cls: 3.437736, grad_norm: 2.065785
+Steps:   0%|          | 263/1000000 [01:08<75:13:43,  3.69it/s, grad_norm=2.07, loss_final=4.36, loss_mean=0.977, loss_mean_cls=3.44, proj_loss=-0.057][2026-03-23 13:57:20] Step: 263, Training Logs: loss_final: 3.840959, loss_mean: 0.999257, proj_loss: -0.056460, loss_mean_cls: 2.898163, grad_norm: 2.413785
+Steps:   0%|          | 264/1000000 [01:09<72:59:25,  3.80it/s, grad_norm=2.41, loss_final=3.84, loss_mean=0.999, loss_mean_cls=2.9, proj_loss=-0.0565][2026-03-23 13:57:20] Step: 264, Training Logs: loss_final: 3.940903, loss_mean: 1.023376, proj_loss: -0.057920, loss_mean_cls: 2.975446, grad_norm: 1.844397
+Steps:   0%|          | 265/1000000 [01:09<71:26:14,  3.89it/s, grad_norm=1.84, loss_final=3.94, loss_mean=1.02, loss_mean_cls=2.98, proj_loss=-0.0579][2026-03-23 13:57:21] Step: 265, Training Logs: loss_final: 4.206155, loss_mean: 0.981618, proj_loss: -0.059432, loss_mean_cls: 3.283970, grad_norm: 2.069561
+Steps:   0%|          | 266/1000000 [01:09<70:20:03,  3.95it/s, grad_norm=2.07, loss_final=4.21, loss_mean=0.982, loss_mean_cls=3.28, proj_loss=-0.0594][2026-03-23 13:57:21] Step: 266, Training Logs: loss_final: 3.804409, loss_mean: 0.990602, proj_loss: -0.060707, loss_mean_cls: 2.874513, grad_norm: 1.578623
+Steps:   0%|          | 267/1000000 [01:09<69:36:48,  3.99it/s, grad_norm=1.58, loss_final=3.8, loss_mean=0.991, loss_mean_cls=2.87, proj_loss=-0.0607][2026-03-23 13:57:21] Step: 267, Training Logs: loss_final: 4.086649, loss_mean: 1.005726, proj_loss: -0.058538, loss_mean_cls: 3.139461, grad_norm: 2.692237
+Steps:   0%|          | 268/1000000 [01:10<69:04:32,  4.02it/s, grad_norm=2.69, loss_final=4.09, loss_mean=1.01, loss_mean_cls=3.14, proj_loss=-0.0585][2026-03-23 13:57:21] Step: 268, Training Logs: loss_final: 3.484965, loss_mean: 1.013161, proj_loss: -0.056986, loss_mean_cls: 2.528790, grad_norm: 2.118618
+Steps:   0%|          | 269/1000000 [01:10<69:49:28,  3.98it/s, grad_norm=2.12, loss_final=3.48, loss_mean=1.01, loss_mean_cls=2.53, proj_loss=-0.057][2026-03-23 13:57:22] Step: 269, Training Logs: loss_final: 3.312028, loss_mean: 1.008822, proj_loss: -0.058315, loss_mean_cls: 2.361521, grad_norm: 1.997808
+Steps:   0%|          | 270/1000000 [01:10<68:55:10,  4.03it/s, grad_norm=2, loss_final=3.31, loss_mean=1.01, loss_mean_cls=2.36, proj_loss=-0.0583][2026-03-23 13:57:22] Step: 270, Training Logs: loss_final: 3.996657, loss_mean: 1.009214, proj_loss: -0.058962, loss_mean_cls: 3.046405, grad_norm: 2.755926
+Steps:   0%|          | 271/1000000 [01:10<68:39:18,  4.04it/s, grad_norm=2.76, loss_final=4, loss_mean=1.01, loss_mean_cls=3.05, proj_loss=-0.059][2026-03-23 13:57:22] Step: 271, Training Logs: loss_final: 4.391388, loss_mean: 1.024230, proj_loss: -0.058988, loss_mean_cls: 3.426146, grad_norm: 2.096098
+Steps:   0%|          | 272/1000000 [01:11<68:23:53,  4.06it/s, grad_norm=2.1, loss_final=4.39, loss_mean=1.02, loss_mean_cls=3.43, proj_loss=-0.059][2026-03-23 13:57:22] Step: 272, Training Logs: loss_final: 3.920087, loss_mean: 1.024717, proj_loss: -0.059831, loss_mean_cls: 2.955201, grad_norm: 2.465345
+Steps:   0%|          | 273/1000000 [01:11<68:12:32,  4.07it/s, grad_norm=2.47, loss_final=3.92, loss_mean=1.02, loss_mean_cls=2.96, proj_loss=-0.0598][2026-03-23 13:57:23] Step: 273, Training Logs: loss_final: 3.677792, loss_mean: 1.019457, proj_loss: -0.056399, loss_mean_cls: 2.714733, grad_norm: 2.047685
+Steps:   0%|          | 274/1000000 [01:11<68:06:48,  4.08it/s, grad_norm=2.05, loss_final=3.68, loss_mean=1.02, loss_mean_cls=2.71, proj_loss=-0.0564][2026-03-23 13:57:23] Step: 274, Training Logs: loss_final: 3.611300, loss_mean: 1.009368, proj_loss: -0.057628, loss_mean_cls: 2.659561, grad_norm: 1.780220
+Steps:   0%|          | 275/1000000 [01:11<68:01:32,  4.08it/s, grad_norm=1.78, loss_final=3.61, loss_mean=1.01, loss_mean_cls=2.66, proj_loss=-0.0576][2026-03-23 13:57:23] Step: 275, Training Logs: loss_final: 3.762297, loss_mean: 1.000715, proj_loss: -0.057146, loss_mean_cls: 2.818728, grad_norm: 2.212818
+Steps:   0%|          | 276/1000000 [01:12<67:56:22,  4.09it/s, grad_norm=2.21, loss_final=3.76, loss_mean=1, loss_mean_cls=2.82, proj_loss=-0.0571][2026-03-23 13:57:23] Step: 276, Training Logs: loss_final: 4.110180, loss_mean: 1.001855, proj_loss: -0.058339, loss_mean_cls: 3.166663, grad_norm: 3.446807
+Steps:   0%|          | 277/1000000 [01:12<67:54:20,  4.09it/s, grad_norm=3.45, loss_final=4.11, loss_mean=1, loss_mean_cls=3.17, proj_loss=-0.0583][2026-03-23 13:57:24] Step: 277, Training Logs: loss_final: 3.589436, loss_mean: 1.014081, proj_loss: -0.058753, loss_mean_cls: 2.634108, grad_norm: 1.866073
+Steps:   0%|          | 278/1000000 [01:12<67:51:55,  4.09it/s, grad_norm=1.87, loss_final=3.59, loss_mean=1.01, loss_mean_cls=2.63, proj_loss=-0.0588][2026-03-23 13:57:24] Step: 278, Training Logs: loss_final: 3.909306, loss_mean: 1.036322, proj_loss: -0.057491, loss_mean_cls: 2.930474, grad_norm: 3.416674
+Steps:   0%|          | 279/1000000 [01:12<67:51:48,  4.09it/s, grad_norm=3.42, loss_final=3.91, loss_mean=1.04, loss_mean_cls=2.93, proj_loss=-0.0575][2026-03-23 13:57:24] Step: 279, Training Logs: loss_final: 3.635005, loss_mean: 1.021109, proj_loss: -0.058334, loss_mean_cls: 2.672229, grad_norm: 2.773069
+Steps:   0%|          | 280/1000000 [01:13<67:50:10,  4.09it/s, grad_norm=2.77, loss_final=3.64, loss_mean=1.02, loss_mean_cls=2.67, proj_loss=-0.0583][2026-03-23 13:57:24] Step: 280, Training Logs: loss_final: 4.731202, loss_mean: 0.976675, proj_loss: -0.056143, loss_mean_cls: 3.810670, grad_norm: 2.045803
+Steps:   0%|          | 281/1000000 [01:13<67:48:49,  4.10it/s, grad_norm=2.05, loss_final=4.73, loss_mean=0.977, loss_mean_cls=3.81, proj_loss=-0.0561][2026-03-23 13:57:25] Step: 281, Training Logs: loss_final: 4.352121, loss_mean: 1.020213, proj_loss: -0.057039, loss_mean_cls: 3.388947, grad_norm: 3.085879
+Steps:   0%|          | 282/1000000 [01:13<67:48:21,  4.10it/s, grad_norm=3.09, loss_final=4.35, loss_mean=1.02, loss_mean_cls=3.39, proj_loss=-0.057][2026-03-23 13:57:25] Step: 282, Training Logs: loss_final: 3.645599, loss_mean: 1.009737, proj_loss: -0.056890, loss_mean_cls: 2.692752, grad_norm: 2.743040
+Steps:   0%|          | 283/1000000 [01:13<67:47:40,  4.10it/s, grad_norm=2.74, loss_final=3.65, loss_mean=1.01, loss_mean_cls=2.69, proj_loss=-0.0569][2026-03-23 13:57:25] Step: 283, Training Logs: loss_final: 3.621672, loss_mean: 1.014742, proj_loss: -0.058281, loss_mean_cls: 2.665210, grad_norm: 2.332768
+Steps:   0%|          | 284/1000000 [01:14<67:48:17,  4.10it/s, grad_norm=2.33, loss_final=3.62, loss_mean=1.01, loss_mean_cls=2.67, proj_loss=-0.0583][2026-03-23 13:57:25] Step: 284, Training Logs: loss_final: 3.634438, loss_mean: 1.014846, proj_loss: -0.058438, loss_mean_cls: 2.678030, grad_norm: 2.844196
+Steps:   0%|          | 285/1000000 [01:14<67:48:15,  4.10it/s, grad_norm=2.84, loss_final=3.63, loss_mean=1.01, loss_mean_cls=2.68, proj_loss=-0.0584][2026-03-23 13:57:26] Step: 285, Training Logs: loss_final: 3.942388, loss_mean: 0.986733, proj_loss: -0.056891, loss_mean_cls: 3.012546, grad_norm: 2.134713
+Steps:   0%|          | 286/1000000 [01:14<67:47:40,  4.10it/s, grad_norm=2.13, loss_final=3.94, loss_mean=0.987, loss_mean_cls=3.01, proj_loss=-0.0569][2026-03-23 13:57:26] Step: 286, Training Logs: loss_final: 3.926682, loss_mean: 1.011870, proj_loss: -0.057176, loss_mean_cls: 2.971987, grad_norm: 2.768237
+Steps:   0%|          | 287/1000000 [01:14<67:46:57,  4.10it/s, grad_norm=2.77, loss_final=3.93, loss_mean=1.01, loss_mean_cls=2.97, proj_loss=-0.0572][2026-03-23 13:57:26] Step: 287, Training Logs: loss_final: 4.047902, loss_mean: 1.002864, proj_loss: -0.059520, loss_mean_cls: 3.104558, grad_norm: 2.354553
+Steps:   0%|          | 287/1000000 [01:14<67:46:57,  4.10it/s, grad_norm=2.35, loss_final=4.05, loss_mean=1, loss_mean_cls=3.1, proj_loss=-0.0595]

REG/wandb/run-20260323_135607-zue1y2ba/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

REG/wandb/run-20260323_135607-zue1y2ba/files/wandb-metadata.json

@@ -0,0 +1,101 @@
+{
+  "os":  "Linux-5.15.0-94-generic-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-03-23T05:56:07.858187Z",
+  "args":  [
+    "--report-to",
+    "wandb",
+    "--allow-tf32",
+    "--mixed-precision",
+    "bf16",
+    "--seed",
+    "0",
+    "--path-type",
+    "linear",
+    "--prediction",
+    "v",
+    "--weighting",
+    "uniform",
+    "--model",
+    "SiT-XL/2",
+    "--enc-type",
+    "dinov2-vit-b",
+    "--encoder-depth",
+    "8",
+    "--proj-coeff",
+    "0.5",
+    "--output-dir",
+    "exps",
+    "--exp-name",
+    "jsflow-experiment-0.75",
+    "--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",
+    "--learning-rate",
+    "0.00005",
+    "--t-c",
+    "0.75",
+    "--cls",
+    "0.2",
+    "--ot-cls"
+  ],
+  "program":  "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/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/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":  "357568126976"
+    }
+  },
+  "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":  "l4vui4vfnl881ctol25fj9y70t6im9l9"
+}

REG/wandb/run-20260323_135607-zue1y2ba/logs/debug-internal.log

@@ -0,0 +1,6 @@
+{"time":"2026-03-23T13:56:08.183465712+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-23T13:56:10.661719633+08:00","level":"INFO","msg":"stream: created new stream","id":"zue1y2ba"}
+{"time":"2026-03-23T13:56:10.661931952+08:00","level":"INFO","msg":"handler: started","stream_id":"zue1y2ba"}
+{"time":"2026-03-23T13:56:10.662874633+08:00","level":"INFO","msg":"stream: started","id":"zue1y2ba"}
+{"time":"2026-03-23T13:56:10.662895027+08:00","level":"INFO","msg":"writer: started","stream_id":"zue1y2ba"}
+{"time":"2026-03-23T13:56:10.662918583+08:00","level":"INFO","msg":"sender: started","stream_id":"zue1y2ba"}

REG/wandb/run-20260323_135607-zue1y2ba/logs/debug.log

@@ -0,0 +1,20 @@
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_setup.py:_flush():81] Configure stats pid to 397944
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_init.py:setup_run_log_directory():717] Logging user logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260323_135607-zue1y2ba/logs/debug.log
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260323_135607-zue1y2ba/logs/debug-internal.log
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_init.py:init():844] calling init triggers
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-23 13:56:07,881 INFO    MainThread:397944 [wandb_init.py:init():892] starting backend
+2026-03-23 13:56:08,167 INFO    MainThread:397944 [wandb_init.py:init():895] sending inform_init request
+2026-03-23 13:56:08,180 INFO    MainThread:397944 [wandb_init.py:init():903] backend started and connected
+2026-03-23 13:56:08,181 INFO    MainThread:397944 [wandb_init.py:init():973] updated telemetry
+2026-03-23 13:56:08,194 INFO    MainThread:397944 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-23 13:56:11,614 INFO    MainThread:397944 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-23 13:56:11,706 INFO    MainThread:397944 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-23 13:56:11,707 INFO    MainThread:397944 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-23 13:56:11,707 INFO    MainThread:397944 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-23 13:56:11,707 INFO    MainThread:397944 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-23 13:56:11,712 INFO    MainThread:397944 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-23 13:56:11,713 INFO    MainThread:397944 [wandb_run.py:_config_callback():1403] config_cb None None {'output_dir': 'exps', 'exp_name': 'jsflow-experiment-0.75', 'logging_dir': 'logs', 'report_to': 'wandb', 'sampling_steps': 2000, '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': 'bf16', 'epochs': 1400, 'max_train_steps': 1000000, 'checkpointing_steps': 10000, 'gradient_accumulation_steps': 1, 'learning_rate': 5e-05, '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.2, 't_c': 0.75, 'ot_cls': True}

REG/wandb/run-20260323_135841-w9holkos/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

REG/wandb/run-20260323_135841-w9holkos/files/wandb-metadata.json

@@ -0,0 +1,101 @@
+{
+  "os":  "Linux-5.15.0-94-generic-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-03-23T05:58:41.322248Z",
+  "args":  [
+    "--report-to",
+    "wandb",
+    "--allow-tf32",
+    "--mixed-precision",
+    "bf16",
+    "--seed",
+    "0",
+    "--path-type",
+    "linear",
+    "--prediction",
+    "v",
+    "--weighting",
+    "uniform",
+    "--model",
+    "SiT-XL/2",
+    "--enc-type",
+    "dinov2-vit-b",
+    "--encoder-depth",
+    "8",
+    "--proj-coeff",
+    "0.5",
+    "--output-dir",
+    "exps",
+    "--exp-name",
+    "jsflow-experiment-0.75",
+    "--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",
+    "--learning-rate",
+    "0.00005",
+    "--t-c",
+    "0.75",
+    "--cls",
+    "0.05",
+    "--ot-cls"
+  ],
+  "program":  "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/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/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":  "357568360448"
+    }
+  },
+  "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":  "nlbia82zbry6kpqgagoidmc6x8szwd5d"
+}

REG/wandb/run-20260323_135841-w9holkos/logs/debug-internal.log

@@ -0,0 +1,19 @@
+{"time":"2026-03-23T13:58:41.647788404+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-23T13:58:42.578470875+08:00","level":"INFO","msg":"stream: created new stream","id":"w9holkos"}
+{"time":"2026-03-23T13:58:42.578676113+08:00","level":"INFO","msg":"handler: started","stream_id":"w9holkos"}
+{"time":"2026-03-23T13:58:42.579473589+08:00","level":"INFO","msg":"stream: started","id":"w9holkos"}
+{"time":"2026-03-23T13:58:42.57951741+08:00","level":"INFO","msg":"sender: started","stream_id":"w9holkos"}
+{"time":"2026-03-23T13:58:42.579478227+08:00","level":"INFO","msg":"writer: started","stream_id":"w9holkos"}
+{"time":"2026-03-23T14:49:13.568442881+08:00","level":"INFO","msg":"api: retrying HTTP error","status":408,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream","body":"\n<html><head>\n<meta http-equiv=\"content-type\" content=\"text/html;charset=utf-8\">\n<title>408 Request Timeout</title>\n</head>\n<body text=#000000 bgcolor=#ffffff>\n<h1>Error: Request Timeout</h1>\n<h2>Your client has taken too long to issue its request.</h2>\n<h2></h2>\n</body></html>\n"}
+{"time":"2026-03-23T14:52:15.597652411+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2026-03-23T14:52:26.072213509+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": write tcp 172.20.98.27:52324->35.186.228.49:443: write: broken pipe"}
+{"time":"2026-03-23T17:02:52.905542765+08:00","level":"INFO","msg":"api: retrying HTTP error","status":408,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream","body":"\n<html><head>\n<meta http-equiv=\"content-type\" content=\"text/html;charset=utf-8\">\n<title>408 Request Timeout</title>\n</head>\n<body text=#000000 bgcolor=#ffffff>\n<h1>Error: Request Timeout</h1>\n<h2>Your client has taken too long to issue its request.</h2>\n<h2></h2>\n</body></html>\n"}
+{"time":"2026-03-23T17:05:55.176103762+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2026-03-23T17:06:10.164453104+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": unexpected EOF"}
+{"time":"2026-03-23T22:05:06.25355716+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": read tcp 172.20.98.27:44154->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2026-03-23T22:05:20.791067182+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": read tcp 172.20.98.27:40392->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2026-03-24T02:18:38.770696332+08:00","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/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"}
+{"time":"2026-03-24T06:25:41.879737278+08:00","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/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"}
+{"time":"2026-03-24T06:30:14.989373032+08:00","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/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"}
+{"time":"2026-03-24T09:05:02.85908394+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": read tcp 172.20.98.27:46722->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2026-03-25T04:41:04.741907157+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/w9holkos/file_stream\": unexpected EOF"}

REG/wandb/run-20260323_135841-w9holkos/logs/debug.log

@@ -0,0 +1,20 @@
+2026-03-23 13:58:41,343 INFO    MainThread:400275 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-23 13:58:41,343 INFO    MainThread:400275 [wandb_setup.py:_flush():81] Configure stats pid to 400275
+2026-03-23 13:58:41,343 INFO    MainThread:400275 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-23 13:58:41,343 INFO    MainThread:400275 [wandb_init.py:setup_run_log_directory():717] Logging user logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260323_135841-w9holkos/logs/debug.log
+2026-03-23 13:58:41,343 INFO    MainThread:400275 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to /gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/wandb/run-20260323_135841-w9holkos/logs/debug-internal.log
+2026-03-23 13:58:41,343 INFO    MainThread:400275 [wandb_init.py:init():844] calling init triggers
+2026-03-23 13:58:41,344 INFO    MainThread:400275 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-23 13:58:41,344 INFO    MainThread:400275 [wandb_init.py:init():892] starting backend
+2026-03-23 13:58:41,630 INFO    MainThread:400275 [wandb_init.py:init():895] sending inform_init request
+2026-03-23 13:58:41,643 INFO    MainThread:400275 [wandb_init.py:init():903] backend started and connected
+2026-03-23 13:58:41,646 INFO    MainThread:400275 [wandb_init.py:init():973] updated telemetry
+2026-03-23 13:58:41,659 INFO    MainThread:400275 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-23 13:58:43,108 INFO    MainThread:400275 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-23 13:58:43,201 INFO    MainThread:400275 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-23 13:58:43,201 INFO    MainThread:400275 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-23 13:58:43,201 INFO    MainThread:400275 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-23 13:58:43,202 INFO    MainThread:400275 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-23 13:58:43,209 INFO    MainThread:400275 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-23 13:58:43,210 INFO    MainThread:400275 [wandb_run.py:_config_callback():1403] config_cb None None {'output_dir': 'exps', 'exp_name': 'jsflow-experiment-0.75', 'logging_dir': 'logs', 'report_to': 'wandb', 'sampling_steps': 2000, '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': 'bf16', 'epochs': 1400, 'max_train_steps': 1000000, 'checkpointing_steps': 10000, 'gradient_accumulation_steps': 1, 'learning_rate': 5e-05, '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.05, 't_c': 0.75, 'ot_cls': True}