update

README.md

@@ -1,27 +1,7 @@
----
-title: CoTyle
-emoji: 🎨
-colorFrom: gray
-colorTo: purple
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-python_version: 3.10
-# 移除 license 字段，因为它不是官方支持的字段
-gpu: true
-suggested_hardware: a100-large
-models:
-  - Kwai-Kolors/CoTyle
-tags:
-  - image-generation
-  - code-to-style
-  - gradio
----
-
 
 # A Style is Worth One Code: Unlocking Code-to-Style Image Generation with Discrete Style Space
 <p align="center"> 
-    <a href="xxx"><img alt="Build" src="https://img.shields.io/badge/arXiv-Paper-da282a.svg"></a>
+    <a href="https://arxiv.org/abs/2511.10555"><img alt="Build" src="https://img.shields.io/badge/arXiv-Paper-da282a.svg"></a>
     <a href="https://Kwai-Kolors.github.io/CoTyle/"><img alt="Build" src="https://img.shields.io/badge/Project%20Page-Homepage-yellow"></a> 
     <a href="https://github.com/Kwai-Kolors/CoTyle"><img alt="Build" src="https://img.shields.io/badge/GitHub-Code-f8f0f0.svg"></a> 
     <a href="https://huggingface.co/spaces/Kwai-Kolors/CoTyle"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Demo-32CD32"></a>
@@ -78,6 +58,8 @@ git clone https://github.com/Kwai-Kolors/CoTyle
 cd CoTyle
 conda create -n cotyle python=3.10 
 conda activate cotyle
+pip install torch==2.6.0 torchvision==0.21.0
+pip install -e git+https://github.com/Lakonik/piFlow.git@b1ef16e5e305251bccdfeac2a0e3d0ef339b974a#egg=lakonlab
 pip install -r requirements.txt
 ```
 

lakonlab/__init__.py

@@ -0,0 +1,20 @@
+import warnings
+
+# suppress warnings from MMCV about optional dependencies
+warnings.filterwarnings(
+    'ignore',
+    category=UserWarning,
+    message=r'^Fail to import ``MultiScaleDeformableAttention`` from ``mmcv\.ops\.multi_scale_deform_attn``.*',
+    module=r'^mmcv\.cnn\.bricks\.transformer$',
+)
+
+# import all modules for registration
+from .apis import *
+from .datasets import *
+from .models import *
+from .ops import *
+from .runner import *
+from .evaluation import *
+from .utils import *
+
+from .version import __version__

lakonlab/apis/__init__.py

@@ -0,0 +1,4 @@
+from .train import train_model
+from .inference import init_model
+
+__all__ = ['init_model', 'train_model']

lakonlab/apis/inference.py

@@ -0,0 +1,56 @@
+import torch
+import mmcv
+from mmcv.runner import load_checkpoint
+from mmgen.models import build_model
+from lakonlab.runner.hooks.ema_hook import get_ori_key
+
+
+def init_model(
+        config, checkpoint=None, device='cuda:0', cfg_options=None,
+        ema_only=True, use_fp16=False, use_bf16=False):
+    if isinstance(config, str):
+        config = mmcv.Config.fromfile(config)
+    elif not isinstance(config, mmcv.Config):
+        raise TypeError('config must be a filename or Config object, '
+                        f'but got {type(config)}')
+    if cfg_options is not None:
+        config.merge_from_dict(cfg_options)
+
+    model = build_model(
+        config.model, train_cfg=config.train_cfg, test_cfg=config.test_cfg)
+
+    if ema_only:
+        module_keys = []
+        for hook in config.get('custom_hooks', []):
+            if hook['type'] in ('ExponentialMovingAverageHookMod', 'ExponentialMovingAverageHook'):
+                if isinstance(hook['module_keys'], str):
+                    module_keys.append(hook['module_keys'])
+                else:
+                    module_keys.extend(hook['module_keys'])
+        for key in module_keys:
+            ori_key = get_ori_key(key)
+            del model._modules[ori_key]
+
+    if checkpoint is not None:
+        load_checkpoint(model, checkpoint, map_location='cpu')
+
+    model._cfg = config  # save the config in the model for convenience
+
+    for module in model.modules():
+        if hasattr(module, 'bake_lora_weights'):
+            module.bake_lora_weights()
+
+    if use_fp16 or use_bf16:
+        for m in model.modules():
+            if hasattr(m, 'autocast_dtype'):
+                setattr(m, 'autocast_dtype', None)
+        if use_fp16:
+            assert not use_bf16
+            model.to(dtype=torch.float16)
+        elif use_bf16:
+            model.to(dtype=torch.bfloat16)
+
+    model.to(device)
+    model.eval()
+
+    return model

lakonlab/apis/train.py

@@ -0,0 +1,166 @@
+# Modified from https://github.com/open-mmlab/mmgeneration
+
+import warnings
+import re
+from copy import deepcopy
+
+from mmcv.parallel import MMDataParallel
+from mmcv.runner import HOOKS, IterBasedRunner, OptimizerHook, build_runner
+from mmcv.utils import build_from_cfg
+
+from mmgen.datasets import build_dataset
+from mmgen.utils import get_root_logger
+
+from lakonlab.parallel import apply_module_wrapper
+from lakonlab.runner.optimizer import build_optimizers
+from lakonlab.runner.checkpoint import exists_ckpt
+from lakonlab.datasets import build_dataloader
+
+
+def train_model(model,
+                dataset,
+                cfg,
+                distributed=False,
+                validate=False,
+                timestamp=None,
+                meta=None):
+    logger = get_root_logger(cfg.log_level)
+
+    # prepare data loaders
+    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
+
+    # default loader config
+    loader_cfg = dict(
+        # cfg.gpus will be ignored if distributed
+        num_gpus=len(cfg.gpu_ids),
+        seed=cfg.seed)
+
+    # The overall dataloader settings
+    loader_cfg.update({
+        k: v
+        for k, v in cfg.data.items()
+        if k not in [
+            'train', 'train_dataloader', 'val_dataloader', 'test_dataloader'
+        ] and not re.fullmatch(r'(val|test)\d*', k)
+    })
+
+    # The specific datalaoder settings
+    train_loader_cfg = {**loader_cfg, **cfg.data.get('train_dataloader', {})}
+
+    data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]
+
+    if cfg.get('apex_amp', None):
+        raise NotImplementedError('Apex AMP is no longer supported.')
+
+    # put model on gpus
+    if distributed:
+        module_wrapper = cfg.get('module_wrapper', None)
+        model = apply_module_wrapper(model, module_wrapper, cfg)
+    else:
+        model = MMDataParallel(model, device_ids=cfg.gpu_ids)
+
+    # build optimizer
+    if cfg.optimizer:
+        optimizer = build_optimizers(model, cfg.optimizer)
+    # In GANs, we allow building optimizer in GAN model.
+    else:
+        optimizer = None
+
+    # allow users to define the runner
+    if cfg.get('runner', None):
+        runner = build_runner(
+            cfg.runner,
+            dict(
+                model=model,
+                optimizer=optimizer,
+                work_dir=cfg.work_dir,
+                logger=logger,
+                use_apex_amp=False,
+                meta=meta))
+    else:
+        runner = IterBasedRunner(
+            model,
+            optimizer=optimizer,
+            work_dir=cfg.work_dir,
+            logger=logger,
+            meta=meta)
+        # set if use dynamic ddp in training
+        # is_dynamic_ddp=cfg.get('is_dynamic_ddp', False))
+    # an ugly walkaround to make the .log and .log.json filenames the same
+    runner.timestamp = timestamp
+
+    # fp16 setting
+    fp16_cfg = cfg.get('fp16', None)
+
+    # In GANs, we can directly optimize parameter in `train_step` function.
+    if cfg.get('optimizer_cfg', None) is None:
+        optimizer_config = None
+    elif fp16_cfg is not None:
+        raise NotImplementedError('Fp16 has not been supported.')
+        # optimizer_config = Fp16OptimizerHook(
+        #     **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
+    # default to use OptimizerHook
+    elif distributed and 'type' not in cfg.optimizer_config:
+        optimizer_config = OptimizerHook(**cfg.optimizer_config)
+    else:
+        optimizer_config = cfg.optimizer_config
+
+    # # update `out_dir` in  ckpt hook
+    # if cfg.checkpoint_config is not None:
+    #     cfg.checkpoint_config['out_dir'] = os.path.join(
+    #         cfg.work_dir, cfg.checkpoint_config.get('out_dir', 'ckpt'))
+
+    # register hooks
+    runner.register_training_hooks(cfg.lr_config, optimizer_config,
+                                   cfg.checkpoint_config, cfg.log_config,
+                                   cfg.get('momentum_config', None))
+
+    # # DistSamplerSeedHook should be used with EpochBasedRunner
+    # if distributed:
+    #     runner.register_hook(DistSamplerSeedHook())
+
+    # In general, we do NOT adopt standard evaluation hook in GAN training.
+    # Thus, if you want a eval hook, you need further define the key of
+    # 'evaluation' in the config.
+    # register eval hooks
+    if validate and cfg.get('evaluation', None) is not None:
+        assert isinstance(cfg.evaluation, list)
+        for eval_cfg_ in cfg.evaluation:
+            val_dataset = build_dataset(cfg.data[eval_cfg_.data])
+            val_loader_cfg = {
+                **loader_cfg, 'shuffle': False,
+                **cfg.data.get('val_dataloader', {})
+            }
+            val_dataloader = build_dataloader(val_dataset, **val_loader_cfg)
+            eval_cfg = deepcopy(eval_cfg_)
+            priority = eval_cfg.pop('priority', 'LOW')
+            eval_cfg.update(dict(dist=distributed, dataloader=val_dataloader))
+            eval_hook = build_from_cfg(eval_cfg, HOOKS)
+            runner.register_hook(eval_hook, priority=priority)
+
+    # user-defined hooks
+    if cfg.get('custom_hooks', None):
+        custom_hooks = cfg.custom_hooks
+        assert isinstance(custom_hooks, list), \
+            f'custom_hooks expect list type, but got {type(custom_hooks)}'
+        for hook_cfg in cfg.custom_hooks:
+            assert isinstance(hook_cfg, dict), \
+                'Each item in custom_hooks expects dict type, but got ' \
+                f'{type(hook_cfg)}'
+            hook_cfg = hook_cfg.copy()
+            priority = hook_cfg.pop('priority', 'NORMAL')
+            hook = build_from_cfg(hook_cfg, HOOKS)
+            runner.register_hook(hook, priority=priority)
+
+    ckpt_kwargs = dict()
+    if distributed and module_wrapper.lower() in ['fsdp', 'fsdp2']:
+        ckpt_kwargs.update(map_location='cpu')
+    if exists_ckpt(cfg.resume_from):
+        runner.resume(cfg.resume_from, **ckpt_kwargs)
+        for data_loader in data_loaders:
+            data_loader.sampler.set_epoch(runner.epoch)
+            data_loader.sampler.set_iter(runner.iter)
+    elif exists_ckpt(cfg.load_from):
+        runner.load_checkpoint(cfg.load_from, **ckpt_kwargs)
+
+    runner.run(data_loaders, cfg.workflow, cfg.total_iters)

lakonlab/datasets/__init__.py

@@ -0,0 +1,8 @@
+from .builder import build_dataloader
+from .imagenet import ImageNet
+from .checkerboard import CheckerboardData
+from .image_prompts import ImagePrompt
+
+__all__ = [
+    'build_dataloader', 'ImageNet', 'CheckerboardData', 'ImagePrompt'
+]

lakonlab/datasets/builder.py

@@ -0,0 +1,61 @@
+import warnings
+from functools import partial
+
+from mmcv.parallel import collate
+from mmcv.runner import get_dist_info
+from mmcv.utils import TORCH_VERSION, digit_version
+from torch.utils.data import DataLoader
+
+from mmgen.datasets.builder import worker_init_fn
+from .samplers import DistributedSampler
+
+
+def build_dataloader(dataset,
+                     samples_per_gpu,
+                     workers_per_gpu,
+                     num_gpus=1,
+                     dist=True,
+                     shuffle=True,
+                     seed=None,
+                     persistent_workers=False,
+                     sampler=None,
+                     **kwargs):
+    rank, world_size = get_dist_info()
+    if dist:
+        assert sampler is None, 'sampler is not supported in distributed mode'
+        sampler = DistributedSampler(
+            dataset,
+            world_size,
+            rank,
+            shuffle=shuffle,
+            samples_per_gpu=samples_per_gpu,
+            seed=seed)
+        shuffle = False
+        batch_size = samples_per_gpu
+        num_workers = workers_per_gpu
+    else:
+        batch_size = num_gpus * samples_per_gpu
+        num_workers = num_gpus * workers_per_gpu
+
+    init_fn = partial(
+        worker_init_fn, num_workers=num_workers, rank=rank,
+        seed=seed) if seed is not None else None
+
+    if (digit_version(TORCH_VERSION) >= digit_version('1.7.0')
+            and TORCH_VERSION != 'parrots'):
+        kwargs['persistent_workers'] = persistent_workers
+    elif persistent_workers is True:
+        warnings.warn('persistent_workers is invalid because your pytorch '
+                      'version is lower than 1.7.0')
+
+    data_loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        num_workers=num_workers,
+        collate_fn=partial(collate, samples_per_gpu=samples_per_gpu),
+        shuffle=shuffle,
+        worker_init_fn=init_fn,
+        **kwargs)
+
+    return data_loader

lakonlab/datasets/checkerboard.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import torch
+
+from torch.utils.data import Dataset
+from mmgen.datasets.builder import DATASETS
+
+
+@DATASETS.register_module()
+class CheckerboardData(Dataset):
+    def __init__(
+            self,
+            n_rc=4,
+            n_samples=1e8,
+            thickness=1.0,
+            scale=1,
+            shift=[0.0, 0.0],
+            rotation=0.0,
+            test_mode=False):
+        super().__init__()
+        self.n_rc = n_rc
+        self.n_samples = int(n_samples)
+        self.thickness = thickness
+        self.scale = scale
+        self.shift = torch.tensor(shift, dtype=torch.float32)
+        self.rotation = rotation
+        white_squares = [(i, j) for i in range(n_rc) for j in range(n_rc) if (i + j) % 2 == 0]
+        self.white_squares = torch.tensor(white_squares, dtype=torch.float32)
+        self.n_squares = len(white_squares)
+        self.samples = self.draw_samples(self.n_samples)
+
+    def draw_samples(self, n_samples):
+        chosen_indices = torch.randint(0, self.n_squares, size=(n_samples, ))
+        chosen_squares = self.white_squares[chosen_indices]
+        square_samples = torch.rand(n_samples, 2, dtype=torch.float32)
+        if self.thickness < 1:
+            square_samples = square_samples - 0.5
+            square_samples_r = square_samples.square().sum(dim=-1, keepdims=True)
+            square_samples_angle = torch.atan2(square_samples[:, 1], square_samples[:, 0]).unsqueeze(-1)
+            max_r = torch.minimum(
+                0.5 / square_samples_angle.cos().abs().clamp(min=1e-6),
+                0.5 / square_samples_angle.sin().abs().clamp(min=1e-6)).square()
+            square_samples_r_scaled = max_r - (max_r - square_samples_r) * self.thickness ** 0.5
+            square_samples *= (square_samples_r_scaled / square_samples_r).sqrt()
+            square_samples = square_samples + 0.5
+        samples = (chosen_squares + square_samples) * (2 / self.n_rc) - 1
+        if self.rotation != 0.0:
+            angle = torch.tensor(self.rotation, dtype=torch.float32) * torch.pi / 180
+            rotation_matrix = torch.tensor([[torch.cos(angle), -torch.sin(angle)],
+                                            [torch.sin(angle), torch.cos(angle)]])
+            samples = samples @ rotation_matrix
+        return samples * self.scale + self.shift
+
+    def __len__(self):
+        return self.n_samples
+
+    def __getitem__(self, idx):
+        data = dict(x=self.samples[idx])
+        return data

lakonlab/datasets/image_prompts.py

@@ -0,0 +1,432 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import logging
+import os
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import zstandard as zstd
+import pickle
+import gzip
+import orjson
+import mmcv
+import torch.storage
+torch.storage.UntypedStorage.dtype = torch.uint8  # hot patch for torch 2.6 deserialization
+
+from io import BytesIO
+from typing import Optional, Tuple, Union
+from torch.utils.data import Dataset
+from datasets import load_dataset, DatasetDict, Dataset as HFDataset
+from mmcv.fileio import FileClient
+from mmcv.parallel import DataContainer as DC
+from mmgen.utils import get_root_logger
+from mmgen.datasets.builder import DATASETS
+from lakonlab.utils.io_utils import load_image
+
+
+@DATASETS.register_module()
+class ImagePrompt(Dataset):
+    """Initialize an image/prompt dataset that reads either cached pickled records
+    (zstd-compressed) or a HuggingFace prompt dataset (optionally paired with images).
+
+    Args:
+        data_root (str): Root path for IO, resolved via `mmcv.FileClient`.
+        cache_dir (Optional[str]): Subdirectory of `data_root` containing `.zst`
+            cache shards. Enables cache mode when provided and exists. Caches must
+            contain pickled dicts with keys `"prompt"` and `"prompt_embed_kwargs"`,
+            and optionally `"latents"` or `"latent_size"`.
+        cache_datalist_path (Optional[str]): Optional datalist path for `cache_dir`.
+            Supports `.jsonl`, `.jsonl.gz`, or `.json`. If not exists, files are
+            discovered by listing the directory.
+        ignore_cached_latents (bool): If True, ignores any cached latents and
+            prioritizes loading images from `image_dir`. Defaults to False.
+        prompt_dataset_kwargs (Optional[dict]): Keyword arguments forwarded to
+            `datasets.load_dataset(...)`. Enables prompt-dataset mode when provided.
+            If a `DatasetDict` is returned, a split (e.g., "train") is selected
+            internally.
+        image_dir (Optional[str]): Subdirectory of `data_root` with images to pair
+            with prompts (used only in prompt-dataset mode).
+        image_datalist_path (Optional[str]): Optional datalist for `image_dir`
+            (same formats as above). When `bucketize=True`, JSONL entries must
+            include `"size_idx"`.
+        image_extension (Optional[str]): Image file extension used to compose
+            paths when `image_dir` is set. Defaults to ".png".
+        image_scale_factor (float): Scale factor applied to image spatial dimensions
+            after loading. Defaults to 1.0 (no scaling).
+        negative_prompt_embeds_path (Optional[str]): Path to a `torch.load`-able
+            file containing keyward arguments forwarded to the diffusion model
+            for negative prompt embeddings. Added to each sample when provided.
+        negative_prompt_kwargs (Optional[dict]): Keyword arguments forwarded to
+            the text encoder for negative prompts. Added to each sample when provided.
+        pad_seq_len (int): If set, pads/truncates `"encoder_hidden_states"` and
+            `"encoder_hidden_states_mask"` along the sequence dimension to this
+            length. Defaults to None.
+        latent_size (Optional[Tuple[int]]): Default latent shape `(C, H, W)` used
+            when no cached latents exist and no image size is provided. Defaults to
+            `(16, 128, 128)`.
+        vae_scale_factor (Optional[Union[int, Tuple[int]]]): Downscale factor(s)
+            applied to image dimensions when deriving latent sizes from image size.
+            If an `int`, applies to each spatial dim; if a `tuple`, its length must
+            match the provided image spatial size (e.g., `(H, W)` or `(T, H, W)` for
+            video VAEs).
+        repeat (int): Virtual repetition factor for each underlying sample.
+            Affects `__len__` and index mapping.
+        start_ind (Optional[int]): Start index (inclusive) into the underlying
+            dataset. Defaults to 0.
+        end_ind (int): End index (exclusive) into the underlying dataset. Defaults
+            to dataset length.
+        bucketize (bool): If True, enables bucketing in `DistributedSampler` so that
+            each rank receives samples of the same size. Expects `"size_idx"` in JSONL
+            datalists and collects bucket ids. Defaults to False.
+        test_mode (bool): If True, return deterministic noise per sample instead of
+            reading/allocating real latents or images.
+    """
+
+    PROMPT_KEY_MAPS = {
+        'prompt_embeds': 'encoder_hidden_states',
+        'prompt_embeds_scale': 'encoder_hidden_states_scale',
+        'pooled_prompt_embeds': 'pooled_projections',
+        'prompt_embeds_mask': 'encoder_hidden_states_mask'
+    }
+
+    def __init__(self,
+                 data_root: str,
+                 cache_dir: Optional[str] = None,
+                 cache_datalist_path: Optional[str] = None,
+                 ignore_cached_latents: bool = False,
+                 prompt_dataset_kwargs: Optional[dict] = None,
+                 image_dir: Optional[str] = None,
+                 image_datalist_path: Optional[str] = None,
+                 image_extension: Optional[str] = '.png',
+                 image_scale_factor: float = 1.0,
+                 negative_prompt_embeds_path: Optional[str] = None,
+                 negative_prompt_kwargs: Optional[dict] = None,
+                 pad_seq_len: int = None,
+                 latent_size: Optional[Tuple[int]] = (16, 128, 128),
+                 vae_scale_factor: Optional[Union[int, Tuple[int]]] = 8,
+                 repeat: int = 1,
+                 start_ind: Optional[int] = None,
+                 end_ind: int = None,
+                 bucketize: bool = False,
+                 test_mode: bool = False):
+        super().__init__()
+        self.data_root = data_root
+        self.file_client = FileClient.infer_client(uri=self.data_root)
+
+        self.pad_seq_len = pad_seq_len
+
+        self.cache_dir_path = self.cache_datalist_path = None
+        self.prompt_dataset = self.image_dir_path = self.image_datalist_path = None
+        self.image_extension = image_extension
+        self.image_scale_factor = image_scale_factor
+        self.ignore_cached_latents = ignore_cached_latents
+        self.bucketize = bucketize
+        bucket_ids = None
+
+        if (cache_dir is not None
+                and self.file_client.isdir(self.file_client.join_path(data_root, cache_dir))
+                and cache_datalist_path is not None
+                and FileClient.infer_client(uri=cache_datalist_path).isfile(cache_datalist_path)):
+            self.cache_dir_path = self.file_client.join_path(data_root, cache_dir)
+            self.cache_datalist, bucket_ids = self.parse_datalist(
+                self.cache_dir_path, cache_datalist_path)
+            dataset_len = len(self.cache_datalist)
+
+        elif prompt_dataset_kwargs is not None:
+            self.prompt_dataset = load_dataset(**prompt_dataset_kwargs)
+            if isinstance(self.prompt_dataset, DatasetDict):
+                split = 'train' if 'train' in self.prompt_dataset else list(self.prompt_dataset.keys())[0]
+                self.prompt_dataset = self.prompt_dataset[split]
+            assert isinstance(self.prompt_dataset, HFDataset), \
+                f"Expected HF Dataset/DatasetDict, got {type(self.prompt_dataset)}."
+            dataset_len = len(self.prompt_dataset)
+
+        else:
+            raise ValueError('Either `cache_dir` or `prompt_dataset_kwargs` must be provided.')
+
+        if image_dir is not None and self.file_client.isdir(self.file_client.join_path(data_root, image_dir)):
+            self.image_dir_path = self.file_client.join_path(data_root, image_dir)
+            self.image_datalist, bucket_ids = self.parse_datalist(
+                self.image_dir_path, image_datalist_path, datalist_must_exist=True)
+            assert dataset_len == len(self.image_datalist)
+
+        if bucket_ids is None and self.bucketize:
+            assert self.prompt_dataset is not None
+            bucket_ids = self.get_bucket_ids_from_prompt_dataset()
+
+        self.negative_prompt_embed_kwargs = None
+        if negative_prompt_embeds_path is not None:
+            negative_prompt_embeds_bytesio = BytesIO(
+                FileClient.infer_client(uri=negative_prompt_embeds_path).get(negative_prompt_embeds_path))
+            self.negative_prompt_embed_kwargs = self.parse_prompt_embeds(
+                torch.load(negative_prompt_embeds_bytesio, map_location='cpu'))
+        self.negative_prompt_kwargs = negative_prompt_kwargs
+
+        self.latent_size = latent_size
+        self.vae_scale_factor = vae_scale_factor
+
+        self.repeat = repeat
+        if start_ind is not None:
+            start_ind = max(min(start_ind, dataset_len - 1), -dataset_len) % dataset_len
+        else:
+            start_ind = 0
+        if end_ind is not None:
+            end_ind = max(min(end_ind - 1, dataset_len - 1), -dataset_len) % dataset_len + 1
+        else:
+            end_ind = dataset_len
+        assert start_ind < end_ind, f'Invalid start_ind and end_ind.'
+        self.start_ind = start_ind
+        self.end_ind = end_ind
+
+        if self.bucketize:
+            assert bucket_ids is not None and len(bucket_ids) == dataset_len
+            self.bucket_ids = [bucket_ids[self._map_idx(i)] for i in range(len(self))]
+
+        self.test_mode = test_mode
+
+    def get_bucket_ids_from_prompt_dataset(self):
+        ds = self.prompt_dataset
+        assert 'height' in ds.column_names and 'width' in ds.column_names, \
+            'When bucketize=True and no datalist is provided, the prompt dataset ' \
+            'must contain `height` and `width` columns.'
+        cols = ['height', 'width']
+        if 'frames' in ds.column_names:
+            cols = ['frames'] + cols
+        ds_arrow = ds.with_format('arrow', columns=cols)
+        batch = ds_arrow[:]
+
+        arrs = [batch[c].combine_chunks().to_numpy(zero_copy_only=False) for c in cols]
+        arrs = np.stack(arrs, axis=1)
+
+        _, inv = np.unique(arrs, axis=0, return_inverse=True)
+        return inv.tolist()
+
+    def parse_datalist(self, dir_path, datalist_path=None, datalist_must_exist=False):
+        logger = get_root_logger()
+
+        if datalist_path is not None and FileClient.infer_client(uri=datalist_path).isfile(datalist_path):
+            filenames = []
+            bucket_ids = []
+
+            datalist_bytesio = BytesIO(FileClient.infer_client(uri=datalist_path).get(datalist_path))
+            if datalist_path.endswith('.jsonl.gz') or datalist_path.endswith('.jsonl'):
+                if datalist_path.endswith('.jsonl.gz'):
+                    with gzip.open(datalist_bytesio, 'rt', encoding='utf-8') as f:
+                        datalist = f.readlines()
+                else:
+                    datalist = datalist_bytesio.read().decode('utf-8').splitlines()
+                for line in datalist:
+                    data_item = orjson.loads(line)
+                    if 'filename' in data_item:
+                        filenames.append(data_item['filename'])
+                    elif 'image_hash' in data_item:
+                        filenames.append(data_item['image_hash'])
+                    else:
+                        raise ValueError('No valid key to identify data item.')
+                    if self.bucketize:
+                        assert 'size_idx' in data_item, 'size_idx must be provided for bucketize.'
+                        bucket_ids.append(data_item['size_idx'])
+            elif datalist_path.endswith('.json'):
+                assert not self.bucketize, 'Bucketize not supported for json datalist.'
+                datalist = orjson.loads(datalist_bytesio.read())
+                for data_item in datalist:
+                    filenames.append(os.path.splitext(os.path.basename(data_item))[0])
+            else:
+                raise ValueError('Datalist file must be .jsonl, .jsonl.gz or .json')
+
+        else:
+            assert not datalist_must_exist, f'Datalist file {datalist_path} does not exist.'
+            assert not self.bucketize, 'Bucketize not supported when datalist is not provided.'
+            mmcv.print_log(
+                f'Datalist file {datalist_path} does not exist, directly list all files in the directory.',
+                logger=logger,
+                level=logging.WARNING)
+            # list all files in the directory
+            filenames = [os.path.splitext(p)[0] for p in self.file_client.list_dir_or_file(dir_path)]
+            filenames.sort()
+            bucket_ids = None
+            # save the datalist if datalist_path is provided
+            if datalist_path is not None:
+                if datalist_path.endswith('.jsonl.gz') or datalist_path.endswith('.jsonl'):
+                    datalist = []
+                    for filename in filenames:
+                        datalist.append(orjson.dumps({'filename': filename}).decode('utf-8'))
+                    datalist_str = '\n'.join(datalist)
+                    if datalist_path.endswith('.jsonl.gz'):
+                        datalist_bytesio = BytesIO()
+                        with gzip.open(datalist_bytesio, 'wt', encoding='utf-8') as f:
+                            f.write(datalist_str)
+                        FileClient.infer_client(uri=datalist_path).put(datalist_bytesio.getvalue(), datalist_path)
+                    else:
+                        FileClient.infer_client(uri=datalist_path).put_text(datalist_str, datalist_path)
+                elif datalist_path.endswith('.json'):
+                    datalist = filenames
+                    FileClient.infer_client(uri=datalist_path).put_text(
+                        orjson.dumps(datalist).decode('utf-8'), datalist_path)
+
+        mmcv.print_log(f'Loaded {len(filenames)} samples.', logger=logger)
+
+        return filenames, bucket_ids
+
+    def pad_prompt_embeds(self, prompt_embeds):
+        if self.pad_seq_len is not None:
+            if prompt_embeds.size(0) > self.pad_seq_len:
+                prompt_embeds = prompt_embeds[:self.pad_seq_len]
+            else:
+                zeros_size = (self.pad_seq_len - prompt_embeds.size(0),) + prompt_embeds.shape[1:]
+                prompt_embeds = torch.cat([prompt_embeds, prompt_embeds.new_zeros(zeros_size)], dim=0)
+        return prompt_embeds
+
+    def parse_prompt_embeds(self, data):
+        prompt_embed_kwargs = data.get('prompt_embed_kwargs', {}).copy()
+
+        # Map legacy keys to new ones if not already present
+        for legacy_key, new_key in self.PROMPT_KEY_MAPS.items():
+            if legacy_key in data and new_key not in prompt_embed_kwargs:
+                prompt_embed_kwargs[new_key] = data[legacy_key]
+
+        # Common post-processing
+        encoder_hidden_states_scale = prompt_embed_kwargs.pop('encoder_hidden_states_scale', None)
+        if 'encoder_hidden_states' in prompt_embed_kwargs:
+            encoder_hidden_states = prompt_embed_kwargs['encoder_hidden_states'].float()
+            if encoder_hidden_states_scale is not None:
+                encoder_hidden_states = encoder_hidden_states * encoder_hidden_states_scale
+            prompt_embed_kwargs['encoder_hidden_states'] = self.pad_prompt_embeds(encoder_hidden_states)
+
+        if 'pooled_projections' in prompt_embed_kwargs:
+            prompt_embed_kwargs['pooled_projections'] = prompt_embed_kwargs['pooled_projections'].float()
+
+        if 'encoder_hidden_states_mask' in prompt_embed_kwargs:
+            prompt_embed_kwargs['encoder_hidden_states_mask'] = self.pad_prompt_embeds(
+                prompt_embed_kwargs['encoder_hidden_states_mask'])
+
+        return prompt_embed_kwargs
+
+    def calculate_latent_size(self, image_spatial_size):
+        if isinstance(self.vae_scale_factor, int):
+            latent_spatial_size = tuple(s // self.vae_scale_factor for s in image_spatial_size)
+        else:
+            assert len(self.vae_scale_factor) == len(image_spatial_size)
+            latent_spatial_size = tuple(
+                s // f for s, f in zip(image_spatial_size, self.vae_scale_factor))
+        latent_size = (self.latent_size[0],) + latent_spatial_size
+        return latent_size
+
+    def calculate_scaled_image_size(self, image_spatial_size):
+        if self.image_scale_factor != 1:
+            if len(image_spatial_size) == 2:
+                new_spatial_size = (int(round(image_spatial_size[0] * self.image_scale_factor)),
+                                    int(round(image_spatial_size[1] * self.image_scale_factor)))
+            elif len(image_spatial_size) == 3:
+                new_spatial_size = (image_spatial_size[0],
+                                    int(round(image_spatial_size[1] * self.image_scale_factor)),
+                                    int(round(image_spatial_size[2] * self.image_scale_factor)))
+            else:
+                raise ValueError(f'Unsupported image spatial size {image_spatial_size}.')
+        else:
+            new_spatial_size = image_spatial_size
+        return new_spatial_size
+
+    def scale_image(self, image):
+        if self.image_scale_factor != 1:
+            new_spatial_size = self.calculate_scaled_image_size(image.shape[1:])
+            if len(new_spatial_size) == 2:
+                image = F.interpolate(
+                    image[None], size=new_spatial_size, mode='bicubic', align_corners=False, antialias=True
+                )[0].clamp(min=0, max=1)
+            elif len(new_spatial_size) == 3:
+                image = F.interpolate(
+                    image, size=new_spatial_size[1:], mode='bicubic', align_corners=False, antialias=True
+                ).clamp(min=0, max=1)
+            else:
+                raise ValueError(f'Unsupported image spatial size {image.shape[1:]}.')
+        return image
+
+    def _map_idx(self, idx):
+        return self.start_ind + (idx // self.repeat)
+
+    def __len__(self):
+        return self.repeat * (self.end_ind - self.start_ind)
+
+    def __getitem__(self, idx):
+        mapped_idx = self._map_idx(idx)
+
+        prompt_data = None
+
+        if self.cache_dir_path is not None:
+            data_path = self.file_client.join_path(
+                self.cache_dir_path, f'{self.cache_datalist[mapped_idx]}.zst')
+            data_bytesio = BytesIO(self.file_client.get(data_path))
+            with zstd.ZstdDecompressor().stream_reader(data_bytesio) as f:
+                raw_data = pickle.load(f)
+            data = dict(
+                ids=DC(idx, cpu_only=True),
+                name=DC(raw_data['prompt'], cpu_only=True),
+                prompt_embed_kwargs=self.parse_prompt_embeds(raw_data))
+
+            if not self.ignore_cached_latents:  # load latents
+                if 'latents' in raw_data:
+                    latents = raw_data['latents']
+                    if self.test_mode:
+                        data['noise'] = torch.randn(
+                            latents.size(), dtype=torch.float32, generator=torch.Generator().manual_seed(idx))
+                    else:
+                        data['latents'] = latents.float()
+                        latents_scale = raw_data.get('latents_scale', None)
+                        if latents_scale is not None:
+                            data['latents'] = data['latents'] * latents_scale
+                else:
+                    latent_size = raw_data.get('latent_size', self.latent_size)
+                    if self.test_mode:
+                        data['noise'] = torch.randn(
+                            latent_size, dtype=torch.float32, generator=torch.Generator().manual_seed(idx))
+                    else:
+                        data['latents'] = torch.empty(latent_size, dtype=torch.float32)
+
+        else:
+            prompt_data = self.prompt_dataset[mapped_idx]
+            if 'prompt_kwargs' in prompt_data:
+                prompt_kwargs = {k: DC(v, cpu_only=True) for k, v in prompt_data['prompt_kwargs'].items()}
+            else:
+                prompt_kwargs = dict(prompt=DC(prompt_data['prompt'], cpu_only=True))
+            data = dict(
+                ids=DC(idx, cpu_only=True),
+                name=DC(prompt_data['prompt'], cpu_only=True),
+                prompt_kwargs=prompt_kwargs)
+
+        if self.image_dir_path is not None:
+            image_path = self.file_client.join_path(
+                self.image_dir_path, self.image_datalist[mapped_idx] + self.image_extension)
+            image = load_image(image_path, self.file_client)
+            image = np.moveaxis(image, -1, 0)  # channel first
+            if self.test_mode:
+                data['noise'] = torch.randn(
+                    self.calculate_latent_size(self.calculate_scaled_image_size(image.shape[1:])),
+                    dtype=torch.float32, generator=torch.Generator().manual_seed(idx))
+            else:
+                images = torch.from_numpy(image)
+                if images.dtype == torch.uint8:
+                    images = images.float() / 255.0
+                assert torch.is_floating_point(images), f'Image dtype {images.dtype} not supported.'
+                data['images'] = self.scale_image(images.float())
+        elif 'latents' not in data and 'noise' not in data:  # allocate latents if not already loaded
+            if prompt_data is not None and 'height' in prompt_data and 'width' in prompt_data:
+                image_spatial_size = (prompt_data['height'], prompt_data['width'])
+                if 'frames' in prompt_data:
+                    image_spatial_size = (prompt_data['frames'],) + image_spatial_size
+                latent_size = self.calculate_latent_size(self.calculate_scaled_image_size(image_spatial_size))
+            else:
+                latent_size = self.latent_size
+            if self.test_mode:
+                data['noise'] = torch.randn(
+                    latent_size, dtype=torch.float32, generator=torch.Generator().manual_seed(idx))
+            else:
+                data['latents'] = torch.empty(latent_size, dtype=torch.float32)
+
+        if self.negative_prompt_embed_kwargs is not None:
+            data.update(negative_prompt_embed_kwargs=self.negative_prompt_embed_kwargs)
+        if self.negative_prompt_kwargs is not None:
+            data.update(negative_prompt_kwargs=self.negative_prompt_kwargs)
+
+        return data

lakonlab/datasets/imagenet.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import os
+
+import numpy as np
+import torch
+import mmcv
+
+from io import BytesIO
+from PIL import Image
+from torch.utils.data import Dataset
+from mmcv.fileio import FileClient
+from mmcv.parallel import DataContainer as DC
+from mmgen.datasets.builder import DATASETS
+from mmgen.utils import get_root_logger
+
+
+def image_preproc(pil_image, image_size, random_flip=False):
+    """
+    Center cropping implementation from ADM.
+    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+    """
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    arr = arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]
+
+    if random_flip and np.random.rand() < 0.5:
+        arr = np.ascontiguousarray(arr[:, ::-1])
+
+    if arr.ndim == 2:
+        arr = np.stack([arr] * 3, axis=-1)
+    elif arr.ndim == 3:
+        if arr.shape[2] == 1:
+            arr = np.concatenate([arr] * 3, axis=-1)
+        elif arr.shape[2] == 4:
+            arr = arr[:, :, :3]
+        else:
+            assert arr.shape[2] == 3
+    else:
+        raise ValueError(f'Unexpected number of dimensions: {arr.ndim}')
+    return arr
+
+
+@DATASETS.register_module()
+class ImageNet(Dataset):
+    def __init__(
+            self,
+            data_root='data/imagenet/train',
+            datalist_path='data/imagenet/train.txt',
+            label2name_path='data/imagenet/imagenet1000_clsidx_to_labels.txt',
+            random_flip=True,
+            negative_label=1000,
+            image_size=256,
+            latent_size=(4, 32, 32),
+            test_label_repeat=1,
+            test_mode=False,
+            num_test_images=50000):
+        super().__init__()
+        self.data_root = data_root
+        self.file_client = FileClient.infer_client(uri=self.data_root)
+
+        self.datalist_path = datalist_path
+        self.label2name_path = label2name_path
+        self.random_flip = random_flip
+        self.negative_label = negative_label
+        self.image_size = image_size
+        self.latent_size = latent_size
+        self.test_label_repeat = test_label_repeat
+        self.test_mode = test_mode
+        self.num_test_images = num_test_images
+
+        self.label2name = {}
+        label2name_text = FileClient.infer_client(uri=self.label2name_path).get_text(self.label2name_path)
+        for line in label2name_text.split('\n'):
+            line = line.strip()
+            if len(line) == 0:
+                continue
+            idx, name = line.split(':')
+            idx, name = idx.strip(), name.strip()
+            if name[-1] == ',':
+                name = name[:-1]
+            if name[0] == '"' and name[-1] == '"':
+                name = name[1:-1]
+            if name[0] == "'" and name[-1] == "'":
+                name = name[1:-1]
+            self.label2name[int(idx)] = name
+
+        if not test_mode:
+            self.all_paths = []
+            self.all_labels = []
+            datalist_text = FileClient.infer_client(uri=self.datalist_path).get_text(self.datalist_path)
+            for line in datalist_text.split('\n'):
+                line = line.strip()
+                if len(line) == 0:
+                    continue
+                path_label = line.split(' ')
+                self.all_paths.append(path_label[0])
+                if len(path_label) > 1:
+                    self.all_labels.append(int(path_label[1]))
+
+            logger = get_root_logger()
+            mmcv.print_log(f'Data root: {self.data_root}', logger=logger)
+            mmcv.print_log(f'Data list path: {self.datalist_path}', logger=logger)
+            mmcv.print_log(f'Number of images: {len(self.all_paths)}', logger=logger)
+
+    def __len__(self):
+        return self.num_test_images if self.test_mode else len(self.all_paths)
+
+    def __getitem__(self, idx):
+        data = dict(ids=DC(idx, cpu_only=True))
+
+        if self.test_mode:
+            label_generator = torch.Generator().manual_seed(idx // self.test_label_repeat)
+            label = torch.randint(0, 1000, (), generator=label_generator).long()
+            noise_generator = torch.Generator().manual_seed(idx + 1000)
+            noise = torch.randn(self.latent_size, generator=noise_generator)
+            data.update(noise=noise)
+
+        else:
+            rel_data_path = self.all_paths[idx]
+            data.update(paths=DC(rel_data_path, cpu_only=True))
+            data_path = self.file_client.join_path(self.data_root, rel_data_path)
+            data_bytesio = BytesIO(self.file_client.get(data_path))
+            ext = os.path.splitext(data_path)[-1]
+            if ext.lower() in ('.pth', '.pt'):
+                torch_data = torch.load(data_bytesio, map_location='cpu')
+                label = torch_data['y'].long()
+                data.update(latents=torch_data['x'].float())
+            elif ext.lower() in ('.jpg', '.jpeg', '.png'):
+                label = torch.tensor(self.all_labels[idx], dtype=torch.long)
+                img_data = Image.open(data_bytesio)
+                data.update(
+                    images=torch.from_numpy(image_preproc(
+                        img_data, self.image_size, random_flip=self.random_flip)).float().permute(2, 0, 1) / 255.0)
+            else:
+                raise ValueError(f'Unsupported file extension: {ext}')
+
+        name = self.label2name[label.item()]
+        data.update(labels=label, name=DC(name, cpu_only=True))
+
+        if self.negative_label is not None:
+            if isinstance(self.negative_label, int):
+                data.update(negative_labels=torch.tensor(self.negative_label, dtype=torch.long))
+            else:
+                raise ValueError(f'Unsupported negative label: {self.negative_label}')
+
+        return data

lakonlab/datasets/samplers/__init__.py

@@ -0,0 +1 @@
+from .distributed_sampler import DistributedSampler

lakonlab/datasets/samplers/distributed_sampler.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import numpy as np
+import torch
+
+from torch.utils.data import DistributedSampler as _DistributedSampler
+from mmgen.utils import sync_random_seed
+
+
+def reverse_index_map(bucket_ids):
+    bucket_map = dict()
+    for data_id, bucket_id in enumerate(bucket_ids):
+        if bucket_id not in bucket_map:
+            bucket_map[bucket_id] = []
+        bucket_map[bucket_id].append(data_id)
+    return bucket_map
+
+
+class DistributedSampler(_DistributedSampler):
+
+    def __init__(self,
+                 dataset,
+                 num_replicas=None,
+                 rank=None,
+                 shuffle=True,
+                 samples_per_gpu=1,
+                 seed=None):
+        super().__init__(dataset, num_replicas=num_replicas, rank=rank)
+
+        self.shuffle = shuffle
+        self.samples_per_gpu = samples_per_gpu
+
+        self.bucket_map = self.total_size_bucketwise = None
+        if hasattr(dataset, 'bucket_ids'):
+            self._init_bucket_sampler(dataset)
+        else:
+            self._init_sampler(dataset)
+
+        self.seed = sync_random_seed(seed)
+        self.skip_iter = 0
+
+    def _init_sampler(self, dataset):
+        data_len = len(dataset)
+        # to avoid padding bug when meeting too small dataset
+        if data_len < self.num_replicas * self.samples_per_gpu:
+            raise ValueError(
+                'You may use too small dataset and our distributed '
+                'sampler cannot pad your dataset correctly. Please '
+                'use fewer GPUs or smaller batch sizes per GPU.')
+
+        num_batches = int(np.ceil(data_len / self.num_replicas / self.samples_per_gpu))
+        self.num_samples = num_batches * self.samples_per_gpu
+        self.total_size = self.num_samples * self.num_replicas
+
+    def _init_bucket_sampler(self, dataset):
+        self.bucket_map = reverse_index_map(dataset.bucket_ids)
+        self.bucket_map = dict(sorted(self.bucket_map.items()))  # sort by bucket_id
+
+        data_len = 0
+        self.total_size_bucketwise = {}
+
+        for bucket_id, data_indices in self.bucket_map.items():
+            _data_len = len(data_indices)
+            if _data_len < self.samples_per_gpu:
+                raise ValueError(
+                    'You may use too small dataset and our distributed '
+                    'sampler cannot pad your dataset correctly. Please '
+                    'use smaller batch sizes per GPU.')
+
+            _total_num_batches = int(np.ceil(_data_len / self.samples_per_gpu))
+            _total_size = _total_num_batches * self.samples_per_gpu
+
+            data_len += _total_size
+            self.total_size_bucketwise[bucket_id] = _total_size
+
+        if data_len < self.num_replicas * self.samples_per_gpu:
+            raise ValueError(
+                'You may use too small dataset and our distributed '
+                'sampler cannot pad your dataset correctly. Please '
+                'use fewer GPUs or smaller batch sizes per GPU.')
+
+        num_batches = int(np.ceil(data_len / self.num_replicas / self.samples_per_gpu))
+        self.num_samples = num_batches * self.samples_per_gpu
+        self.total_size = self.num_samples * self.num_replicas
+
+    def update_sampler(self, dataset, samples_per_gpu=None):
+        self.dataset = dataset
+        if samples_per_gpu is not None:
+            self.samples_per_gpu = samples_per_gpu
+        self.bucket_map = self.total_size_bucketwise = None
+        if hasattr(dataset, 'bucket_ids'):
+            self._init_bucket_sampler(dataset)
+        else:
+            self._init_sampler(dataset)
+
+    def set_iter(self, iteration):
+        num_batches = self.num_samples // self.samples_per_gpu
+        self.skip_iter = iteration % num_batches
+
+    def __iter__(self):
+        if self.bucket_map is None:
+            if self.shuffle:
+                g = torch.Generator()
+                g.manual_seed(self.seed + self.epoch)
+                indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            else:
+                indices = torch.arange(len(self.dataset)).tolist()
+            # add extra samples to make it evenly divisible
+            indices += indices[:(self.total_size - len(indices))]
+            assert len(indices) == self.total_size
+            # subsample
+            indices = indices[self.rank:self.total_size:self.num_replicas]
+
+        else:  # guarantees that batch samples are from the same bucket
+            if self.shuffle:
+                g = torch.Generator()
+                g.manual_seed(self.seed + self.epoch)
+            else:
+                g = None
+            indices = []
+            for bucket_id, data_indices in self.bucket_map.items():
+                data_indices = torch.tensor(data_indices)
+                if g is not None:
+                    data_indices = data_indices[torch.randperm(len(data_indices), generator=g)]
+                pad = self.total_size_bucketwise[bucket_id] - data_indices.numel()
+                if pad:
+                    data_indices = torch.cat([data_indices, data_indices[:pad]], dim=0)
+                assert data_indices.numel() == self.total_size_bucketwise[bucket_id]
+                _total_num_batches = self.total_size_bucketwise[bucket_id] // self.samples_per_gpu
+                _num_batches = _total_num_batches // self.num_replicas
+                _total_leftover_batches = _total_num_batches % self.num_replicas
+                # data_indices_a: evenly split batches for full round-robins across replicas
+                # data_indices_b: the leftover partial round-robin
+                data_indices_a = data_indices[:(_num_batches * self.num_replicas * self.samples_per_gpu)].reshape(
+                    _num_batches, self.samples_per_gpu, self.num_replicas
+                ).permute(0, 2, 1).reshape(
+                    _num_batches * self.num_replicas, self.samples_per_gpu)
+                data_indices_b = data_indices[(_num_batches * self.num_replicas * self.samples_per_gpu):].reshape(
+                    self.samples_per_gpu, _total_leftover_batches
+                ).permute(1, 0)
+                indices.extend([data_indices_a, data_indices_b])
+            indices = torch.cat(indices, dim=0)  # (total_num_batches, samples_per_gpu)
+            if g is not None:
+                indices = indices[torch.randperm(indices.size(0), generator=g)]
+            total_num_batches = self.total_size // self.samples_per_gpu
+            pad = total_num_batches - indices.size(0)
+            if pad:
+                indices = torch.cat([indices, indices[:pad]], dim=0)
+            assert indices.numel() == self.total_size
+            indices = indices[self.rank:total_num_batches:self.num_replicas].flatten().tolist()
+
+        assert len(indices) == self.num_samples
+        skip_len = self.skip_iter * self.samples_per_gpu
+        assert skip_len < self.num_samples
+        indices = indices[skip_len:]
+        self.skip_iter = 0
+
+        return iter(indices)

lakonlab/evaluation/__init__.py

@@ -0,0 +1,8 @@
+from .metrics import FIDKID, PR, InceptionMetrics, ColorStats, HPSv2, CLIPSimilarity
+from .vqa_score import VQAScore
+from .hpsv3 import HPSv3
+from .eval_hooks import GenerativeEvalHook
+
+__all__ = ['GenerativeEvalHook', 'FIDKID', 'PR',
+           'InceptionMetrics', 'ColorStats', 'HPSv2', 'VQAScore', 'CLIPSimilarity',
+           'HPSv3']

lakonlab/evaluation/eval_hooks.py

@@ -0,0 +1,318 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import sys
+import os
+import re
+import unicodedata
+import numpy as np
+import torch
+import torch.distributed as dist
+import mmcv
+
+from copy import deepcopy
+from concurrent.futures import ThreadPoolExecutor
+from mmcv.runner import HOOKS, get_dist_info
+from mmcv.fileio import FileClient
+from mmgen.models.architectures.common import get_module_device
+from mmgen.core import GenerativeEvalHook as _GenerativeEvalHook
+from lakonlab.utils.io_utils import save_image, save_video, load_images_parallel
+from lakonlab.runner.timer import default_timers
+from lakonlab.utils import gc_context
+from lakonlab.ui.media_viewer import write_html
+
+default_timers.add_timer('total time')
+
+_reserved = {"CON", "PRN", "AUX", "NUL", *(f"COM{i}" for i in range(1, 10)), *(f"LPT{i}" for i in range(1, 10))}
+_invalid = re.compile(r'[<>:"/\\|?*\:%]+')
+
+
+def flatten_list(lst):
+    for item in lst:
+        if isinstance(item, list):
+            yield from flatten_list(item)  # recurse into sub-list
+        else:
+            yield item
+
+
+def _safe_name(name: str, max_len: int = 160) -> str:
+    s = unicodedata.normalize("NFKC", str(name))
+    s = "".join(c for c in s if 32 <= ord(c) != 127)  # drop control chars
+    s = _invalid.sub("", s).strip(" ._")  # strip invalid chars
+    if s.startswith("."):
+        s = s.lstrip(".")  # avoid hidden files
+    if s.upper() in _reserved:
+        s += "_"  # avoid reserved names
+    s = (s or "untitled")[:max_len].rstrip(" .")  # truncate & clean
+    return s or "untitled"
+
+
+def evaluate(model, dataloader, metrics=None,
+             feed_batch_size=32, viz_dir=None, viz_num=None, sample_kwargs=dict(),
+             fps=16, enable_timers=False, reuse_viz=False):
+    has_metrics = metrics is not None and len(metrics) > 0
+    if has_metrics:
+        for metric in metrics:
+            if hasattr(metric, 'load_to_gpu'):
+                metric.load_to_gpu()
+
+    if enable_timers:
+        default_timers.enable_all()
+
+    batch_size = dataloader.batch_size
+    rank, ws = get_dist_info()
+    total_batch_size = batch_size * ws
+
+    max_num_fakes = len(dataloader.dataset)
+
+    if viz_dir is not None:
+        html_entries = []
+        saved_data_ids = set()
+        file_client = FileClient.infer_client(uri=viz_dir)
+        executor = ThreadPoolExecutor(max_workers=(os.cpu_count() or 4) * 4)
+    else:
+        html_entries = saved_data_ids = file_client = executor = None
+
+    if rank == 0:
+        mmcv.print_log(
+            f'Generate {max_num_fakes} fake samples for evaluation', 'mmgen')
+        pbar = mmcv.ProgressBar(max_num_fakes)
+
+    log_vars = dict()
+    batch_size_list = []
+
+    for i, data in enumerate(dataloader):
+        can_reuse = False
+        loaded_imgs_tensor = None
+        loaded_html_entries = None
+
+        if viz_dir is not None and reuse_viz:
+            batch_names = list(flatten_list(data['name'].data))
+            ids = list(flatten_list(data['ids'].data))
+
+            reuse_candidates = []
+            all_png_exist = True
+
+            for name, data_id in zip(batch_names, ids):
+                safe_name = _safe_name(name)
+                filename = f'{data_id:09d}_{safe_name}.png'
+                filepath = file_client.join_path(viz_dir, filename)
+                if file_client.isfile(filepath):
+                    reuse_candidates.append((data_id, name, filename, filepath))
+                else:
+                    all_png_exist = False
+                    break
+
+            if all_png_exist:
+                filepaths = []
+                loaded_html_entries = []
+                for (data_id, name, filename, filepath) in reuse_candidates:
+                    filepaths.append(filepath)
+                    rel_filepath = file_client.join_path(os.path.basename(viz_dir), filename)
+                    loaded_html_entries.append((data_id, rel_filepath, name))
+                try:
+                    loaded_imgs_tensor = torch.from_numpy(
+                        np.stack(load_images_parallel(filepaths, file_client), axis=0)
+                    ).permute(0, 3, 1, 2).float() / 255.0
+                    can_reuse = True
+                except Exception as e:
+                    can_reuse = False
+                    loaded_imgs_tensor = None
+                    loaded_html_entries = None
+
+        if can_reuse:
+            outputs_dict = dict(
+                pred_imgs=loaded_imgs_tensor,  # [0,1] float32
+                num_samples=loaded_imgs_tensor.size(0))
+            html_entries.extend(loaded_html_entries)
+
+        else:
+            # fall back to generation
+            sample_kwargs_ = deepcopy(sample_kwargs)
+
+            with default_timers['total time']:
+                outputs_dict = model.val_step(data, show_pbar=rank == 0, **sample_kwargs_)
+
+            if viz_dir is not None:
+                batch_names = list(flatten_list(data['name'].data))
+                for batch_id, data_id in enumerate(flatten_list(data['ids'].data)):
+                    if (viz_num is not None and data_id >= viz_num) or (data_id in saved_data_ids):
+                        continue
+                    name = batch_names[batch_id]
+                    safe_name = _safe_name(name)
+
+                    image_viz = (outputs_dict['pred_imgs'][batch_id] * 255).round().to(torch.uint8)
+                    if image_viz.dim() == 3:  # image
+                        image_viz = image_viz.permute(1, 2, 0).cpu().numpy()
+                        filename = f'{data_id:09d}_{safe_name}.png'
+                        executor.submit(
+                            save_image,
+                            image_viz, file_client.join_path(viz_dir, filename), file_client)
+                    elif image_viz.dim() == 4:  # video
+                        image_viz = image_viz.permute(1, 2, 3, 0).cpu().numpy()  # (t, h, w, c)
+                        filename = f'{data_id:09d}_{safe_name}.mp4'
+                        executor.submit(
+                            save_video,
+                            image_viz, file_client.join_path(viz_dir, filename), file_client, fps)
+                    else:
+                        raise ValueError(f'Unsupported image dimension: {image_viz.dim()}')
+                    rel_filepath = file_client.join_path(os.path.basename(viz_dir), filename)
+                    html_entries.append((data_id, rel_filepath, name))
+                    saved_data_ids.add(data_id)
+
+        if 'log_vars' in outputs_dict:
+            for k, v in outputs_dict['log_vars'].items():
+                if k in log_vars:
+                    log_vars[k].append(outputs_dict['log_vars'][k])
+                else:
+                    log_vars[k] = [outputs_dict['log_vars'][k]]
+        batch_size_list.append(outputs_dict['num_samples'])
+
+        if has_metrics:
+            pred_imgs = outputs_dict['pred_imgs'].split(feed_batch_size, dim=0)
+            real_imgs = None
+            if 'images' in data:
+                real_imgs = data['images']
+            elif 'target_imgs' in outputs_dict:
+                real_imgs = outputs_dict['target_imgs']
+            if real_imgs is not None:
+                real_imgs = real_imgs.split(feed_batch_size, dim=0)
+            requires_prompt = False
+            for metric in metrics:
+                requires_prompt |= getattr(metric, 'requires_prompt', False)
+            if requires_prompt:
+                prompts = list(flatten_list(data['name'].data))  # list of prompts
+                prompts = [prompts[i:i + feed_batch_size] for i in range(0, len(prompts), feed_batch_size)]
+            for metric in metrics:
+                for batch_id, batch_imgs in enumerate(pred_imgs):
+                    if getattr(metric, 'requires_prompt', False):
+                        metric.feed(
+                            dict(imgs=batch_imgs * 2 - 1, prompts=prompts[batch_id]), 'fakes')
+                        if real_imgs is not None:
+                            metric.feed(
+                                dict(imgs=real_imgs[batch_id] * 2 - 1, prompts=prompts[batch_id]), 'reals')
+                    else:
+                        metric.feed(batch_imgs * 2 - 1, 'fakes')
+                        if real_imgs is not None:
+                            metric.feed(real_imgs[batch_id] * 2 - 1, 'reals')
+
+        if rank == 0:
+            pbar.update(total_batch_size)
+
+    if ws > 1:
+        device = get_module_device(model)
+        batch_size_list = torch.tensor(batch_size_list, dtype=torch.float, device=device)
+        batch_size_sum = torch.sum(batch_size_list)
+        dist.all_reduce(batch_size_sum, op=dist.ReduceOp.SUM)
+        for k, v in log_vars.items():
+            weigted_values = torch.tensor(log_vars[k], dtype=torch.float, device=device) * batch_size_list
+            weigted_values_sum = torch.sum(weigted_values)
+            dist.all_reduce(weigted_values_sum, op=dist.ReduceOp.SUM)
+            log_vars[k] = float(weigted_values_sum / batch_size_sum)
+    else:
+        for k, v in log_vars.items():
+            log_vars[k] = np.average(log_vars[k], weights=batch_size_list)
+
+    if viz_dir is not None:
+        if ws > 1:
+            gathered = [None for _ in range(ws)]
+            dist.all_gather_object(gathered, html_entries)
+            if rank == 0:
+                html_entries = [e for sub in gathered for e in (sub or [])]
+        if rank == 0:
+            unique_entries = dict()
+            for entry in html_entries:
+                data_id = entry[0]
+                if data_id not in unique_entries:
+                    unique_entries[data_id] = entry
+            html_entries = list(unique_entries.values())
+            html_entries.sort(key=lambda item: item[0])
+            html_path = file_client.join_path(os.path.dirname(viz_dir), os.path.basename(viz_dir) + '.html')
+            write_html(html_path, html_entries, file_client)
+        executor.shutdown(wait=True)
+
+    return log_vars
+
+
+@HOOKS.register_module(force=True)
+class GenerativeEvalHook(_GenerativeEvalHook):
+    greater_keys = ['acc', 'top', 'AR@', 'auc', 'precision', 'mAP', 'is', 'test_ssim', 'test_psnr']
+    less_keys = ['loss', 'fid', 'kid', 'test_lpips']
+    _supported_best_metrics = ['fid', 'kid', 'is', 'test_ssim', 'test_psnr', 'test_lpips']
+
+    def __init__(self,
+                 *args,
+                 data='',
+                 viz_dir=None,
+                 feed_batch_size=32,
+                 viz_num=None,
+                 clear_reals=False,
+                 prefix='',
+                 metric_cpu_offload=False,
+                 **kwargs):
+        super(GenerativeEvalHook, self).__init__(*args, **kwargs)
+        self.data = data
+        self.viz_dir = viz_dir
+        self.file_client = FileClient.infer_client(
+            uri=viz_dir) if viz_dir is not None else None
+        self.feed_batch_size = feed_batch_size
+        self.viz_num = viz_num
+        self.clear_reals = clear_reals
+        self.prefix = prefix
+        self.metric_cpu_offload = metric_cpu_offload
+
+    @torch.no_grad()
+    def after_train_iter(self, runner):
+        with gc_context(enable=True):
+            interval = self.get_current_interval(runner)
+            if not self.every_n_iters(runner, interval):
+                return
+
+            runner.model.eval()
+            rank, ws = get_dist_info()
+
+            if self.viz_dir is not None:
+                viz_dir = self.file_client.join_path(self.viz_dir, str(runner.iter + 1))
+                if rank == 0:
+                    if self.file_client.exists(viz_dir):
+                        for name in self.file_client.list_dir_or_file(viz_dir):
+                            self.file_client.remove(self.file_client.join_path(viz_dir, name))
+                if ws > 1:
+                    dist.barrier()
+            else:
+                viz_dir = None
+            log_vars = evaluate(
+                runner.model, self.dataloader, self.metrics, self.feed_batch_size,
+                viz_dir, self.viz_num, self.sample_kwargs)
+
+            if len(runner.log_buffer.output) == 0:
+                runner.log_buffer.clear()
+
+            # a dirty walkround to change the line at the end of pbar
+            if rank == 0:
+                sys.stdout.write('\n')
+                for metric in self.metrics:
+                    metric.summary()
+                    for name, val in metric._result_dict.items():
+                        prefix_name = self.prefix + '_' + name if len(self.prefix) > 0 else name
+                        runner.log_buffer.output[self.data + '_' + prefix_name] = val
+                        # record best metric and save the best ckpt
+                        if self.save_best_ckpt and name in self.best_metric:
+                            self._save_best_ckpt(runner, val, name)
+                for name, val in log_vars.items():
+                    prefix_name = self.prefix + '_' + name if len(self.prefix) > 0 else name
+                    # print(self.data + '_' + prefix_name + ' = {}'.format(val))
+                    runner.log_buffer.output[self.data + '_' + prefix_name] = val
+                    # record best metric and save the best ckpt
+                    if self.save_best_ckpt and name in self.best_metric:
+                        self._save_best_ckpt(runner, val, name)
+                runner.log_buffer.ready = True
+
+            runner.model.train()
+
+            for metric in self.metrics:
+                metric.clear(clear_reals=self.clear_reals)
+                if self.metric_cpu_offload:
+                    if hasattr(metric, 'offload_to_cpu'):
+                        metric.offload_to_cpu()
+
+            torch.cuda.empty_cache()

lakonlab/evaluation/hpsv3.py

@@ -0,0 +1,587 @@
+# Modified from https://github.com/MizzenAI/HPSv3
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+import mmcv
+
+from typing import List, Optional, Union
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, FullyShardedDataParallel
+from torch.distributed.fsdp.wrap import ModuleWrapPolicy
+from accelerate import init_empty_weights
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.image_utils import (
+    OPENAI_CLIP_MEAN,
+    OPENAI_CLIP_STD,
+)
+from transformers.utils import TensorType
+from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLVisionBlock, Qwen2VLDecoderLayer
+from mmcv.runner import get_dist_info
+from mmgen.core.registry import METRICS
+from mmgen.core.evaluation.metrics import Metric
+from lakonlab.runner.checkpoint import _load_checkpoint
+
+
+INSTRUCTION = """
+You are tasked with evaluating a generated image based on Visual Quality and Text Alignment and give a overall score to estimate the human preference. Please provide a rating from 0 to 10, with 0 being the worst and 10 being the best. 
+
+**Visual Quality:**  
+Evaluate the overall visual quality of the image. The following sub-dimensions should be considered:
+- **Reasonableness:** The image should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups.
+- **Clarity:** Evaluate the sharpness and visibility of the image. The image should be clear and easy to interpret, with no blurring or indistinct areas.
+- **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows).
+- **Aesthetic and Creativity:** Assess the artistic aspects of the image, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance.
+- **Safety:** The image should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. 
+
+**Text Alignment:**  
+Assess how well the image matches the textual prompt across the following sub-dimensions:
+- **Subject Relevance** Evaluate how accurately the subject(s) in the image (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior.
+- **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the image adheres to this style.
+- **Contextual Consistency**: Assess whether the background, setting, and surrounding elements in the image logically fit the scenario described in the prompt. The environment should support and enhance the subject without contradictions.
+- **Attribute Fidelity**: Check if specific attributes mentioned in the prompt (e.g., colors, clothing, accessories, expressions, actions) are faithfully represented in the image. Minor deviations may be acceptable, but critical attributes should be preserved.
+- **Semantic Coherence**: Evaluate whether the overall meaning and intent of the prompt are captured in the image. The generated content should not introduce elements that conflict with or distort the original description.
+Textual prompt - {text_prompt}
+
+
+"""
+
+prompt_with_special_token = """
+Please provide the overall ratings of this image: <|Reward|>
+
+END
+"""
+
+prompt_without_special_token = """
+Please provide the overall ratings of this image: 
+"""
+
+
+def smart_resize(
+        height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280):
+    """Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+
+    """
+    if height < factor or width < factor:
+        raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+    elif max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = round(height / factor) * factor
+    w_bar = round(width / factor) * factor
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = math.floor(height / beta / factor) * factor
+        w_bar = math.floor(width / beta / factor) * factor
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+    return h_bar, w_bar
+
+
+class Qwen2VLImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]
+
+    def __init__(
+            self,
+            do_resize: bool = True,
+            do_normalize: bool = True,
+            image_mean: Optional[Union[float, List[float]]] = None,
+            image_std: Optional[Union[float, List[float]]] = None,
+            min_pixels: int = 256 * 28 * 28,
+            max_pixels: int = 256 * 28 * 28,
+            patch_size: int = 14,
+            temporal_patch_size: int = 2,
+            merge_size: int = 2,
+            **kwargs):
+        super().__init__(**kwargs)
+        self.do_resize = do_resize
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
+        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.merge_size = merge_size
+
+    def _preprocess(
+            self,
+            images: torch.Tensor,
+            do_resize: bool = None,
+            do_normalize: bool = None,
+            image_mean: Optional[Union[float, List[float]]] = None,
+            image_std: Optional[Union[float, List[float]]] = None):
+        batch_size, channel, height, width = images.size()
+        if do_resize:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=self.patch_size * self.merge_size,
+                min_pixels=self.min_pixels,
+                max_pixels=self.max_pixels,
+            )
+            images = F.interpolate(
+                images,
+                size=(resized_height, resized_width),
+                mode='bicubic',
+                align_corners=False,
+                antialias=True,
+            ).clamp(min=0, max=1)
+        else:
+            resized_height, resized_width = height, width
+
+        if do_normalize:
+            mean = torch.tensor(image_mean, device=images.device, dtype=images.dtype).view(-1, 1, 1)
+            std = torch.tensor(image_std, device=images.device, dtype=images.dtype).view(-1, 1, 1)
+            images = (images - mean) / std
+
+        patches = images.unsqueeze(1).expand(-1, self.temporal_patch_size, -1, -1, -1)
+
+        grid_t = 1
+        grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
+
+        patches = patches.reshape(
+            batch_size * grid_t,
+            self.temporal_patch_size,
+            channel,
+            grid_h // self.merge_size,
+            self.merge_size,
+            self.patch_size,
+            grid_w // self.merge_size,
+            self.merge_size,
+            self.patch_size,
+        )
+        patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
+        flatten_patches = patches.reshape(
+            batch_size * grid_t * grid_h * grid_w, channel * self.temporal_patch_size * self.patch_size * self.patch_size
+        )
+
+        return flatten_patches, np.array((grid_t, grid_h, grid_w)).reshape(1, 3).repeat(batch_size, axis=0)
+
+    def preprocess(
+            self,
+            images: torch.Tensor,
+            return_tensors: Optional[Union[str, TensorType]] = None):
+        pixel_values, vision_grid_thws = self._preprocess(
+            images,
+            do_resize=self.do_resize,
+            do_normalize=self.do_normalize,
+            image_mean=self.image_mean,
+            image_std=self.image_std,
+        )
+        data = {"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration):
+
+    def __init__(
+        self,
+        config,
+        output_dim=4,
+        reward_token="last",
+        special_token_ids=None,
+        rm_head_type="default",
+        rm_head_kwargs=None,
+    ):
+        super().__init__(config)
+        # pdb.set_trace()
+        self.output_dim = output_dim
+        if rm_head_type == "default":
+            self.rm_head = nn.Linear(config.hidden_size, output_dim, bias=False)
+        elif rm_head_type == "ranknet":
+            if rm_head_kwargs is not None:
+                for layer in range(rm_head_kwargs.get("num_layers", 3)):
+                    if layer == 0:
+                        self.rm_head = nn.Sequential(
+                            nn.Linear(config.hidden_size, rm_head_kwargs["hidden_size"]),
+                            nn.ReLU(),
+                            nn.Dropout(rm_head_kwargs.get("dropout", 0.1)),
+                        )
+                    elif layer < rm_head_kwargs.get("num_layers", 3) - 1:
+                        self.rm_head.add_module(
+                            f"layer_{layer}",
+                            nn.Sequential(
+                                nn.Linear(rm_head_kwargs["hidden_size"], rm_head_kwargs["hidden_size"]),
+                                nn.ReLU(),
+                                nn.Dropout(rm_head_kwargs.get("dropout", 0.1)),
+                            ),
+                        )
+                    else:
+                        self.rm_head.add_module(
+                            f"output_layer",
+                            nn.Linear(rm_head_kwargs["hidden_size"], output_dim, bias=rm_head_kwargs.get("bias", False)),
+                        )
+
+            else:
+                self.rm_head = nn.Sequential(
+                    nn.Linear(config.hidden_size, 1024),
+                    nn.ReLU(),
+                    nn.Dropout(0.05),
+                    nn.Linear(1024, 16),
+                    nn.ReLU(),
+                    nn.Linear(16, output_dim),
+                )
+
+        self.rm_head.to(torch.float32)
+        self.reward_token = reward_token
+
+        self.special_token_ids = special_token_ids
+        if self.special_token_ids is not None:
+            self.reward_token = "special"
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        rope_deltas: Optional[torch.LongTensor] = None,
+    ):
+        # modified from the origin class Qwen2VLForConditionalGeneration
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        # pdb.set_trace()
+        if inputs_embeds is None:
+            inputs_embeds = self.model.language_model.embed_tokens(input_ids)
+            if pixel_values is not None:
+                pixel_values = pixel_values.type(self.model.visual.get_dtype())
+                image_embeds = self.model.visual(pixel_values, grid_thw=image_grid_thw)
+                image_mask = (
+                    (input_ids == self.config.image_token_id)
+                    .unsqueeze(-1)
+                    .expand_as(inputs_embeds)
+                )
+                image_embeds = image_embeds.to(
+                    inputs_embeds.device, inputs_embeds.dtype
+                )
+                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+            if pixel_values_videos is not None:
+                pixel_values_videos = pixel_values_videos.type(self.model.visual.get_dtype())
+                video_embeds = self.model.visual(pixel_values_videos, grid_thw=video_grid_thw)
+                video_mask = (
+                    (input_ids == self.config.video_token_id)
+                    .unsqueeze(-1)
+                    .expand_as(inputs_embeds)
+                )
+                video_embeds = video_embeds.to(
+                    inputs_embeds.device, inputs_embeds.dtype
+                )
+                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(inputs_embeds.device)
+
+        outputs = self.model.language_model(
+            input_ids=None,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]  # [B, L, D]
+        with torch.autocast(device_type='cuda', dtype=torch.float32):
+            logits = self.rm_head(hidden_states)  # [B, L, N]
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        # get sequence length
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError(
+                "Cannot handle batch sizes > 1 if no padding token is defined."
+            )
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
+                sequence_lengths = (
+                    torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+                )
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(logits.device)
+            else:
+                sequence_lengths = -1
+
+        # get the last token's logits
+        if self.reward_token == "last":
+            pooled_logits = logits[
+                torch.arange(batch_size, device=logits.device), sequence_lengths
+            ]
+        elif self.reward_token == "mean":
+            # get the mean of all valid tokens' logits
+            valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1)
+            pooled_logits = torch.stack(
+                [logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)]
+            )
+        elif self.reward_token == "special":
+            # special_token_ids = self.tokenizer.convert_tokens_to_ids(self.special_tokens)
+            # create a mask for special tokens
+            special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+            for special_token_id in self.special_token_ids:
+                special_token_mask = special_token_mask | (
+                    input_ids == special_token_id
+                )
+            pooled_logits = logits[special_token_mask, ...]
+            pooled_logits = pooled_logits.view(
+                batch_size, 1, -1
+            )  # [B, 3, N] assert 3 attributes
+            pooled_logits = pooled_logits.view(batch_size, -1)
+
+            # pdb.set_trace()
+        else:
+            raise ValueError("Invalid reward_token")
+
+        return {"logits": pooled_logits}
+
+
+_hpsv3_cache = {}
+
+
+def load_hpsv3(device, dtype, use_fsdp=True):
+    # Create cache key from arguments
+    cache_key = f"{device}_{dtype}_{use_fsdp}"
+
+    # Check if model is already cached
+    if cache_key in _hpsv3_cache:
+        return _hpsv3_cache[cache_key]
+
+    processor = AutoProcessor.from_pretrained(
+        'Qwen/Qwen2-VL-7B-Instruct', padding_side='right',
+    )
+    processor.image_processor = Qwen2VLImageProcessor()
+    special_tokens = ['<|Reward|>']
+    processor.tokenizer.add_special_tokens(
+        {'additional_special_tokens': special_tokens}
+    )
+    special_token_ids = processor.tokenizer.convert_tokens_to_ids(special_tokens)
+
+    with init_empty_weights():
+        config = Qwen2VLRewardModelBT.config_class.from_pretrained(
+            'Qwen/Qwen2-VL-7B-Instruct',
+        )
+        model = Qwen2VLRewardModelBT(
+            config,
+            output_dim=2,
+            reward_token='special',
+            special_token_ids=special_token_ids,
+            rm_head_type='ranknet',
+        )
+        model.requires_grad_(False)
+
+    model.resize_token_embeddings(len(processor.tokenizer))
+
+    model.config.tokenizer_padding_side = processor.tokenizer.padding_side
+    model.config.pad_token_id = processor.tokenizer.pad_token_id
+
+    state_dict = _load_checkpoint(
+        'huggingface://MizzenAI/HPSv3/HPSv3.safetensors', map_location='cpu'
+    )
+    new_state_dict = dict()
+    for k, v in state_dict.items():  # fix transformers version mismatch
+        if k.startswith('model.'):
+            new_k = 'model.language_model.' + k[len('model.'):]
+        elif k.startswith('visual.'):
+            new_k = 'model.visual.' + k[len('visual.'):]
+        else:
+            new_k = k
+        new_state_dict[new_k] = v
+    model.load_state_dict(new_state_dict, strict=True, assign=True)
+    model.rm_head.to(torch.float32)
+
+    if use_fsdp:
+        mmcv.print_log('Wrapping HPSv3 model with FSDP.')
+        ignored_states = []
+        for p in model.rm_head.parameters():
+            p.data = p.data.cuda()
+            ignored_states.append(p)
+        model = FullyShardedDataParallel(
+            model,
+            device_id=torch.cuda.current_device(),
+            use_orig_params=False,
+            mixed_precision=MixedPrecision(
+                param_dtype=dtype,
+                reduce_dtype=dtype,
+                buffer_dtype=dtype,
+                cast_root_forward_inputs=False),
+            sharding_strategy=ShardingStrategy.HYBRID_SHARD,
+            auto_wrap_policy=ModuleWrapPolicy([Qwen2VLVisionBlock, Qwen2VLDecoderLayer]),
+            ignored_states=ignored_states
+        )
+    else:
+        model.to(device)
+
+    result = model, processor
+    _hpsv3_cache[cache_key] = result
+    return result
+
+
+@METRICS.register_module()
+class HPSv3(Metric):
+    name = 'HPSv3'
+    requires_prompt = True
+
+    def __init__(self,
+                 num_images=None,
+                 use_fsdp=True):
+        super().__init__(num_images)
+        use_fsdp = use_fsdp and torch.cuda.is_available() and dist.is_initialized() and dist.get_world_size() > 0
+
+        self.use_fsdp = use_fsdp
+        self.dtype = torch.bfloat16
+        self.device = 'cuda' if use_fsdp else 'cpu'
+
+        self.model, self.processor = load_hpsv3(device=self.device, dtype=self.dtype, use_fsdp=use_fsdp)
+        self.model.eval()
+
+    def prepare(self):
+        self.scores = []
+
+    @torch.no_grad()
+    def feed_op(self, batch, mode):
+        imgs = batch['imgs']
+        prompts = batch['prompts']
+
+        imgs = (imgs.to(device=self.device, dtype=torch.float32) / 2 + 0.5).clamp(0, 1)
+
+        message_list = []
+        for text in prompts:
+            out_message = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image",
+                            "min_pixels": self.processor.image_processor.min_pixels,
+                            "max_pixels": self.processor.image_processor.max_pixels,
+                        },
+                        {
+                            "type": "text",
+                            "text": (
+                                INSTRUCTION.format(text_prompt=text)
+                                + prompt_with_special_token
+                            ),
+                        },
+                    ],
+                }
+            ]
+            message_list.append(out_message)
+
+        batch = self.processor(
+            text=self.processor.apply_chat_template(message_list, tokenize=False, add_generation_prompt=True),
+            images=imgs,
+            padding=True,
+            return_tensors="pt",
+            videos_kwargs={"do_rescale": True})
+        batch = {k: v.to(self.device) for k, v in batch.items()}
+        rewards = self.model(
+            return_dict=True,
+            **batch
+        )["logits"][:, 0]
+
+        if dist.is_initialized():
+            ws = dist.get_world_size()
+            placeholder = [torch.empty_like(rewards) for _ in range(ws)]
+            dist.all_gather(placeholder, rewards)
+            rewards = torch.cat(placeholder, dim=0)
+
+        if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
+            self.scores.append(rewards.float().cpu())
+
+    def feed(self, batch, mode):
+        if mode == 'reals':
+            return 0
+
+        if self.num_images is None:
+            self.feed_op(batch, mode)
+
+        else:
+            _, ws = get_dist_info()
+
+            if self.num_fake_feeded == self.num_fake_need:
+                return 0
+
+            if isinstance(batch, dict):
+                batch_size = len(list(batch.values())[0])
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = {k: v[:end] for k, v in batch.items()}
+            else:
+                batch_size = batch.shape[0]
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = batch[:end]
+
+            global_end = min(batch_size * ws,
+                             self.num_fake_need - self.num_fake_feeded)
+            self.feed_op(batch_to_feed, mode)
+            self.num_fake_feeded += global_end
+            return end
+
+    @torch.no_grad()
+    def summary(self):
+        scores = torch.cat(self.scores, dim=0)
+        if self.num_images is not None:
+            assert scores.shape[0] >= self.num_images
+            scores = scores[:self.num_images]
+        mean_score = scores.mean().item()
+        self._result_dict = dict(hpsv3=mean_score)
+        self._result_str = f'HPSv3: {mean_score:.4f}'
+        return mean_score
+
+    def clear_fake_data(self):
+        self.scores = []
+        self.num_fake_feeded = 0
+
+    def clear(self, clear_reals=False):
+        self.clear_fake_data()
+
+    def load_to_gpu(self):
+        if torch.cuda.is_available() and not isinstance(self.model, FullyShardedDataParallel):
+            self.model.cuda()
+            self.device = 'cuda'
+
+    def offload_to_cpu(self):
+        if not isinstance(self.model, FullyShardedDataParallel):
+            self.model.cpu()
+            self.device = 'cpu'

lakonlab/evaluation/metrics.py

@@ -0,0 +1,1329 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import os
+import sys
+import logging
+import pickle
+import warnings
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import mmcv
+import hashlib
+
+from copy import deepcopy
+from contextlib import contextmanager, redirect_stdout, nullcontext
+from scipy import linalg
+from scipy.stats import entropy
+from torchvision import models
+from mmcv.runner import get_dist_info, load_checkpoint
+from mmgen.utils import get_root_logger
+from mmgen.core.registry import METRICS
+from mmgen.core.evaluation.metrics import (
+    Metric, TERO_INCEPTION_URL, _load_inception_torch, MMGEN_CACHE_DIR)
+from mmgen.core.evaluation.metrics import FID as _FID
+from mmgen.core.evaluation.metrics import PR as _PR
+from open_clip import get_tokenizer, create_model
+from lakonlab.utils.io_utils import download_from_huggingface, download_from_url
+
+
+# Global caches for model loading
+_inception_cache = {}
+_hpsv2_cache = {}
+_clip_cache = {}
+
+
+def _argv_ctx(argv):
+    class _Argv:
+
+        def __enter__(self):
+            self._old = sys.argv
+            sys.argv = argv
+
+        def __exit__(self, exc_type, exc, tb):
+            sys.argv = self._old
+
+    return _Argv()
+
+
+def _redirect_stdout(to_buf):
+    return redirect_stdout(to_buf) if to_buf is not None else nullcontext()
+
+
+@contextmanager
+def _quarantine_openclip_logging():
+    """
+    Guard against open_clip (and friends) mutating global logging.
+    Snapshots root handlers/level, runs the block, then removes any
+    NEW handlers and restores the level. Also disables propagation
+    for the open_clip logger so logs don’t bubble to root.
+    """
+    root = logging.getLogger()
+    before_handlers = tuple(root.handlers)   # snapshot by identity
+    before_ids = {id(h) for h in before_handlers}
+    before_level = root.level
+
+    try:
+        yield
+    finally:
+        # Remove only handlers that were added during the block
+        for h in list(root.handlers):
+            if id(h) not in before_ids:
+                root.removeHandler(h)
+                try:
+                    h.close()
+                except Exception:
+                    pass
+        root.setLevel(before_level)
+
+        # Clamp open_clip logger so it won’t re-emit to root
+        oc = logging.getLogger("open_clip")
+        oc.propagate = False
+        oc.handlers.clear()
+
+
+def _load_inception_from_path(inception_path, map_location=None):
+    mmcv.print_log(
+        'Try to load Tero\'s Inception Model from '
+        f'\'{inception_path}\'.', 'mmgen')
+    try:
+        model = torch.jit.load(inception_path, map_location=map_location)
+        mmcv.print_log('Load Tero\'s Inception Model successfully.', 'mmgen')
+    except Exception as e:
+        model = None
+        mmcv.print_log(
+            'Load Tero\'s Inception Model failed. '
+            f'\'{e}\' occurs.', 'mmgen')
+    return model
+
+
+def _load_inception_from_url(inception_url, map_location=None):
+    """
+    Fix multi-node downloading issue in MMGen.
+    """
+    inception_url = inception_url if inception_url else TERO_INCEPTION_URL
+    mmcv.print_log(f'Try to download Inception Model from {inception_url}...',
+                   'mmgen')
+    try:
+        path = download_from_url(inception_url, dest_dir=MMGEN_CACHE_DIR)
+        mmcv.print_log('Download Finished.')
+        return _load_inception_from_path(path, map_location=map_location)
+    except Exception as e:
+        mmcv.print_log(f'Download Failed. {e} occurs.')
+        return None
+
+
+def load_inception(inception_args, metric, map_location=None):
+    """
+    Fix multi-node downloading issue in MMGen.
+    """
+    if not isinstance(inception_args, dict):
+        raise TypeError('Receive invalid \'inception_args\': '
+                        f'\'{inception_args}\'')
+
+    # Create cache key from arguments
+    cache_key = hashlib.md5(str(sorted(inception_args.items())).encode()).hexdigest()
+    cache_key += f"_{metric}"
+    
+    # Check if model is already cached
+    if cache_key in _inception_cache:
+        return _inception_cache[cache_key]
+
+    _inception_args = deepcopy(inception_args)
+    inceptoin_type = _inception_args.pop('type', None)
+
+    if torch.__version__ < '1.6.0':
+        mmcv.print_log(
+            'Current Pytorch Version not support script module, load '
+            'Inception Model from torch model zoo. If you want to use '
+            'Tero\' script model, please update your Pytorch higher '
+            f'than \'1.6\' (now is {torch.__version__})', 'mmgen')
+        result = _load_inception_torch(_inception_args, metric), 'pytorch'
+        _inception_cache[cache_key] = result
+        return result
+
+    # load pytorch version is specific
+    if inceptoin_type != 'StyleGAN':
+        result = _load_inception_torch(_inception_args, metric), 'pytorch'
+        _inception_cache[cache_key] = result
+        return result
+
+    # try to load Tero's version
+    path = _inception_args.get('inception_path', TERO_INCEPTION_URL)
+
+    # try to parse `path` as web url and download
+    if 'http' not in path:
+        model = _load_inception_from_path(path, map_location=map_location)
+        if isinstance(model, torch.nn.Module):
+            result = model, 'StyleGAN'
+            _inception_cache[cache_key] = result
+            return result
+
+    # try to parse `path` as path on disk
+    model = _load_inception_from_url(path, map_location=map_location)
+    if isinstance(model, torch.nn.Module):
+        result = model, 'StyleGAN'
+        _inception_cache[cache_key] = result
+        return result
+
+    raise RuntimeError('Cannot Load Inception Model, please check the input '
+                       f'`inception_args`: {inception_args}')
+
+
+def load_hpsv2(hps_version, device='cpu', precision='fp16'):
+    assert hps_version in ['v2', 'v2.1']
+    
+    # Create cache key from arguments
+    cache_key = f"{hps_version}_{device}_{precision}"
+    
+    # Check if model is already cached
+    if cache_key in _hpsv2_cache:
+        return _hpsv2_cache[cache_key]
+
+    with _quarantine_openclip_logging():
+        model = create_model(
+            'ViT-H-14-quickgelu',
+            precision=precision,
+            device=device,
+            output_dict=True)
+        model.requires_grad_(False)
+        tokenizer = get_tokenizer('ViT-H-14')
+    load_checkpoint(
+        model,
+        f'huggingface://xswu/HPSv2/HPS_{hps_version}_compressed.pt',
+        map_location='cpu', strict=True)
+
+    result = model, tokenizer
+    _hpsv2_cache[cache_key] = result
+    return result
+
+
+def load_openclip(
+        model_name='ViT-L-14-336-quickgelu',
+        pretrained='openai',
+        device='cpu',
+        precision='fp16'):
+    cache_key = f'{model_name}_{pretrained}_{device}_{precision}'
+    if cache_key in _clip_cache:
+        return _clip_cache[cache_key]
+
+    with _quarantine_openclip_logging():
+        model = create_model(
+            model_name,
+            pretrained=pretrained,
+            precision=precision,
+            device=device,
+            output_dict=True)
+        model.requires_grad_(False)
+        tokenizer = get_tokenizer(model_name)
+    _clip_cache[cache_key] = (model, tokenizer)
+    return _clip_cache[cache_key]
+
+
+def compute_pr_distances(row_features,
+                         col_features,
+                         col_batch_size=10000):
+    dist_batches = []
+    for col_batch in col_features.split(col_batch_size):
+        dist_batch = torch.cdist(
+            row_features.unsqueeze(0), col_batch.unsqueeze(0))[0]
+        dist_batches.append(dist_batch.cpu())
+    return torch.cat(dist_batches, dim=1)
+
+
+@METRICS.register_module(force=True)
+class PR(_PR):
+
+    def __init__(
+            self,
+            num_images=None,
+            image_shape=None,
+            feats_pkl=None,
+            k=3,
+            bgr2rgb=True,
+            vgg16_script=None,
+            inception_args=None,
+            row_batch_size=10000,
+            col_batch_size=10000):
+        super(_PR, self).__init__(num_images, image_shape)
+
+        self.feats_pkl = feats_pkl
+
+        self.vgg16 = self.inception_net = None
+        self.device = 'cpu'
+
+        if vgg16_script is not None:
+            mmcv.print_log('loading vgg16 for improved precision and recall...',
+                           'mmgen')
+            if os.path.isfile(vgg16_script):
+                self.vgg16 = torch.jit.load('work_dirs/cache/vgg16.pt', map_location=self.device).eval()
+                self.use_tero_scirpt = True
+            else:
+                mmcv.print_log(
+                    'Cannot load Tero\'s script module. Use official '
+                    'vgg16 instead', 'mmgen')
+                self.vgg16 = models.vgg16(pretrained=True).eval()
+                self.use_tero_scirpt = False
+        elif inception_args is not None:
+            self.inception_net, self.inception_style = load_inception(
+                inception_args, 'FID')
+        else:
+            raise ValueError('Please provide either vgg16_script or inception_args')
+
+        self.k = k
+        self.bgr2rgb = bgr2rgb
+        self.row_batch_size = row_batch_size
+        self.col_batch_size = col_batch_size
+
+    def prepare(self):
+        self.features_of_reals = []
+        self.features_of_fakes = []
+        if self.feats_pkl is not None:
+            assert mmcv.is_filepath(self.feats_pkl)
+            with open(self.feats_pkl, 'rb') as f:
+                reference = pickle.load(f)
+                self.features_of_reals = [torch.from_numpy(feat) for feat in reference['features_of_reals']]
+                self.num_real_feeded = reference['num_real_feeded']
+                mmcv.print_log(
+                    f'Load reference inception pkl from {self.feats_pkl}',
+                    'mmgen')
+
+    def extract_features(self, batch):
+        if self.vgg16 is not None:
+            if self.use_tero_scirpt:
+                batch = (batch * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+                feat = self.vgg16(batch, return_features=True)
+            else:
+                batch = F.interpolate(batch, size=(224, 224))
+                before_fc = self.vgg16.features(batch)
+                before_fc = before_fc.view(-1, 7 * 7 * 512)
+                feat = self.vgg16.classifier[:4](before_fc)
+        else:
+            if self.inception_style == 'StyleGAN':
+                batch = (batch * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+                feat = self.inception_net(batch, return_features=True)
+            else:
+                feat = self.inception_net(batch)[0].view(batch.shape[0], -1)
+        return feat
+
+    @torch.no_grad()
+    def feed_op(self, batch, mode):
+        batch = batch.to(self.device)
+        if self.bgr2rgb:
+            batch = batch[:, [2, 1, 0]]
+
+        feat = self.extract_features(batch)
+
+        if dist.is_initialized():
+            ws = dist.get_world_size()
+            placeholder = [torch.zeros_like(feat) for _ in range(ws)]
+            dist.all_gather(placeholder, feat)
+            feat = torch.stack(placeholder, dim=1).reshape(feat.size(0) * ws, *feat.shape[1:])
+
+        if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
+            if mode == 'reals':
+                self.features_of_reals.append(feat)
+            elif mode == 'fakes':
+                self.features_of_fakes.append(feat)
+            else:
+                raise ValueError(f'{mode} is not a implemented feed mode.')
+
+    def feed(self, batch, mode):
+        if self.num_images is not None:
+            return super().feed(batch, mode)
+        else:
+            self.feed_op(batch, mode)
+
+    @torch.no_grad()
+    def summary(self):
+        gen_features = torch.cat(self.features_of_fakes)
+        real_features = torch.cat(self.features_of_reals).to(device=gen_features.device)
+        if self.num_images is not None:
+            assert gen_features.shape[0] >= self.num_images
+            gen_features = gen_features[:self.num_images]
+            if self.feats_pkl is None:  # real feats not pre-calculated
+                assert real_features.shape[0] >= self.num_images
+                real_features = real_features[:self.num_images]
+
+        self._result_dict = {}
+
+        for name, manifold, probes in [
+            ('precision', real_features, gen_features),
+            ('recall', gen_features, real_features)
+        ]:
+            kth = []
+            for manifold_batch in manifold.split(self.row_batch_size):
+                distance = compute_pr_distances(
+                    row_features=manifold_batch,
+                    col_features=manifold,
+                    col_batch_size=self.col_batch_size)
+                kth.append(
+                    distance.to(torch.float32).kthvalue(self.k + 1).values.to(torch.float16))
+            kth = torch.cat(kth)
+            pred = []
+            for probes_batch in probes.split(self.row_batch_size):
+                distance = compute_pr_distances(
+                    row_features=probes_batch,
+                    col_features=manifold,
+                    col_batch_size=self.col_batch_size)
+                pred.append((distance <= kth).any(dim=1))
+            self._result_dict[name] = float(torch.cat(pred).to(torch.float32).mean())
+
+        precision = self._result_dict['precision']
+        recall = self._result_dict['recall']
+        self._result_str = f'precision: {precision}, recall:{recall}'
+        return self._result_dict
+
+    def clear_fake_data(self):
+        self.features_of_fakes = []
+        self.num_fake_feeded = 0
+
+    def clear(self, clear_reals=False):
+        self.clear_fake_data()
+        if clear_reals:
+            self.features_of_reals = []
+            self.num_real_feeded = 0
+
+    def load_to_gpu(self):
+        """Move models to GPU."""
+        if torch.cuda.is_available():
+            if self.vgg16 is not None:
+                self.vgg16 = self.vgg16.cuda()
+            elif self.inception_net is not None:
+                self.inception_net.cuda()
+            self.device = 'cuda'
+
+    def offload_to_cpu(self):
+        """Move models to CPU."""
+        if self.vgg16 is not None:
+            self.vgg16 = self.vgg16.cpu()
+        elif self.inception_net is not None:
+            self.inception_net.cpu()
+        self.device = 'cpu'
+
+
+@METRICS.register_module(force=True)
+class FID(_FID):
+
+    def __init__(self,
+                 num_images=None,
+                 image_shape=None,
+                 inception_pkl=None,
+                 bgr2rgb=True,
+                 inception_args=dict(normalize_input=False)):
+        super().__init__(
+            num_images,
+            image_shape=image_shape,
+            inception_pkl=inception_pkl,
+            bgr2rgb=bgr2rgb,
+            inception_args=inception_args)
+
+    def prepare(self):
+        if self.inception_pkl is not None:
+            assert mmcv.is_filepath(self.inception_pkl)
+            if self.inception_pkl.startswith('huggingface://'):
+                self.inception_pkl = download_from_huggingface(self.inception_pkl)
+            elif self.inception_pkl.startswith(('http://', 'https://')):
+                self.inception_pkl = download_from_url(self.inception_pkl)
+            with open(self.inception_pkl, 'rb') as f:
+                reference = pickle.load(f)
+                self.real_mean = reference['mean']
+                self.real_cov = reference['cov']
+                mmcv.print_log(
+                    f'Load reference inception pkl from {self.inception_pkl}',
+                    'mmgen')
+            self.num_real_feeded = self.num_images
+
+    @torch.no_grad()
+    def summary(self):
+        # calculate reference inception stat
+        if self.real_mean is None:
+            feats = torch.cat(self.real_feats, dim=0)
+            if self.num_images is not None:
+                assert feats.shape[0] >= self.num_images
+                feats = feats[:self.num_images]
+            feats_np = feats.numpy()
+            self.real_mean = np.mean(feats_np, 0)
+            self.real_cov = np.cov(feats_np, rowvar=False)
+
+        # calculate fake inception stat
+        fake_feats = torch.cat(self.fake_feats, dim=0)
+        if self.num_images is not None:
+            assert fake_feats.shape[0] >= self.num_images
+            fake_feats = fake_feats[:self.num_images]
+        fake_feats_np = fake_feats.numpy()
+        fake_mean = np.mean(fake_feats_np, 0)
+        fake_cov = np.cov(fake_feats_np, rowvar=False)
+
+        # calculate distance between real and fake statistics
+        fid, mean, cov = self._calc_fid(fake_mean, fake_cov, self.real_mean, self.real_cov)
+
+        # results for print/table
+        self._result_str = (f'{fid:.4f} ({mean:.5f}/{cov:.5f})')
+        # results for log_buffer
+        self._result_dict = dict(fid=fid, fid_mean=mean, fid_cov=cov)
+
+        return fid, mean, cov
+
+    def feed(self, batch, mode):
+        if self.num_images is not None:
+            return super().feed(batch, mode)
+        else:
+            self.feed_op(batch, mode)
+
+
+@METRICS.register_module()
+class FIDKID(FID):
+    name = 'FIDKID'
+
+    def __init__(self,
+                 num_images=None,
+                 num_subsets=100,
+                 max_subset_size=1000,
+                 **kwargs):
+        super().__init__(num_images=num_images, **kwargs)
+        self.num_subsets = num_subsets
+        self.max_subset_size = max_subset_size
+        self.real_feats_np = None
+
+    def prepare(self):
+        if self.inception_pkl is not None:
+            assert mmcv.is_filepath(self.inception_pkl)
+            with open(self.inception_pkl, 'rb') as f:
+                reference = pickle.load(f)
+                self.real_mean = reference['mean']
+                self.real_cov = reference['cov']
+                self.real_feats_np = reference['feats_np']
+                mmcv.print_log(
+                    f'Load reference inception pkl from {self.inception_pkl}',
+                    'mmgen')
+            self.num_real_feeded = self.num_images
+
+    @staticmethod
+    def _calc_kid(real_feat, fake_feat, num_subsets, max_subset_size):
+        """Refer to the implementation from:
+        https://github.com/NVlabs/stylegan2-ada-pytorch/blob/main/metrics/kernel_inception_distance.py#L18  # noqa
+        Args:
+            real_feat (np.array): Features of the real samples.
+            fake_feat (np.array): Features of the fake samples.
+            num_subsets (int): Number of subsets to calculate KID.
+            max_subset_size (int): The max size of each subset.
+        Returns:
+            float: The calculated kid metric.
+        """
+        n = real_feat.shape[1]
+        m = min(min(real_feat.shape[0], fake_feat.shape[0]), max_subset_size)
+        t = 0
+        for _ in range(num_subsets):
+            x = fake_feat[np.random.choice(
+                fake_feat.shape[0], m, replace=False)]
+            y = real_feat[np.random.choice(
+                real_feat.shape[0], m, replace=False)]
+            a = (x @ x.T / n + 1)**3 + (y @ y.T / n + 1)**3
+            b = (x @ y.T / n + 1)**3
+            t += (a.sum() - np.diag(a).sum()) / (m - 1) - b.sum() * 2 / m
+
+        kid = t / num_subsets / m
+        return float(kid)
+
+    @torch.no_grad()
+    def summary(self):
+        if self.real_feats_np is None:
+            feats = torch.cat(self.real_feats, dim=0)
+            if self.num_images is not None:
+                assert feats.shape[0] >= self.num_images
+                feats = feats[:self.num_images]
+            feats_np = feats.numpy()
+            self.real_feats_np = feats_np
+            self.real_mean = np.mean(feats_np, 0)
+            self.real_cov = np.cov(feats_np, rowvar=False)
+
+        fake_feats = torch.cat(self.fake_feats, dim=0)
+        if self.num_images is not None:
+            assert fake_feats.shape[0] >= self.num_images
+            fake_feats = fake_feats[:self.num_images]
+        fake_feats_np = fake_feats.numpy()
+        fake_mean = np.mean(fake_feats_np, 0)
+        fake_cov = np.cov(fake_feats_np, rowvar=False)
+
+        fid, mean, cov = self._calc_fid(fake_mean, fake_cov, self.real_mean,
+                                        self.real_cov)
+        kid = self._calc_kid(self.real_feats_np, fake_feats_np, self.num_subsets,
+                             self.max_subset_size) * 1000
+
+        self._result_str = f'{fid:.4f} ({mean:.5f}/{cov:.5f}), {kid:.4f}'
+        self._result_dict = dict(fid=fid, fid_mean=mean, fid_cov=cov, kid=kid)
+
+        return fid, mean, cov, kid
+
+
+@METRICS.register_module()
+class InceptionMetrics(Metric):
+    name = 'InceptionMetrics'
+
+    def __init__(self,
+                 num_images=None,
+                 reference_pkl=None,
+                 bgr2rgb=False,
+                 center_crop=False,  # SDXL-Lightning patch FID
+                 resize=True,
+                 inception_args=dict(
+                    type='StyleGAN',
+                    inception_path=TERO_INCEPTION_URL),
+                 use_kid=False,
+                 use_pr=True,
+                 use_is=True,
+                 kid_num_subsets=100,
+                 kid_max_subset_size=1000,
+                 pr_k=3,
+                 pr_row_batch_size=10000,
+                 pr_col_batch_size=10000,
+                 is_splits=10,
+                 prefix=''):
+        super().__init__(num_images)
+        self.reference_pkl = reference_pkl
+        self.real_feats = []
+        self.fake_feats = []
+        self.preds = []
+        self.real_mean = None
+        self.real_cov = None
+        self.bgr2rgb = bgr2rgb
+        self.center_crop = center_crop
+        self.resize = resize
+        self.device = 'cpu'
+
+        if self.center_crop and self.resize:
+            warnings.warn('`center_crop` is set to True, `resize` will be ignored.')
+
+        logger = get_root_logger()
+        ori_level = logger.level
+        logger.setLevel('ERROR')
+        self.inception_net, self.inception_style = load_inception(
+            inception_args, 'FID', map_location=self.device)
+        logger.setLevel(ori_level)
+
+        self.inception_net.eval()
+
+        self.use_kid = use_kid
+        self.use_pr = use_pr
+        self.use_is = use_is
+        self.kid_num_subsets = kid_num_subsets
+        self.kid_max_subset_size = kid_max_subset_size
+        self.real_feats_np = None
+
+        self.pr_k = pr_k
+        self.pr_row_batch_size = pr_row_batch_size
+        self.pr_col_batch_size = pr_col_batch_size
+
+        self.is_splits = is_splits
+        self.prefix = prefix
+
+    def prepare(self):
+        self.real_feats = []
+        self.real_feats_np = None
+        self.fake_feats = []
+        self.preds = []
+        if self.reference_pkl is not None:
+            assert mmcv.is_filepath(self.reference_pkl)
+            if self.reference_pkl.startswith('huggingface://'):
+                self.reference_pkl = download_from_huggingface(self.reference_pkl)
+            elif self.reference_pkl.startswith(('http://', 'https://')):
+                self.reference_pkl = download_from_url(self.reference_pkl)
+            with open(self.reference_pkl, 'rb') as f:
+                reference = pickle.load(f)
+                self.real_mean = reference['mean']
+                self.real_cov = reference['cov']
+                self.real_feats_np = reference['real_feats_np']
+                self.real_feats = [torch.from_numpy(reference['real_feats_np'])]
+                self.num_real_feeded = reference['num_real_feeded']
+
+    @staticmethod
+    def _calc_fid(sample_mean, sample_cov, real_mean, real_cov, eps=1e-6):
+        """Refer to the implementation from:
+
+        https://github.com/rosinality/stylegan2-pytorch/blob/master/fid.py#L34
+        """
+        cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)
+
+        if not np.isfinite(cov_sqrt).all():
+            print('product of cov matrices is singular')
+            offset = np.eye(sample_cov.shape[0]) * eps
+            cov_sqrt = linalg.sqrtm(
+                (sample_cov + offset) @ (real_cov + offset))
+
+        if np.iscomplexobj(cov_sqrt):
+            if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
+                m = np.max(np.abs(cov_sqrt.imag))
+
+                raise ValueError(f'Imaginary component {m}')
+
+            cov_sqrt = cov_sqrt.real
+
+        mean_diff = sample_mean - real_mean
+        mean_norm = mean_diff @ mean_diff
+
+        trace = np.trace(sample_cov) + np.trace(
+            real_cov) - 2 * np.trace(cov_sqrt)
+
+        fid = mean_norm + trace
+
+        return fid, mean_norm, trace
+
+    @staticmethod
+    def _calc_kid(real_feat, fake_feat, num_subsets, max_subset_size):
+        """Refer to the implementation from:
+        https://github.com/NVlabs/stylegan2-ada-pytorch/blob/main/metrics/kernel_inception_distance.py#L18  # noqa
+        Args:
+            real_feat (np.array): Features of the real samples.
+            fake_feat (np.array): Features of the fake samples.
+            num_subsets (int): Number of subsets to calculate KID.
+            max_subset_size (int): The max size of each subset.
+        Returns:
+            float: The calculated kid metric.
+        """
+        n = real_feat.shape[1]
+        m = min(min(real_feat.shape[0], fake_feat.shape[0]), max_subset_size)
+        t = 0
+        for _ in range(num_subsets):
+            x = fake_feat[np.random.choice(
+                fake_feat.shape[0], m, replace=False)]
+            y = real_feat[np.random.choice(
+                real_feat.shape[0], m, replace=False)]
+            a = (x @ x.T / n + 1)**3 + (y @ y.T / n + 1)**3
+            b = (x @ y.T / n + 1)**3
+            t += (a.sum() - np.diag(a).sum()) / (m - 1) - b.sum() * 2 / m
+
+        kid = t / num_subsets / m
+        return float(kid)
+
+    def extract_features(self, batch):
+        if self.center_crop:
+            crop_size = 299
+            h, w = batch.shape[2], batch.shape[3]
+            assert h >= crop_size and w >= crop_size
+            h_offset = (h - crop_size) // 2
+            w_offset = (w - crop_size) // 2
+            batch = batch[:, :, h_offset:h_offset + crop_size, w_offset:w_offset + crop_size]
+        elif self.resize:
+            batch = F.interpolate(
+                batch, size=(299, 299), mode='bicubic', align_corners=False, antialias=True).clamp(min=-1, max=1)
+        assert self.inception_style == 'StyleGAN'
+        batch = (batch * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        feat = self.inception_net(batch, return_features=True)
+        pred = F.linear(feat, self.inception_net.output.weight).softmax(dim=1)
+        return feat, pred
+
+    @torch.no_grad()
+    def feed_op(self, batch, mode):
+        if self.bgr2rgb:
+            batch = batch[:, [2, 1, 0]]
+        batch = batch.to(self.device)
+
+        feat, pred = self.extract_features(batch)
+
+        if dist.is_initialized():
+            ws = dist.get_world_size()
+            placeholder = [torch.zeros_like(feat) for _ in range(ws)]
+            dist.all_gather(placeholder, feat)
+            feat = torch.stack(placeholder, dim=1).reshape(feat.size(0) * ws, *feat.shape[1:])
+            if mode == 'fakes':
+                placeholder = [torch.zeros_like(pred) for _ in range(ws)]
+                dist.all_gather(placeholder, pred)
+                pred = torch.stack(placeholder, dim=1).reshape(pred.size(0) * ws, *pred.shape[1:])
+
+        # in distributed training, we only collect features at rank-0.
+        if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
+            if mode == 'reals':
+                self.real_feats.append(feat.cpu())
+            elif mode == 'fakes':
+                self.fake_feats.append(feat.cpu())
+                self.preds.append(pred.cpu().numpy())
+            else:
+                raise ValueError(
+                    f"The expected mode should be set to 'reals' or 'fakes,\
+                    but got '{mode}'")
+
+    def feed(self, batch, mode):
+        if self.num_images is None:
+            self.feed_op(batch, mode)
+
+        else:
+            _, ws = get_dist_info()
+            if mode == 'reals':
+                if self.num_real_feeded == self.num_real_need:
+                    return 0
+
+                if isinstance(batch, dict):
+                    batch_size = len(list(batch.values())[0])
+                    end = min(batch_size, self.num_real_need - self.num_real_feeded)
+                    batch_to_feed = {k: v[:end] for k, v in batch.items()}
+                else:
+                    batch_size = batch.shape[0]
+                    end = min(batch_size, self.num_real_need - self.num_real_feeded)
+                    batch_to_feed = batch[:end]
+
+                global_end = min(batch_size * ws,
+                                 self.num_real_need - self.num_real_feeded)
+                self.feed_op(batch_to_feed, mode)
+                self.num_real_feeded += global_end
+                return end
+
+            elif mode == 'fakes':
+                if self.num_fake_feeded == self.num_fake_need:
+                    return 0
+
+                if isinstance(batch, dict):
+                    batch_size = len(list(batch.values())[0])
+                    end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                    batch_to_feed = {k: v[:end] for k, v in batch.items()}
+                else:
+                    batch_size = batch.shape[0]
+                    end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                    batch_to_feed = batch[:end]
+
+                global_end = min(batch_size * ws,
+                                 self.num_fake_need - self.num_fake_feeded)
+                self.feed_op(batch_to_feed, mode)
+                self.num_fake_feeded += global_end
+                return end
+            else:
+                raise ValueError(
+                    'The expected mode should be set to \'reals\' or \'fakes\','
+                    f'but got \'{mode}\'')
+
+    @torch.no_grad()
+    def summary(self):
+        real_feats = torch.cat(self.real_feats, dim=0)
+        fake_feats = torch.cat(self.fake_feats, dim=0)
+        if self.num_images is not None:
+            assert fake_feats.shape[0] >= self.num_images
+            fake_feats = fake_feats[:self.num_images]
+            if self.reference_pkl is None:  # real feats not pre-calculated
+                assert real_feats.shape[0] >= self.num_images
+                real_feats = real_feats[:self.num_images]
+
+        if self.real_feats_np is None:
+            real_feats_np = real_feats.numpy()
+            self.real_feats_np = real_feats_np
+            self.real_mean = np.mean(real_feats_np, 0)
+            self.real_cov = np.cov(real_feats_np, rowvar=False)
+
+        self._result_dict = dict()
+
+        prefix = self.prefix + '_' if len(self.prefix) > 0 else ''
+
+        # FID
+        fake_feats_np = fake_feats.numpy()
+        fake_mean = np.mean(fake_feats_np, 0)
+        fake_cov = np.cov(fake_feats_np, rowvar=False)
+        fid, mean, cov = self._calc_fid(fake_mean, fake_cov, self.real_mean,
+                                        self.real_cov)
+        self._result_dict.update({f'{prefix}fid': fid})
+        _result_str = f'{prefix}FID: {fid:.4f} ({mean:.4f}/{cov:.4f})'
+
+        # KID
+        if self.use_kid:
+            kid = self._calc_kid(self.real_feats_np, fake_feats_np, self.kid_num_subsets,
+                                 self.kid_max_subset_size) * 1000
+            self._result_dict.update({f'{prefix}kid': kid})
+            _result_str += f', {prefix}KID: {kid:.4f}'
+        else:
+            kid = None
+
+        # PR
+        if self.use_pr:
+            for name, manifold, probes in [
+                (f'{prefix}precision', real_feats, fake_feats),
+                (f'{prefix}recall', fake_feats, real_feats)
+            ]:
+                kth = []
+                for manifold_batch in manifold.split(self.pr_row_batch_size):
+                    distance = compute_pr_distances(
+                        row_features=manifold_batch,
+                        col_features=manifold,
+                        col_batch_size=self.pr_col_batch_size)
+                    kth.append(
+                        distance.to(torch.float32).kthvalue(self.pr_k + 1).values.to(torch.float16))
+                kth = torch.cat(kth)
+                pred = []
+                for probes_batch in probes.split(self.pr_row_batch_size):
+                    distance = compute_pr_distances(
+                        row_features=probes_batch,
+                        col_features=manifold,
+                        col_batch_size=self.pr_col_batch_size)
+                    pred.append((distance <= kth).any(dim=1))
+                self._result_dict[name] = float(torch.cat(pred).to(torch.float32).mean())
+            precision = self._result_dict[f'{prefix}precision']
+            recall = self._result_dict[f'{prefix}recall']
+            _result_str += f', {prefix}Precision: {precision:.5f}, {prefix}Recall:{recall:.5f}'
+        else:
+            precision = recall = None
+
+        # IS
+        if self.use_is:
+            split_scores = []
+            self.preds = np.concatenate(self.preds, axis=0)
+            if self.num_images is not None:
+                assert self.preds.shape[0] >= self.num_images
+                self.preds = self.preds[:self.num_images]
+            num_preds = self.preds.shape[0]
+            for k in range(self.is_splits):
+                part = self.preds[k * (num_preds // self.is_splits):(k + 1) * (num_preds // self.is_splits), :]
+                py = np.mean(part, axis=0)
+                scores = []
+                for i in range(part.shape[0]):
+                    pyx = part[i, :]
+                    scores.append(entropy(pyx, py))
+                split_scores.append(np.exp(np.mean(scores)))
+            is_mean = np.mean(split_scores)
+            self._result_dict.update({f'{prefix}is': is_mean})
+            _result_str += f', {prefix}IS: {is_mean:.2f}'
+        else:
+            is_mean = None
+
+        self._result_str = _result_str
+
+        return fid, kid, precision, recall, is_mean
+
+    def clear_fake_data(self):
+        self.fake_feats = []
+        self.preds = []
+        self.num_fake_feeded = 0
+
+    def clear(self, clear_reals=False):
+        self.clear_fake_data()
+        if clear_reals:
+            self.real_feats = []
+            self.real_feats_np = None
+            self.num_real_feeded = 0
+
+    def load_to_gpu(self):
+        """Move models to GPU."""
+        if torch.cuda.is_available():
+            self.inception_net.cuda()
+            self.device = 'cuda'
+
+    def offload_to_cpu(self):
+        """Move models to CPU."""
+        self.inception_net.cpu()
+        self.device = 'cpu'
+
+
+@METRICS.register_module()
+class ColorStats(Metric):
+    name = 'ColorStats'
+
+    def __init__(self,
+                 num_images=None):
+        super().__init__(num_images)
+
+    def prepare(self):
+        self.stats = []
+
+    @staticmethod
+    def srgb_to_linear(c):
+        threshold = 0.04045
+        below = c <= threshold
+        out = torch.where(
+            below, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
+        return out
+
+    @staticmethod
+    def linear_to_srgb(c):
+        threshold = 0.0031308
+        below = c <= threshold
+        out = torch.where(
+            below, 12.92 * c, 1.055 * c ** (1.0 / 2.4) - 0.055)
+        return out
+
+    def rgb_to_grayscale_srgb(self, img_srgb):
+        img_lin = self.srgb_to_linear(img_srgb)
+        R_lin, G_lin, B_lin = img_lin.unbind(dim=1)
+        Y_lin = 0.2126 * R_lin + 0.7152 * G_lin + 0.0722 * B_lin
+        gray_srgb = self.linear_to_srgb(Y_lin)
+        return gray_srgb
+
+    @staticmethod
+    def srgb_to_hsv_saturation(img_srgb):
+        c_max = torch.amax(img_srgb, dim=1)
+        c_min = torch.amin(img_srgb, dim=1)
+        delta = c_max - c_min
+        sat = delta / c_max.clamp(min=1e-5)
+        return sat
+
+    def compute_stats(self, batch):
+        batch = (batch / 2 + 0.5).clamp(0, 1)
+        gray = self.rgb_to_grayscale_srgb(batch).flatten(1)
+        contrast, brightness = torch.std_mean(gray, dim=1)
+        saturation = self.srgb_to_hsv_saturation(batch).flatten(1).mean(dim=1)
+        return torch.stack([brightness, contrast, saturation], dim=-1)
+
+    @torch.no_grad()
+    def feed_op(self, batch, mode):
+        stats = self.compute_stats(batch)
+
+        if dist.is_initialized():
+            ws = dist.get_world_size()
+            placeholder = [torch.zeros_like(stats) for _ in range(ws)]
+            dist.all_gather(placeholder, stats)
+            stats = torch.stack(placeholder, dim=1).reshape(stats.size(0) * ws, *stats.shape[1:])
+
+        # in distributed training, we only collect features at rank-0.
+        if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
+            self.stats.append(stats.cpu())
+
+    def feed(self, batch, mode):
+        if mode == 'reals':
+            return 0
+
+        if self.num_images is None:
+            self.feed_op(batch, mode)
+
+        else:
+            _, ws = get_dist_info()
+
+            if self.num_fake_feeded == self.num_fake_need:
+                return 0
+
+            if isinstance(batch, dict):
+                batch_size = len(list(batch.values())[0])
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = {k: v[:end] for k, v in batch.items()}
+            else:
+                batch_size = batch.shape[0]
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = batch[:end]
+
+            global_end = min(batch_size * ws,
+                             self.num_fake_need - self.num_fake_feeded)
+            self.feed_op(batch_to_feed, mode)
+            self.num_fake_feeded += global_end
+            return end
+
+    @torch.no_grad()
+    def summary(self):
+        stats = torch.cat(self.stats, dim=0)
+        if self.num_images is not None:
+            assert stats.shape[0] >= self.num_images
+            stats = stats[:self.num_images]
+        stats = stats.mean(dim=0)
+        brightness, contrast, saturation = stats.tolist()
+        self._result_dict = dict(
+            brightness=brightness, contrast=contrast, saturation=saturation)
+        self._result_str = f'Brightness: {brightness:.4f}, Contrast: {contrast:.4f}, Saturation: {saturation:.4f}'
+        return brightness, contrast, saturation
+
+    def clear_fake_data(self):
+        self.stats = []
+        self.num_fake_feeded = 0
+
+    def clear(self, clear_reals=False):
+        self.clear_fake_data()
+
+
+@METRICS.register_module()
+class HPSv2(Metric):
+    name = 'HPSv2'
+    requires_prompt = True
+
+    def __init__(self,
+                 num_images=None,
+                 hps_version='v2.1'):
+        super().__init__(num_images)
+        self.hps_version = hps_version
+        self.device = 'cpu'  # Initialize on CPU
+        self.dtype = torch.float16
+        self.model, self.tokenizer = load_hpsv2(hps_version, device=self.device, precision='fp16')
+        self.model.eval()
+        image_size = self.model.visual.image_size
+        if isinstance(image_size, tuple):
+            assert len(image_size) == 2 and image_size[0] == image_size[1]
+            image_size = image_size[0]
+        self.image_size = image_size
+        self.image_mean = torch.tensor(self.model.visual.image_mean, device=self.device).view(3, 1, 1)
+        self.image_std = torch.tensor(self.model.visual.image_std, device=self.device).view(3, 1, 1)
+
+    def prepare(self):
+        self.scores = []
+
+    def resize(self, imgs):
+        h, w = imgs.shape[2:]
+        scale = self.image_size / float(max(h, w))
+        if scale != 1.0:
+            h = int(round(h * scale))
+            w = int(round(w * scale))
+            imgs = F.interpolate(imgs, size=(h, w), mode='bicubic', align_corners=False, antialias=True).clamp(0, 1)
+        if h != w:
+            pad_h = self.image_size - h
+            pad_w = self.image_size - w
+            imgs = F.pad(
+                imgs, (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2), mode='constant', value=0)
+        return imgs
+
+    @torch.no_grad()
+    def feed_op(self, batch, mode):
+        imgs = batch['imgs']
+        prompts = batch['prompts']
+
+        imgs = (imgs.to(device=self.device, dtype=torch.float32) / 2 + 0.5).clamp(0, 1)
+        imgs = ((self.resize(imgs) - self.image_mean) / self.image_std).to(dtype=self.dtype)
+        prompts = self.tokenizer(prompts).to(device=self.device)
+
+        outputs = self.model(imgs, prompts)
+        image_features, text_features = outputs['image_features'], outputs['text_features']
+        hps_scores = (image_features * text_features).sum(dim=-1)  # (bs, )
+
+        if dist.is_initialized():
+            ws = dist.get_world_size()
+            placeholder = [torch.empty_like(hps_scores) for _ in range(ws)]
+            dist.all_gather(placeholder, hps_scores)
+            hps_scores = torch.stack(placeholder, dim=1).reshape(hps_scores.size(0) * ws)
+
+        if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
+            self.scores.append(hps_scores.float().cpu())
+
+    def feed(self, batch, mode):
+        if mode == 'reals':
+            return 0
+
+        if self.num_images is None:
+            self.feed_op(batch, mode)
+
+        else:
+            _, ws = get_dist_info()
+
+            if self.num_fake_feeded == self.num_fake_need:
+                return 0
+
+            if isinstance(batch, dict):
+                batch_size = len(list(batch.values())[0])
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = {k: v[:end] for k, v in batch.items()}
+            else:
+                batch_size = batch.shape[0]
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = batch[:end]
+
+            global_end = min(batch_size * ws,
+                             self.num_fake_need - self.num_fake_feeded)
+            self.feed_op(batch_to_feed, mode)
+            self.num_fake_feeded += global_end
+            return end
+
+    @torch.no_grad()
+    def summary(self):
+        scores = torch.cat(self.scores, dim=0)
+        if self.num_images is not None:
+            assert scores.shape[0] >= self.num_images
+            scores = scores[:self.num_images]
+        mean_score = scores.mean().item()
+        self._result_dict = dict(hpsv2=mean_score)
+        self._result_str = f'HPSv2: {mean_score:.4f}'
+        return mean_score
+
+    def clear_fake_data(self):
+        self.scores = []
+        self.num_fake_feeded = 0
+
+    def clear(self, clear_reals=False):
+        self.clear_fake_data()
+
+    def load_to_gpu(self):
+        if torch.cuda.is_available():
+            self.model.cuda()
+            self.image_mean = self.image_mean.cuda()
+            self.image_std = self.image_std.cuda()
+            self.device = 'cuda'
+
+    def offload_to_cpu(self):
+        self.model.cpu()
+        self.image_mean = self.image_mean.cpu()
+        self.image_std = self.image_std.cpu()
+        self.device = 'cpu'
+
+
+@METRICS.register_module()
+class CLIPSimilarity(Metric):
+    """
+    Average image–text CLIP cosine similarity (↑ better).
+    Preprocess emulates OpenAI CLIP for ViT-L/14@336:
+      - Resize so min(H, W) = 336 (bicubic, antialias), keep aspect ratio
+      - Center crop to 336x336
+      - Normalize with model.visual.image_mean/std
+    Expects batch = {'imgs': (B,3,H,W) in [-1,1], 'prompts': List[str]}
+    """
+    name = 'CLIPSimilarity'
+    requires_prompt = True
+
+    def __init__(
+        self,
+        num_images=None,
+        model_name='ViT-L-14-336-quickgelu',
+        pretrained='openai',
+        precision='fp16',   # 'fp16' | 'fp32' | 'bf16'
+    ):
+        super().__init__(num_images)
+        self.model_name = model_name
+        self.pretrained = pretrained
+        self.precision = precision
+
+        self.device = 'cpu'
+        self.dtype = {
+            'fp16': torch.float16,
+            'bf16': torch.bfloat16,
+            'fp32': torch.float32
+        }.get(precision, torch.float16)
+
+        self.model, self.tokenizer = load_openclip(
+            model_name=model_name,
+            pretrained=pretrained,
+            device=self.device,
+            precision=precision,
+        )
+        self.model.eval()
+
+        # OpenAI ViT-L/14@336 uses square 336 input
+        image_size = self.model.visual.image_size
+        if isinstance(image_size, tuple):
+            assert len(image_size) == 2 and image_size[0] == image_size[1]
+            image_size = image_size[0]
+        self.image_size = int(image_size)  # 336
+
+        # Use the model's own stats for normalization
+        self.image_mean = torch.tensor(self.model.visual.image_mean, device=self.device).view(3, 1, 1)
+        self.image_std = torch.tensor(self.model.visual.image_std, device=self.device).view(3, 1, 1)
+
+    def prepare(self):
+        self.scores = []
+
+    def _resize_min_side_then_center_crop(self, imgs):
+        """
+        imgs: (B,3,H,W) in [0,1], float32, on self.device
+        1) Resize so min(H,W) == self.image_size, preserve AR (bicubic, antialias)
+        2) Center-crop to (self.image_size, self.image_size)
+        3) Normalize with model mean/std
+        4) Cast to self.dtype
+        """
+        _, _, H, W = imgs.shape
+        target = self.image_size
+
+        # Scale factor so that the shorter side becomes 'target'
+        short, long = (H, W) if H < W else (W, H)
+        if short == 0:
+            raise ValueError("Invalid image with zero dimension.")
+        scale = target / float(short)
+
+        new_h = max(1, int(round(H * scale)))
+        new_w = max(1, int(round(W * scale)))
+        if new_h != H or new_w != W:
+            imgs = F.interpolate(
+                imgs, size=(new_h, new_w),
+                mode='bicubic', align_corners=False, antialias=True
+            ).clamp(0, 1)
+
+        # Center crop to target x target
+        top = max(0, (new_h - target) // 2)
+        left = max(0, (new_w - target) // 2)
+        imgs = imgs[:, :, top:top + target, left:left + target]
+
+        imgs = (imgs - self.image_mean) / self.image_std
+        return imgs.to(dtype=self.dtype)
+
+    @torch.no_grad()
+    def feed_op(self, batch, mode):
+        if mode == 'reals':
+            return 0
+
+        imgs = batch['imgs']
+        prompts = batch['prompts']
+
+        # [-1,1] -> [0,1]
+        imgs = (imgs.to(device=self.device, dtype=torch.float32) / 2 + 0.5).clamp(0, 1)
+        imgs = self._resize_min_side_then_center_crop(imgs)
+
+        # Tokenize on device
+        text = self.tokenizer(prompts).to(device=self.device)
+
+        # Forward (create_model(..., output_dict=True)) => dict w/ features
+        out = self.model(imgs, text)
+        if isinstance(out, dict) and ('image_features' in out and 'text_features' in out):
+            img_feat = out['image_features']
+            txt_feat = out['text_features']
+        else:
+            img_feat = self.model.encode_image(imgs)
+            txt_feat = self.model.encode_text(text)
+
+        # Cosine similarity per pair
+        img_feat = F.normalize(img_feat, dim=-1)
+        txt_feat = F.normalize(txt_feat, dim=-1)
+        sim = (img_feat * txt_feat).sum(dim=-1).to(torch.float32)  # (B,)
+
+        # DDP gather
+        if dist.is_initialized():
+            ws = dist.get_world_size()
+            bucket = [torch.empty_like(sim) for _ in range(ws)]
+            dist.all_gather(bucket, sim)
+            sim = torch.stack(bucket, dim=1).reshape(sim.size(0) * ws)
+
+        if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
+            self.scores.append(sim.cpu())
+
+    def feed(self, batch, mode):
+        if mode == 'reals':
+            return 0
+
+        if self.num_images is None:
+            self.feed_op(batch, mode)
+
+        else:
+            _, ws = get_dist_info()
+
+            if self.num_fake_feeded == self.num_fake_need:
+                return 0
+
+            if isinstance(batch, dict):
+                batch_size = len(list(batch.values())[0])
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = {k: v[:end] for k, v in batch.items()}
+            else:
+                batch_size = batch.shape[0]
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = batch[:end]
+
+            global_end = min(batch_size * ws, self.num_fake_need - self.num_fake_feeded)
+            self.feed_op(batch_to_feed, mode)
+            self.num_fake_feeded += global_end
+            return end
+
+    @torch.no_grad()
+    def summary(self):
+        sims = torch.cat(self.scores, dim=0)
+        if self.num_images is not None:
+            assert sims.shape[0] >= self.num_images
+            sims = sims[:self.num_images]
+        mean_sim = sims.mean().item()
+
+        self._result_dict = dict(clipsim=mean_sim)  # raw cosine in [-1,1]
+        self._result_str = f'CLIPSim: {mean_sim:.4f}'
+        return mean_sim
+
+    def clear_fake_data(self):
+        self.scores = []
+        self.num_fake_feeded = 0
+
+    def clear(self, clear_reals=False):
+        self.clear_fake_data()
+
+    def load_to_gpu(self):
+        if torch.cuda.is_available():
+            self.model.cuda()
+            self.image_mean = self.image_mean.cuda()
+            self.image_std = self.image_std.cuda()
+            self.device = 'cuda'
+
+    def offload_to_cpu(self):
+        self.model.cpu()
+        self.image_mean = self.image_mean.cpu()
+        self.image_std = self.image_std.cpu()
+        self.device = 'cpu'

lakonlab/evaluation/vqa_score.py

@@ -0,0 +1,667 @@
+# Modified from https://github.com/linzhiqiu/t2v_metrics
+# Copyright 2023 Zhiqiu Lin
+
+import os
+import re
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+import mmcv
+
+from typing import List, Optional, Tuple, Union
+from dataclasses import dataclass, field
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, FullyShardedDataParallel
+from torch.distributed.fsdp.wrap import ModuleWrapPolicy
+from transformers import (
+    AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, T5Config, T5ForConditionalGeneration,
+    CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig)
+from transformers.models.t5.modeling_t5 import T5Block
+from transformers.modeling_outputs import Seq2SeqLMOutput
+from mmcv.runner import get_dist_info
+from mmgen.core.registry import METRICS
+from mmgen.core.evaluation.metrics import Metric
+
+IMAGE_TOKEN_INDEX = -200
+CONTEXT_LEN = 2048
+SYSTEM_MSG = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions."
+IGNORE_INDEX = -100
+DEFAULT_IMAGE_TOKEN = "<image>"
+
+default_question_template = 'Does this figure show "{}"? Please answer yes or no.'
+default_answer_template = "Yes"
+
+
+def t5_tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
+
+    def insert_separator(X, sep):
+        return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1]
+
+    input_ids = []
+    # Since there's no bos_token_id, simply concatenate the tokenized prompt_chunks with the image_token_index
+    for x in insert_separator(prompt_chunks, [image_token_index]):
+        input_ids.extend(x)
+
+    if return_tensors is not None:
+        if return_tensors == 'pt':
+            return torch.tensor(input_ids, dtype=torch.long)
+        raise ValueError(f'Unsupported tensor type: {return_tensors}')
+    return input_ids
+
+
+def format_question(question, conversation_style='plain'):
+    if conversation_style == 't5_plain':  # for 1st stage t5 model
+        question = DEFAULT_IMAGE_TOKEN + question
+    elif conversation_style == 't5_chat':  # for 2nd stage t5 model
+        question = SYSTEM_MSG + " USER: " + DEFAULT_IMAGE_TOKEN + "\n" + question + " ASSISTANT: "
+    elif conversation_style == 't5_chat_no_system':  # for 2nd stage t5 model
+        question = "USER: " + DEFAULT_IMAGE_TOKEN + "\n" + question + " ASSISTANT: "
+    elif conversation_style == 't5_chat_no_system_no_user':  # for 2nd stage t5 model
+        question = "" + DEFAULT_IMAGE_TOKEN + "\n" + question + " : "
+    # elif conversation_style == 't5_chat_ood_system': # for 2nd stage t5 model
+    #     question = SYSTEM_MSG + " HUMAN: " + DEFAULT_IMAGE_TOKEN + "\n" + question + " GPT: "
+    else:
+        raise NotImplementedError()
+    return question
+
+
+def format_answer(answer, conversation_style='plain'):
+    return answer
+
+
+class CLIPVisionTower(nn.Module):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__()
+
+        self.is_loaded = False
+
+        self.vision_tower_name = vision_tower
+        self.select_layer = args.mm_vision_select_layer
+        self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')
+
+        if not delay_load:
+            self.load_model()
+        else:
+            self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
+
+    def load_model(self):
+        self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
+        self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name)
+        self.vision_tower.requires_grad_(False)
+
+        self.is_loaded = True
+
+    def feature_select(self, image_forward_outs):
+        image_features = image_forward_outs.hidden_states[self.select_layer]
+        if self.select_feature == 'patch':
+            image_features = image_features[:, 1:]
+        elif self.select_feature == 'cls_patch':
+            image_features = image_features
+        else:
+            raise ValueError(f'Unexpected select feature: {self.select_feature}')
+        return image_features
+
+    @torch.no_grad()
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0),
+                                                      output_hidden_states=True)
+                image_feature = self.feature_select(image_forward_out).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype),
+                                                   output_hidden_states=True)
+            image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device
+
+    @property
+    def config(self):
+        if self.is_loaded:
+            return self.vision_tower.config
+        else:
+            return self.cfg_only
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2
+
+
+class IdentityMap(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    @property
+    def config(self):
+        return {"mm_projector_type": 'identity'}
+
+
+def build_vision_tower(vision_tower_cfg, **kwargs):
+    vision_tower = getattr(vision_tower_cfg, 'mm_vision_tower', getattr(vision_tower_cfg, 'vision_tower', None))
+    is_absolute_path_exists = os.path.exists(vision_tower)
+    if is_absolute_path_exists or vision_tower.startswith("openai") or vision_tower.startswith("laion"):
+        return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)
+
+    raise ValueError(f'Unknown vision tower: {vision_tower}')
+
+
+def build_vision_projector(config, delay_load=False, **kwargs):
+    projector_type = getattr(config, 'mm_projector_type', 'linear')
+
+    if projector_type == 'linear':
+        return nn.Linear(config.mm_hidden_size, config.hidden_size)
+
+    mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
+    if mlp_gelu_match:
+        mlp_depth = int(mlp_gelu_match.group(1))
+        modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+        for _ in range(1, mlp_depth):
+            modules.append(nn.GELU())
+            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+        return nn.Sequential(*modules)
+
+    if projector_type == 'identity':
+        return IdentityMap()
+
+    raise ValueError(f'Unknown projector type: {projector_type}')
+
+
+@dataclass
+class ModelArguments:
+    tune_mm_mlp_adapter: bool = field(default=False)
+    vision_tower: Optional[str] = field(default='openai/clip-vit-large-patch14-336')
+    mm_vision_select_layer: Optional[int] = field(default=-2)  # default to the second last layer in llava1.5
+    pretrain_mm_mlp_adapter: Optional[str] = field(default=None)
+    mm_projector_type: Optional[str] = field(default='mlp2x_gelu')
+    mm_vision_select_feature: Optional[str] = field(default="patch")
+
+
+class CLIPT5Config(T5Config):
+    model_type = "clip_t5"
+
+
+class CLIPT5ForConditionalGeneration(T5ForConditionalGeneration):
+    # This class supports both T5 and FlanT5
+    config_class = CLIPT5Config
+
+    def __init__(self, config):
+        super(CLIPT5ForConditionalGeneration, self).__init__(config)
+        self.embed_tokens = self.encoder.embed_tokens
+        if hasattr(config, "mm_vision_tower"):
+            self.vision_tower = build_vision_tower(config, delay_load=False)
+            self.mm_projector = build_vision_projector(config)
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, 'vision_tower', None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def get_model(self):
+        return self  # for compatibility with LlavaMetaForCausalLM
+
+    def prepare_inputs_labels_for_multimodal(
+            self, input_ids, attention_mask, decoder_attention_mask, past_key_values, labels, images
+    ):
+        # The labels are now separated from the input_ids.
+        vision_tower = self.get_vision_tower()
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            raise NotImplementedError()
+
+        if type(images) is list or images.ndim == 5:
+            concat_images = torch.cat([image for image in images], dim=0)
+            image_features = self.encode_images(concat_images)
+            split_sizes = [image.shape[0] for image in images]
+            image_features = torch.split(image_features, split_sizes, dim=0)
+            image_features = [x.flatten(0, 1) for x in image_features]
+        else:
+            image_features = self.encode_images(images)
+
+        new_input_embeds = []
+        cur_image_idx = 0
+        for _, cur_input_ids in enumerate(input_ids):
+            if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0:
+                # multimodal LLM, but the current sample is not multimodal
+                raise NotImplementedError()
+            image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
+            cur_new_input_embeds = []
+            while image_token_indices.numel() > 0:
+                cur_image_features = image_features[cur_image_idx]
+                image_token_start = image_token_indices[0]
+                cur_new_input_embeds.append(self.embed_tokens(cur_input_ids[:image_token_start]))
+                cur_new_input_embeds.append(cur_image_features)
+                cur_image_idx += 1
+                cur_input_ids = cur_input_ids[image_token_start + 1:]
+                image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
+            if cur_input_ids.numel() > 0:
+                cur_new_input_embeds.append(self.embed_tokens(cur_input_ids))
+            cur_new_input_embeds = [x.to(device=self.device) for x in cur_new_input_embeds]
+            cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)
+            new_input_embeds.append(cur_new_input_embeds)
+
+        if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds):
+            max_len = max(x.shape[0] for x in new_input_embeds)
+
+            new_input_embeds_align = []
+            _input_embeds_lengths = []
+            for cur_new_embed in new_input_embeds:
+                _input_embeds_lengths.append(cur_new_embed.shape[0])
+                cur_new_embed = torch.cat((cur_new_embed,
+                                           torch.zeros((max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]),
+                                                       dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0)
+                new_input_embeds_align.append(cur_new_embed)
+            new_input_embeds = torch.stack(new_input_embeds_align, dim=0)
+
+            if attention_mask is not None:
+                new_attention_mask = []
+                for cur_attention_mask, _input_embeds_length in zip(attention_mask, _input_embeds_lengths):
+                    new_attn_mask_pad_left = torch.full((_input_embeds_length - input_ids.shape[1],), True,
+                                                        dtype=attention_mask.dtype, device=attention_mask.device)
+                    new_attn_mask_pad_right = torch.full((new_input_embeds.shape[1] - _input_embeds_length,), False,
+                                                         dtype=attention_mask.dtype, device=attention_mask.device)
+                    cur_new_attention_mask = torch.cat(
+                        (new_attn_mask_pad_left, cur_attention_mask, new_attn_mask_pad_right), dim=0)
+                    new_attention_mask.append(cur_new_attention_mask)
+                attention_mask = torch.stack(new_attention_mask, dim=0)
+                assert attention_mask.shape == new_input_embeds.shape[:2]
+        else:
+            new_input_embeds = torch.stack(new_input_embeds, dim=0)
+
+            if attention_mask is not None:
+                new_attn_mask_pad_left = torch.full(
+                    (attention_mask.shape[0], new_input_embeds.shape[1] - input_ids.shape[1]), True,
+                    dtype=attention_mask.dtype, device=attention_mask.device)
+                attention_mask = torch.cat((new_attn_mask_pad_left, attention_mask), dim=1)
+                assert attention_mask.shape == new_input_embeds.shape[:2]
+
+        return None, attention_mask, decoder_attention_mask, past_key_values, new_input_embeds, labels
+
+    def encode_images(self, images):
+        image_features = self.get_vision_tower()(images)
+        image_features = self.mm_projector(image_features)
+        return image_features
+
+    def initialize_vision_modules(self, model_args, fsdp=None):
+        vision_tower = model_args.vision_tower
+        mm_vision_select_layer = model_args.mm_vision_select_layer
+        mm_vision_select_feature = model_args.mm_vision_select_feature
+        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
+
+        self.config.mm_vision_tower = vision_tower
+        self.config.pretrain_mm_mlp_adapter = pretrain_mm_mlp_adapter
+
+        if self.get_vision_tower() is None:
+            vision_tower = build_vision_tower(model_args)
+
+            if fsdp is not None and len(fsdp) > 0:
+                self.vision_tower = [vision_tower]
+            else:
+                self.vision_tower = vision_tower
+        else:
+            if fsdp is not None and len(fsdp) > 0:
+                vision_tower = self.vision_tower[0]
+            else:
+                vision_tower = self.vision_tower
+            if not vision_tower.is_loaded:
+                vision_tower.load_model()
+
+        self.config.use_mm_proj = True
+        self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'mlp2x_gelu')
+        self.config.mm_hidden_size = vision_tower.hidden_size
+        self.config.mm_vision_select_layer = mm_vision_select_layer
+        self.config.mm_vision_select_feature = mm_vision_select_feature
+
+        if getattr(self, 'mm_projector', None) is None:
+            self.mm_projector = build_vision_projector(self.config)
+
+        if pretrain_mm_mlp_adapter is not None:
+            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu')
+
+            def get_w(weights, keyword):
+                return {k.split(keyword + '.')[1]: v for k, v in weights.items() if keyword in k}
+
+            self.mm_projector.load_state_dict(get_w(mm_projector_weights, 'mm_projector'))
+
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            decoder_attention_mask: Optional[torch.Tensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            images: Optional[torch.FloatTensor] = None,
+            return_dict: Optional[bool] = None,
+            **kwargs,
+    ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if inputs_embeds is None:
+            _, attention_mask, decoder_attention_mask, past_key_values, inputs_embeds, labels = \
+                self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, decoder_attention_mask,
+                                                          past_key_values, labels, images)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = super(CLIPT5ForConditionalGeneration, self).forward(
+            input_ids=None,  # will be None if inputs_embeds is not None
+            attention_mask=attention_mask,
+            decoder_attention_mask=decoder_attention_mask,
+            labels=labels,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs,
+        )
+
+        return outputs
+
+    @torch.no_grad()
+    def generate(
+            self,
+            inputs: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            images: Optional[torch.Tensor] = None,
+            **kwargs,
+    ):
+        assert images is not None, "images must be provided"
+        assert inputs is not None, "inputs must be provided"
+        assert attention_mask is not None, "attention_mask must be provided"
+        _, attention_mask, _, _, inputs_embeds, _ = \
+            self.prepare_inputs_labels_for_multimodal(inputs, attention_mask, None, None, None, images)
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = super(CLIPT5ForConditionalGeneration, self).generate(
+            input_ids=None,  # will be None if inputs_embeds is not None
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+        )
+        return outputs
+
+    def prepare_inputs_for_generation(
+            self,
+            input_ids,
+            past_key_values=None,
+            attention_mask=None,
+            head_mask=None,
+            decoder_head_mask=None,
+            decoder_attention_mask=None,
+            cross_attn_head_mask=None,
+            use_cache=None,
+            encoder_outputs=None,
+            inputs_embeds=None,
+            **kwargs,
+    ):
+        # cut decoder_input_ids if past_key_values is used
+        if past_key_values is not None:
+            past_length = past_key_values[0][0].shape[2]
+
+            # Some generation methods already pass only the last input ID
+            if input_ids.shape[1] > past_length:
+                remove_prefix_length = past_length
+            else:
+                # Default to old behavior: keep only final ID
+                remove_prefix_length = input_ids.shape[1] - 1
+
+            input_ids = input_ids[:, remove_prefix_length:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update({
+            "decoder_input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "encoder_outputs": encoder_outputs,
+            "attention_mask": attention_mask,
+            "head_mask": head_mask,
+            "decoder_head_mask": decoder_head_mask,
+            "decoder_attention_mask": decoder_attention_mask,
+            "cross_attn_head_mask": cross_attn_head_mask,
+            "use_cache": use_cache,
+        })
+        return model_inputs
+
+
+AutoConfig.register("clip_t5", CLIPT5Config)
+AutoModelForSeq2SeqLM.register(CLIPT5Config, CLIPT5ForConditionalGeneration)
+
+
+_vqascore_cache = {}
+
+
+def load_vqascore(device, dtype, use_fsdp=True):
+    # Create cache key from arguments
+    cache_key = f"{device}_{dtype}_{use_fsdp}"
+
+    # Check if model is already cached
+    if cache_key in _vqascore_cache:
+        return _vqascore_cache[cache_key]
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        'google/flan-t5-xxl', use_fast=False, model_max_length=2048)
+    model = CLIPT5ForConditionalGeneration.from_pretrained(
+        'zhiqiulin/clip-flant5-xxl',
+        torch_dtype=dtype,
+        use_cache=False,
+        freeze_mm_mlp_adapter=True)
+    model.requires_grad_(False)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if use_fsdp:
+        mmcv.print_log('Wrapping VQAScore model with FSDP.')
+        model = FullyShardedDataParallel(
+            model,
+            device_id=torch.cuda.current_device(),
+            use_orig_params=False,
+            mixed_precision=MixedPrecision(
+                param_dtype=dtype,
+                reduce_dtype=dtype,
+                buffer_dtype=dtype,
+                cast_root_forward_inputs=False),
+            sharding_strategy=ShardingStrategy.HYBRID_SHARD,
+            auto_wrap_policy=ModuleWrapPolicy([T5Block]))
+    else:
+        model.to(device)
+
+    result = model, tokenizer
+    _vqascore_cache[cache_key] = result
+    return result
+
+
+@METRICS.register_module()
+class VQAScore(Metric):
+    name = 'VQAScore'
+    requires_prompt = True
+
+    def __init__(self,
+                 num_images=None,
+                 use_fsdp=True):
+        super().__init__(num_images)
+        use_fsdp = use_fsdp and torch.cuda.is_available() and dist.is_initialized() and dist.get_world_size() > 0
+
+        self.use_fsdp = use_fsdp
+        self.dtype = torch.bfloat16
+        self.device = 'cuda' if use_fsdp else 'cpu'
+
+        self.model, self.tokenizer = load_vqascore(device=self.device, dtype=self.dtype, use_fsdp=use_fsdp)
+        self.model.eval()
+        image_processor = self.model.get_vision_tower().image_processor
+        image_size = tuple(image_processor.crop_size.values())
+        assert len(image_size) == 2 and image_size[0] == image_size[1]
+        self.image_size = image_size[0]
+        self.image_mean = torch.tensor(image_processor.image_mean, device=self.device).view(3, 1, 1)
+        self.image_std = torch.tensor(image_processor.image_std, device=self.device).view(3, 1, 1)
+        self.clamp_high = (1 - self.image_mean) / self.image_std
+        self.clamp_low = -self.image_mean / self.image_std
+
+    def prepare(self):
+        self.scores = []
+
+    def resize(self, imgs):
+        h, w = imgs.shape[2:]
+        if h != w:
+            pad_size = max(h, w)
+            pad_h = pad_size - h
+            pad_w = pad_size - w
+            imgs = F.pad(
+                imgs, (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2), mode='constant', value=0)
+            h = w = pad_size
+        if h != self.image_size:
+            imgs = F.interpolate(imgs, size=self.image_size, mode='bicubic', align_corners=False, antialias=True)
+            imgs = torch.maximum(torch.minimum(imgs, self.clamp_high), self.clamp_low)
+        return imgs
+
+    @torch.no_grad()
+    def feed_op(self, batch, mode):
+        imgs = batch['imgs']
+        prompts = batch['prompts']
+
+        imgs = (imgs.to(device=self.device, dtype=torch.float32) / 2 + 0.5).clamp(0, 1)
+        imgs = self.resize((imgs - self.image_mean) / self.image_std).to(dtype=self.dtype)
+
+        # ========= preprocess prompts =========
+        questions = [default_question_template.format(prompt) for prompt in prompts]
+        answers = [default_answer_template.format(prompt) for prompt in prompts]
+
+        questions = [format_question(question, conversation_style='t5_chat') for question in questions]
+        answers = [format_answer(answer, conversation_style='t5_chat') for answer in answers]
+
+        input_ids = [t5_tokenizer_image_token(question, self.tokenizer, return_tensors='pt') for question in questions]
+        labels = [t5_tokenizer_image_token(answer, self.tokenizer, return_tensors='pt') for answer in answers]
+
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=0)[:, :self.tokenizer.model_max_length]
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX)[:, :self.tokenizer.model_max_length]
+
+        input_ids = input_ids.to(device=self.device)
+        labels = labels.to(device=self.device)
+
+        attention_mask = input_ids.ne(self.tokenizer.pad_token_id).to(device=self.device)
+        decoder_attention_mask = labels.ne(IGNORE_INDEX).to(device=self.device)
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_attention_mask=decoder_attention_mask,
+            labels=labels,
+            images=imgs,
+            past_key_values=None,
+            inputs_embeds=None,
+            use_cache=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=True,
+        )
+
+        logits = outputs.logits
+        bs, seq_len, vocab_size = logits.size()
+        vqa_score = (-F.cross_entropy(
+            logits.reshape(bs * seq_len, vocab_size), labels.reshape(bs * seq_len), reduction='none'
+        ).reshape(bs, 2).mean(dim=1)).exp()  # (bs, )
+
+        if dist.is_initialized():
+            ws = dist.get_world_size()
+            placeholder = [torch.empty_like(vqa_score) for _ in range(ws)]
+            dist.all_gather(placeholder, vqa_score)
+            vqa_score = torch.stack(placeholder, dim=1).reshape(vqa_score.size(0) * ws)
+
+        if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized():
+            self.scores.append(vqa_score.float().cpu())
+
+    def feed(self, batch, mode):
+        if mode == 'reals':
+            return 0
+
+        if self.num_images is None:
+            self.feed_op(batch, mode)
+
+        else:
+            _, ws = get_dist_info()
+
+            if self.num_fake_feeded == self.num_fake_need:
+                return 0
+
+            if isinstance(batch, dict):
+                batch_size = len(list(batch.values())[0])
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = {k: v[:end] for k, v in batch.items()}
+            else:
+                batch_size = batch.shape[0]
+                end = min(batch_size, self.num_fake_need - self.num_fake_feeded)
+                batch_to_feed = batch[:end]
+
+            global_end = min(batch_size * ws,
+                             self.num_fake_need - self.num_fake_feeded)
+            self.feed_op(batch_to_feed, mode)
+            self.num_fake_feeded += global_end
+            return end
+
+    @torch.no_grad()
+    def summary(self):
+        scores = torch.cat(self.scores, dim=0)
+        if self.num_images is not None:
+            assert scores.shape[0] >= self.num_images
+            scores = scores[:self.num_images]
+        mean_score = scores.mean().item()
+        self._result_dict = dict(vqascore=mean_score)
+        self._result_str = f'VQAScore: {mean_score:.4f}'
+        return mean_score
+
+    def clear_fake_data(self):
+        self.scores = []
+        self.num_fake_feeded = 0
+
+    def clear(self, clear_reals=False):
+        self.clear_fake_data()
+
+    def load_to_gpu(self):
+        if torch.cuda.is_available() and not isinstance(self.model, FullyShardedDataParallel):
+            self.model.cuda()
+            self.image_mean = self.image_mean.cuda()
+            self.image_std = self.image_std.cuda()
+            self.clamp_high = self.clamp_high.cuda()
+            self.clamp_low = self.clamp_low.cuda()
+            self.device = 'cuda'
+
+    def offload_to_cpu(self):
+        if not isinstance(self.model, FullyShardedDataParallel):
+            self.model.cpu()
+            self.image_mean = self.image_mean.cpu()
+            self.image_std = self.image_std.cpu()
+            self.clamp_high = self.clamp_high.cpu()
+            self.clamp_low = self.clamp_low.cpu()
+            self.device = 'cpu'

lakonlab/models/__init__.py

@@ -0,0 +1,6 @@
+from .losses import *
+from .architecture import *
+from .diffusions import *
+from .diffusion_2d import Diffusion2D
+from .latent_diffusion_class_image import LatentDiffusionClassImage
+from .latent_diffusion_text_image import LatentDiffusionTextImage

lakonlab/models/architecture/__init__.py

@@ -0,0 +1,4 @@
+from .ddpm import *
+from .diffusers import *
+from .gmflow import *
+from .dxflow import *

lakonlab/models/architecture/ddpm/__init__.py

@@ -0,0 +1,6 @@
+from .denoising import DenoisingUnetMod
+from .modules import (
+    MultiHeadAttentionMod, DenoisingResBlockMod, DenoisingDownsampleMod, DenoisingUpsampleMod)
+
+__all__ = ['DenoisingUnetMod', 'MultiHeadAttentionMod', 'DenoisingResBlockMod',
+           'DenoisingDownsampleMod', 'DenoisingUpsampleMod']