Add files using upload-large-folder tool

REG/preprocessing/README.md

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

REG/preprocessing/dataset_image_encoder.py

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

REG/preprocessing/dataset_prepare_convert.sh

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

REG/preprocessing/dataset_prepare_encode.sh

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

REG/preprocessing/dataset_tools.py

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

REG/preprocessing/dnnlib/__init__.py

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

REG/preprocessing/dnnlib/util.py

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

REG/preprocessing/encoders.py

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

REG/preprocessing/torch_utils/__init__.py

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

REG/preprocessing/torch_utils/distributed.py

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

REG/preprocessing/torch_utils/misc.py

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

REG/preprocessing/torch_utils/persistence.py

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

REG/preprocessing/torch_utils/training_stats.py

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

REG/wandb/debug-internal.log

@@ -0,0 +1,21 @@
+{"time":"2026-04-08T18:26:46.552297532+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-04-08T18:26:47.20146143+08:00","level":"INFO","msg":"stream: created new stream","id":"xtwg5t5s"}
+{"time":"2026-04-08T18:26:47.201551011+08:00","level":"INFO","msg":"handler: started","stream_id":"xtwg5t5s"}
+{"time":"2026-04-08T18:26:47.202423643+08:00","level":"INFO","msg":"stream: started","id":"xtwg5t5s"}
+{"time":"2026-04-08T18:26:47.202450453+08:00","level":"INFO","msg":"writer: started","stream_id":"xtwg5t5s"}
+{"time":"2026-04-08T18:26:47.202479681+08:00","level":"INFO","msg":"sender: started","stream_id":"xtwg5t5s"}
+{"time":"2026-04-09T00:59:33.394616937+08:00","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/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-04-09T15:26:36.673675921+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream\": read tcp 172.20.98.30:37630->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2026-04-09T15:32:51.675782111+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream\": read tcp 172.20.98.30:55710->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2026-04-09T15:33:36.688517829+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream\": EOF"}
+{"time":"2026-04-10T00:33:41.365462236+08:00","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/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-04-10T06:11:35.438909216+08:00","level":"INFO","msg":"api: retrying HTTP error","status":429,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream","body":"{\"error\":\"rate limit exceeded: per_run limit on filestream requests\"}"}
+{"time":"2026-04-11T02:04:06.260667043+08:00","level":"INFO","msg":"api: retrying HTTP error","status":500,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream","body":"{\"error\":\"context deadline exceeded\"}"}
+{"time":"2026-04-11T10:00:44.531212038+08:00","level":"INFO","msg":"api: retrying HTTP error","status":500,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream","body":"{\"error\":\"context deadline exceeded\"}"}
+{"time":"2026-04-11T10:20:26.360393211+08:00","level":"INFO","msg":"api: retrying HTTP error","status":500,"url":"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream","body":"{\"error\":\"context deadline exceeded\"}"}
+{"time":"2026-04-12T21:59:44.458847327+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2026-04-13T00:04:28.494081102+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/2365972933-teleai/REG/xtwg5t5s/file_stream\": read tcp 172.20.98.30:35484->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2026-04-13T02:39:57.535775934+08:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2026-04-13T02:39:58.493368195+08:00","level":"INFO","msg":"handler: closed","stream_id":"xtwg5t5s"}
+{"time":"2026-04-13T02:39:58.494772782+08:00","level":"INFO","msg":"sender: closed","stream_id":"xtwg5t5s"}
+{"time":"2026-04-13T02:39:58.49521181+08:00","level":"INFO","msg":"stream: closed","id":"xtwg5t5s"}

REG/wandb/debug.log

@@ -0,0 +1,22 @@
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [wandb_setup.py:_flush():81] Configure stats pid to 128263
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [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-20260408_182646-xtwg5t5s/logs/debug.log
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [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-20260408_182646-xtwg5t5s/logs/debug-internal.log
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [wandb_init.py:init():844] calling init triggers
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-04-08 18:26:46,224 INFO    MainThread:128263 [wandb_init.py:init():892] starting backend
+2026-04-08 18:26:46,532 INFO    MainThread:128263 [wandb_init.py:init():895] sending inform_init request
+2026-04-08 18:26:46,548 INFO    MainThread:128263 [wandb_init.py:init():903] backend started and connected
+2026-04-08 18:26:46,551 INFO    MainThread:128263 [wandb_init.py:init():973] updated telemetry
+2026-04-08 18:26:46,572 INFO    MainThread:128263 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-04-08 18:26:47,862 INFO    MainThread:128263 [wandb_init.py:init():1042] starting run threads in backend
+2026-04-08 18:26:47,956 INFO    MainThread:128263 [wandb_run.py:_console_start():2524] atexit reg
+2026-04-08 18:26:47,956 INFO    MainThread:128263 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-04-08 18:26:47,956 INFO    MainThread:128263 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-04-08 18:26:47,956 INFO    MainThread:128263 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-04-08 18:26:48,108 INFO    MainThread:128263 [wandb_init.py:init():1082] run started, returning control to user process
+2026-04-08 18:26:48,108 INFO    MainThread:128263 [wandb_run.py:_config_callback():1403] config_cb None None {'output_dir': 'exps', 'exp_name': 'jsflow-experiment-0.75-0.01-one-step', 'logging_dir': 'logs', 'report_to': 'wandb', 'sampling_steps': 2000, 'resume_step': 0, 'resume_from_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', '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': 14000, 'max_train_steps': 10000000, '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.005, 't_c': 0.75, 'ot_cls': True, 'tc_velocity_loss_coeff': 2.0}
+2026-04-13 02:35:32,832 INFO    wandb-AsyncioManager-main:128263 [service_client.py:_forward_responses():134] Reached EOF.
+2026-04-13 02:35:32,833 INFO    wandb-AsyncioManager-main:128263 [mailbox.py:close():155] Closing mailbox, abandoning 1 handles.

REG/wandb/run-20260322_141726-2yw08kz9/files/config.yaml

@@ -0,0 +1,203 @@
+_wandb:
+    value:
+        cli_version: 0.25.0
+        e:
+            257k9ot60u1bv0aiwlacsvutj9c72h7y:
+                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: "357556633600"
+                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-22T06:17:26.670763Z"
+                writerId: 257k9ot60u1bv0aiwlacsvutj9c72h7y
+        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
+                - 61
+            "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: 10000
+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_141726-2yw08kz9/files/output.log

@@ -0,0 +1,27 @@
+Steps:   0%|                                                                                      | 1/1000000 [00:02<614:34:39,  2.21s/it][2026-03-22 14:17:31] Generating EMA samples done.
+[2026-03-22 14:17:31] Step: 1, Training Logs: loss_final: 3.278940, loss_mean: 1.706308, proj_loss: 0.001541, loss_mean_cls: 1.571091, grad_norm: 1.481672
+Steps:   0%| | 2/1000000 [00:02<289:06:04,  1.04s/it, grad_norm=1.48, loss_final=3.28, loss_mean=1.71, loss_mean_cls=1.57, proj_loss=0.001[2026-03-22 14:17:31] Step: 2, Training Logs: loss_final: 3.211831, loss_mean: 1.688932, proj_loss: -0.010287, loss_mean_cls: 1.533185, grad_norm: 1.055476
+Steps:   0%| | 3/1000000 [00:02<187:48:39,  1.48it/s, grad_norm=1.06, loss_final=3.21, loss_mean=1.69, loss_mean_cls=1.53, proj_loss=-0.01[2026-03-22 14:17:31] Step: 3, Training Logs: loss_final: 3.201248, loss_mean: 1.663205, proj_loss: -0.019184, loss_mean_cls: 1.557227, grad_norm: 1.116387
+Steps:   0%| | 4/1000000 [00:02<140:12:43,  1.98it/s, grad_norm=1.12, loss_final=3.2, loss_mean=1.66, loss_mean_cls=1.56, proj_loss=-0.019[2026-03-22 14:17:32] Step: 4, Training Logs: loss_final: 3.198367, loss_mean: 1.682051, proj_loss: -0.026376, loss_mean_cls: 1.542691, grad_norm: 0.722294
+Steps:   0%| | 5/1000000 [00:03<113:52:43,  2.44it/s, grad_norm=0.722, loss_final=3.2, loss_mean=1.68, loss_mean_cls=1.54, proj_loss=-0.02[2026-03-22 14:17:32] Step: 5, Training Logs: loss_final: 3.140483, loss_mean: 1.679105, proj_loss: -0.034564, loss_mean_cls: 1.495943, grad_norm: 0.811589
+Steps:   0%| | 6/1000000 [00:03<97:59:40,  2.83it/s, grad_norm=0.812, loss_final=3.14, loss_mean=1.68, loss_mean_cls=1.5, proj_loss=-0.034[2026-03-22 14:17:32] Step: 6, Training Logs: loss_final: 2.988440, loss_mean: 1.682339, proj_loss: -0.039506, loss_mean_cls: 1.345606, grad_norm: 0.931524
+Steps:   0%| | 7/1000000 [00:03<87:55:00,  3.16it/s, grad_norm=0.932, loss_final=2.99, loss_mean=1.68, loss_mean_cls=1.35, proj_loss=-0.03[2026-03-22 14:17:32] Step: 7, Training Logs: loss_final: 3.111949, loss_mean: 1.690802, proj_loss: -0.042757, loss_mean_cls: 1.463904, grad_norm: 0.830852
+Steps:   0%| | 8/1000000 [00:03<81:19:20,  3.42it/s, grad_norm=0.831, loss_final=3.11, loss_mean=1.69, loss_mean_cls=1.46, proj_loss=-0.04[2026-03-22 14:17:33] Step: 8, Training Logs: loss_final: 3.278931, loss_mean: 1.660797, proj_loss: -0.045011, loss_mean_cls: 1.663145, grad_norm: 0.847438
+Steps:   0%| | 9/1000000 [00:04<76:56:10,  3.61it/s, grad_norm=0.847, loss_final=3.28, loss_mean=1.66, loss_mean_cls=1.66, proj_loss=-0.04[2026-03-22 14:17:33] Step: 9, Training Logs: loss_final: 3.221569, loss_mean: 1.658834, proj_loss: -0.046031, loss_mean_cls: 1.608767, grad_norm: 0.909827
+Steps:   0%| | 10/1000000 [00:04<73:57:18,  3.76it/s, grad_norm=0.91, loss_final=3.22, loss_mean=1.66, loss_mean_cls=1.61, proj_loss=-0.04[2026-03-22 14:17:33] Step: 10, Training Logs: loss_final: 3.216744, loss_mean: 1.665229, proj_loss: -0.047761, loss_mean_cls: 1.599277, grad_norm: 1.014574
+Steps:   0%| | 11/1000000 [00:04<71:52:01,  3.87it/s, grad_norm=1.01, loss_final=3.22, loss_mean=1.67, loss_mean_cls=1.6, proj_loss=-0.047[2026-03-22 14:17:33] Step: 11, Training Logs: loss_final: 3.216658, loss_mean: 1.649915, proj_loss: -0.049347, loss_mean_cls: 1.616090, grad_norm: 1.028789
+Steps:   0%| | 12/1000000 [00:04<70:26:20,  3.94it/s, grad_norm=1.03, loss_final=3.22, loss_mean=1.65, loss_mean_cls=1.62, proj_loss=-0.04[2026-03-22 14:17:34] Step: 12, Training Logs: loss_final: 3.155676, loss_mean: 1.624463, proj_loss: -0.049856, loss_mean_cls: 1.581069, grad_norm: 1.231291
+Steps:   0%| | 13/1000000 [00:05<69:25:29,  4.00it/s, grad_norm=1.23, loss_final=3.16, loss_mean=1.62, loss_mean_cls=1.58, proj_loss=-0.04Traceback (most recent call last):
+  File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 527, in <module>
+    main(args)
+  File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 415, in main
+    "loss_final": accelerator.gather(loss).mean().detach().item(),
+                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+KeyboardInterrupt
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 527, in <module>
+[rank0]:     main(args)
+[rank0]:   File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 415, in main
+[rank0]:     "loss_final": accelerator.gather(loss).mean().detach().item(),
+[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: KeyboardInterrupt

REG/wandb/run-20260322_141726-2yw08kz9/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_141726-2yw08kz9/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-22T06:17:26.670763Z",
+  "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":  "357556633600"
+    }
+  },
+  "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":  "257k9ot60u1bv0aiwlacsvutj9c72h7y"
+}

REG/wandb/run-20260322_141726-2yw08kz9/files/wandb-summary.json

@@ -0,0 +1 @@
+{"loss_mean_cls":1.5810688734054565,"_timestamp":1.7741602540511734e+09,"_runtime":5.247627056,"loss_mean":1.6244629621505737,"proj_loss":-0.04985573887825012,"grad_norm":1.2312908172607422,"_wandb":{"runtime":5},"_step":12,"loss_final":3.1556761264801025}

REG/wandb/run-20260322_141726-2yw08kz9/logs/debug-internal.log

@@ -0,0 +1,7 @@
+{"time":"2026-03-22T14:17:27.013311984+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-22T14:17:28.347732261+08:00","level":"INFO","msg":"stream: created new stream","id":"2yw08kz9"}
+{"time":"2026-03-22T14:17:28.347960938+08:00","level":"INFO","msg":"handler: started","stream_id":"2yw08kz9"}
+{"time":"2026-03-22T14:17:28.348671928+08:00","level":"INFO","msg":"stream: started","id":"2yw08kz9"}
+{"time":"2026-03-22T14:17:28.348731034+08:00","level":"INFO","msg":"sender: started","stream_id":"2yw08kz9"}
+{"time":"2026-03-22T14:17:28.348748525+08:00","level":"INFO","msg":"writer: started","stream_id":"2yw08kz9"}
+{"time":"2026-03-22T14:17:34.316421629+08:00","level":"INFO","msg":"stream: closing","id":"2yw08kz9"}

REG/wandb/run-20260322_141726-2yw08kz9/logs/debug.log

@@ -0,0 +1,22 @@
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [wandb_setup.py:_flush():81] Configure stats pid to 316313
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [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_141726-2yw08kz9/logs/debug.log
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [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_141726-2yw08kz9/logs/debug-internal.log
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [wandb_init.py:init():844] calling init triggers
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-22 14:17:26,691 INFO    MainThread:316313 [wandb_init.py:init():892] starting backend
+2026-03-22 14:17:26,994 INFO    MainThread:316313 [wandb_init.py:init():895] sending inform_init request
+2026-03-22 14:17:27,008 INFO    MainThread:316313 [wandb_init.py:init():903] backend started and connected
+2026-03-22 14:17:27,011 INFO    MainThread:316313 [wandb_init.py:init():973] updated telemetry
+2026-03-22 14:17:27,025 INFO    MainThread:316313 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-22 14:17:29,067 INFO    MainThread:316313 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-22 14:17:29,158 INFO    MainThread:316313 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-22 14:17:29,158 INFO    MainThread:316313 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-22 14:17:29,158 INFO    MainThread:316313 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-22 14:17:29,159 INFO    MainThread:316313 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-22 14:17:29,163 INFO    MainThread:316313 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-22 14:17:29,163 INFO    MainThread:316313 [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': 10000, 'resume_step': 0, 'model': 'SiT-XL/2', 'num_classes': 1000, 'encoder_depth': 8, 'fused_attn': True, 'qk_norm': False, 'ops_head': 16, 'data_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256', 'semantic_features_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0', 'resolution': 256, 'batch_size': 256, 'allow_tf32': True, 'mixed_precision': '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 14:17:34,316 INFO    wandb-AsyncioManager-main:316313 [service_client.py:_forward_responses():134] Reached EOF.
+2026-03-22 14:17:34,316 INFO    wandb-AsyncioManager-main:316313 [mailbox.py:close():155] Closing mailbox, abandoning 1 handles.

REG/wandb/run-20260322_141833-vm0y8t9t/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_141833-vm0y8t9t/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-22T06:18:33.208941Z",
+  "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":  "357556703232"
+    }
+  },
+  "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":  "gklxguwapb72cxij4696gj37bh1rbthi"
+}

REG/wandb/run-20260322_141833-vm0y8t9t/logs/debug-internal.log

@@ -0,0 +1,6 @@
+{"time":"2026-03-22T14:18:33.472940651+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-22T14:18:35.380852704+08:00","level":"INFO","msg":"stream: created new stream","id":"vm0y8t9t"}
+{"time":"2026-03-22T14:18:35.381056887+08:00","level":"INFO","msg":"handler: started","stream_id":"vm0y8t9t"}
+{"time":"2026-03-22T14:18:35.382108345+08:00","level":"INFO","msg":"writer: started","stream_id":"vm0y8t9t"}
+{"time":"2026-03-22T14:18:35.382119604+08:00","level":"INFO","msg":"stream: started","id":"vm0y8t9t"}
+{"time":"2026-03-22T14:18:35.382161533+08:00","level":"INFO","msg":"sender: started","stream_id":"vm0y8t9t"}

REG/wandb/run-20260322_141833-vm0y8t9t/logs/debug.log

@@ -0,0 +1,20 @@
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [wandb_setup.py:_flush():81] Configure stats pid to 318585
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [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_141833-vm0y8t9t/logs/debug.log
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [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_141833-vm0y8t9t/logs/debug-internal.log
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [wandb_init.py:init():844] calling init triggers
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-22 14:18:33,237 INFO    MainThread:318585 [wandb_init.py:init():892] starting backend
+2026-03-22 14:18:33,460 INFO    MainThread:318585 [wandb_init.py:init():895] sending inform_init request
+2026-03-22 14:18:33,470 INFO    MainThread:318585 [wandb_init.py:init():903] backend started and connected
+2026-03-22 14:18:33,472 INFO    MainThread:318585 [wandb_init.py:init():973] updated telemetry
+2026-03-22 14:18:33,485 INFO    MainThread:318585 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-22 14:18:36,829 INFO    MainThread:318585 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-22 14:18:36,920 INFO    MainThread:318585 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-22 14:18:36,920 INFO    MainThread:318585 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-22 14:18:36,921 INFO    MainThread:318585 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-22 14:18:36,921 INFO    MainThread:318585 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-22 14:18:36,924 INFO    MainThread:318585 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-22 14:18:36,924 INFO    MainThread:318585 [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': 10000, 'resume_step': 0, 'model': 'SiT-XL/2', 'num_classes': 1000, 'encoder_depth': 8, 'fused_attn': True, 'qk_norm': False, 'ops_head': 16, 'data_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256', 'semantic_features_dir': '/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0', 'resolution': 256, 'batch_size': 256, 'allow_tf32': True, 'mixed_precision': '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-20260322_150022-yhxc5cgu/files/output.log

@@ -0,0 +1,19 @@
+Steps:   0%|          | 1/1000000 [00:02<652:30:07,  2.35s/it][2026-03-22 15:00:28] Generating EMA samples for evaluation (t=1→0 and t=0.5)...
+Traceback (most recent call last):
+  File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 628, in <module>
+    main(args)
+  File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 425, in main
+    cls_init = torch.randn(n_samples, base_model.semantic_channels, device=device)
+                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
+    raise AttributeError(
+AttributeError: 'SiT' object has no attribute 'semantic_channels'
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 628, in <module>
+[rank0]:     main(args)
+[rank0]:   File "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/REG/train.py", line 425, in main
+[rank0]:     cls_init = torch.randn(n_samples, base_model.semantic_channels, device=device)
+[rank0]:                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/gemini/space/zhaozy/guzhenyu/envs/envs/SiT/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
+[rank0]:     raise AttributeError(
+[rank0]: AttributeError: 'SiT' object has no attribute 'semantic_channels'

REG/wandb/run-20260322_150022-yhxc5cgu/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_150022-yhxc5cgu/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:00:22.092510Z",
+  "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":  "357557354496"
+    }
+  },
+  "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":  "ucanic8s891x6sl28vnbha78lzoecw66"
+}

REG/wandb/run-20260322_150022-yhxc5cgu/logs/debug-internal.log

@@ -0,0 +1,7 @@
+{"time":"2026-03-22T15:00:22.432399726+08:00","level":"INFO","msg":"stream: starting","core version":"0.25.0"}
+{"time":"2026-03-22T15:00:25.799578446+08:00","level":"INFO","msg":"stream: created new stream","id":"yhxc5cgu"}
+{"time":"2026-03-22T15:00:25.799734466+08:00","level":"INFO","msg":"handler: started","stream_id":"yhxc5cgu"}
+{"time":"2026-03-22T15:00:25.80075778+08:00","level":"INFO","msg":"stream: started","id":"yhxc5cgu"}
+{"time":"2026-03-22T15:00:25.800786229+08:00","level":"INFO","msg":"writer: started","stream_id":"yhxc5cgu"}
+{"time":"2026-03-22T15:00:25.800837858+08:00","level":"INFO","msg":"sender: started","stream_id":"yhxc5cgu"}
+{"time":"2026-03-22T15:00:28.913273863+08:00","level":"INFO","msg":"stream: closing","id":"yhxc5cgu"}

REG/wandb/run-20260322_150022-yhxc5cgu/logs/debug.log

@@ -0,0 +1,22 @@
+2026-03-22 15:00:22,124 INFO    MainThread:323629 [wandb_setup.py:_flush():81] Current SDK version is 0.25.0
+2026-03-22 15:00:22,124 INFO    MainThread:323629 [wandb_setup.py:_flush():81] Configure stats pid to 323629
+2026-03-22 15:00:22,124 INFO    MainThread:323629 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-03-22 15:00:22,124 INFO    MainThread:323629 [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_150022-yhxc5cgu/logs/debug.log
+2026-03-22 15:00:22,124 INFO    MainThread:323629 [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_150022-yhxc5cgu/logs/debug-internal.log
+2026-03-22 15:00:22,125 INFO    MainThread:323629 [wandb_init.py:init():844] calling init triggers
+2026-03-22 15:00:22,125 INFO    MainThread:323629 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-03-22 15:00:22,125 INFO    MainThread:323629 [wandb_init.py:init():892] starting backend
+2026-03-22 15:00:22,416 INFO    MainThread:323629 [wandb_init.py:init():895] sending inform_init request
+2026-03-22 15:00:22,429 INFO    MainThread:323629 [wandb_init.py:init():903] backend started and connected
+2026-03-22 15:00:22,431 INFO    MainThread:323629 [wandb_init.py:init():973] updated telemetry
+2026-03-22 15:00:22,447 INFO    MainThread:323629 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-03-22 15:00:26,403 INFO    MainThread:323629 [wandb_init.py:init():1042] starting run threads in backend
+2026-03-22 15:00:26,494 INFO    MainThread:323629 [wandb_run.py:_console_start():2524] atexit reg
+2026-03-22 15:00:26,494 INFO    MainThread:323629 [wandb_run.py:_redirect():2373] redirect: wrap_raw
+2026-03-22 15:00:26,494 INFO    MainThread:323629 [wandb_run.py:_redirect():2442] Wrapping output streams.
+2026-03-22 15:00:26,495 INFO    MainThread:323629 [wandb_run.py:_redirect():2465] Redirects installed.
+2026-03-22 15:00:26,500 INFO    MainThread:323629 [wandb_init.py:init():1082] run started, returning control to user process
+2026-03-22 15:00:26,500 INFO    MainThread:323629 [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:00:28,913 INFO    wandb-AsyncioManager-main:323629 [service_client.py:_forward_responses():134] Reached EOF.
+2026-03-22 15:00:28,913 INFO    wandb-AsyncioManager-main:323629 [mailbox.py:close():155] Closing mailbox, abandoning 1 handles.