Initial UniSpace demo

.gitattributes

@@ -37,3 +37,7 @@ examples/background_change-input.png filter=lfs diff=lfs merge=lfs -text
 examples/color_alter-input.png filter=lfs diff=lfs merge=lfs -text
 examples/material_alter-input.png filter=lfs diff=lfs merge=lfs -text
 examples/motion_change-input.png filter=lfs diff=lfs merge=lfs -text
+examples/background_change.png filter=lfs diff=lfs merge=lfs -text
+examples/color_alter.png filter=lfs diff=lfs merge=lfs -text
+examples/material_alter.png filter=lfs diff=lfs merge=lfs -text
+examples/subject-add.png filter=lfs diff=lfs merge=lfs -text

LICENSE-UNISPACE

@@ -0,0 +1,29 @@
+MIT License
+
+Copyright (c) 2026 UniSpace Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+Note: The `unispace/` subdirectory (the unified multimodal model) contains
+code originally derived from the ByteDance Bagel project and is licensed
+under the Apache License, Version 2.0. See unispace/LICENSE-APACHE and
+unispace/NOTICE for details. The remaining components are licensed under the
+MIT License above.

README.md

@@ -1,33 +1,65 @@
 ---
 title: UniSpace
-emoji: 🎨
+emoji: 🛰️
 colorFrom: blue
 colorTo: gray
 sdk: gradio
 sdk_version: 6.25.0
 app_file: app.py
-short_description: UniSpace unified image generation and editing
+short_description: Unified text-to-image generation and instruction editing
 python_version: "3.12"
 startup_duration_timeout: 1h
+pinned: false
+license: mit
 ---
 
-# UniSpace: Unified Visual Representation and Scalable Multimodal Modeling
+# UniSpace
 
-This Space demos **UniSpace** — a unified vision model supporting text-to-image generation and instruction-based image editing using a Qwen3-8B Mixture-of-Transformers backbone with a patch-reparameterized visual tokenizer.
+Interactive demo of **UniSpace: Unified Visual Representation and Scalable
+Multimodal Modeling** — an 8B Qwen3 Mixture-of-Transformers built on a single
+patch-reparameterized ("PR") visual tokenizer that handles both image
+generation and instruction-based image editing.
 
-## Modes
+* Paper: <https://huggingface.co/papers/2608.08676>
+* Weights: <https://huggingface.co/yjb6/UniSpace>
+* Code: <https://github.com/yjb6/UniSpace>
 
-- **Text-to-Image Generation**: Provide a text prompt to generate an image from scratch.
-- **Image Editing**: Provide an input image and an editing instruction.
+## What runs here
 
-## Model
+The Space loads the released `unispace-sft-0012000` checkpoint together with
+the `pr-qwen-vit` tokenizer/decoder and exposes the authors' two reference
+inference paths:
 
-- **Checkpoint**: `yjb6/UniSpace` (SFT step 0012000, ~61 GB bf16)
-- **VAE**: PR-Qwen-ViT patch-reparameterized tokenizer
-- **Backbone**: Qwen3-8B with Mixture-of-Transformers (separate gen/und expert branches)
+| Tab | Reference script | Settings |
+|---|---|---|
+| Text to image | `eval/gen/gen_images_qwen3_unified_mot.py` | 1024², 50 steps, CFG 10, `timestep_shift=0.112`, `cfg_renorm_min=0.9` |
+| Instruction editing | `eval/gen/gen_images_imgedit_qwen3_unified_mot.py` | `--match-ref-size --ref-use-mar` (mar_1024 buckets), 50 steps, CFG 10, `timestep_shift=0.112`, `cfg_renorm_min=0.0` |
 
-## References
+Both share one model instance built with the release configuration
+(`max_latent_size=96`, `share_unified2llm=False`, `use_spatial_merge_und=True`,
+`use_spatial_merge_gen=False`, `latent_patch_size=1`).
 
-- [Paper](https://huggingface.co/papers/2608.08676)
-- [GitHub](https://github.com/yjb6/UniSpace)
-- [Project Page](https://yjb6.github.io/UniSpace/)
+The sampling code in `unispace_infer.py` is copied verbatim from those two
+scripts; only the argparse/torchrun driver was removed. Model construction
+differs from the scripts in two ways that are required by the hosting
+environment, not by the model:
+
+1. The base Qwen3-8B weights are never downloaded — every `language_model.*`
+   tensor is present in the UniSpace checkpoint, so only its config and
+   tokenizer are fetched.
+2. The 17-shard fp32 checkpoint is streamed one shard at a time, cast to
+   bf16 and released, and the PatchReparam encoder loads only the Qwen3-VL
+   vision tower (instead of the full 17.5 GB VLM with `device_map="cuda:N"`,
+   which is incompatible with ZeroGPU).
+
+## Credits & licensing
+
+UniSpace code and weights are MIT-licensed (© 2026 UniSpace Authors); the
+vendored `unispace/` and `patch-reparameterization/` sources carry the
+upstream Apache-2.0 / MIT notices (`LICENSE-UNISPACE`, `unispace/NOTICE`,
+`THIRD_PARTY.md`).
+
+The four editing examples in `examples/` are the input images from the
+UniSpace project page gallery, which come from **GEdit-Bench** (MIT, see
+`THIRD_PARTY.md`). Text-to-image example prompts are taken from the UniSpace
+repository's own test prompt set.

THIRD_PARTY.md

@@ -0,0 +1,19 @@
+# Third-party evaluation components
+
+The repository-level MIT license does not relicense third-party benchmarks,
+datasets, models, or services. Obtain external assets from their official
+sources and follow their terms.
+
+| Component | How this repository uses it | Upstream terms |
+|---|---|---|
+| GenEval | Adapted prompt/evaluation utilities under `unispace/eval/gen/geneval/` | MIT notices are retained in the source files. |
+| DPG-Bench | Evaluation package under `unispace/eval/gen/dpg/` | See the bundled `unispace/eval/gen/dpg/LICENSE`. |
+| VIEScore | Adapted evaluator utilities under `unispace/eval/gen/gedit/viescore/` | MIT; the upstream license is bundled in that directory. |
+| GEdit-Bench | Downloaded externally at evaluation time; the dataset is not redistributed here. | The official `stepfun-ai/GEdit-Bench` dataset card declares MIT. |
+| OneIG-Bench | `score.sh` invokes a separately obtained official checkout; official prompts, models, and scorer code are not redistributed here. | The official dataset card declares CC BY-NC 4.0. The official GitHub scorer repository did not provide a root license when checked for this preview, so users must review its current terms before use or redistribution. |
+| ImgEdit | The repository contains adapted benchmark orchestration and score aggregation; benchmark data is external. | The official `PKU-YuanGroup/ImgEdit` repository did not provide a root license when checked for this preview. Obtain the benchmark separately and review its current data, model, and judge terms before use or redistribution. |
+
+Model dependencies such as Qwen3, SigLIP2, DINOv2, and evaluator checkpoints
+are not redistributed by this source tree and remain subject to their own
+licenses. API-based judging also remains subject to the selected provider's
+terms.

app.py

@@ -1,766 +1,554 @@
-"""UniSpace: Unified Visual Representation and Scalable Multimodal Modeling.
+"""UniSpace — unified text-to-image generation and instruction-based image editing.
 
-Gradio demo for text-to-image generation and instruction-based image editing
-using the UniSpace model (Qwen3-8B MoT + PR-Qwen-ViT tokenizer).
+Paper: UniSpace: Unified Visual Representation and Scalable Multimodal Modeling
+Model: https://huggingface.co/yjb6/UniSpace
+Code:  https://github.com/yjb6/UniSpace
 """
 
 import os
-# Allocator config for memory pressure (large model, transient allocation spikes)
+
+# Must be set before torch / the UniSpace modeling code is imported.
+os.environ.setdefault("ATTENTION_BACKEND", "torch_sdpa")
+os.environ.setdefault("USE_FLEX_ATTENTION", "0")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-import sys
-import time
-import gc
-import shutil
-from pathlib import Path
-from dataclasses import dataclass
-
-import spaces  # MUST come before torch / any CUDA-touching import
-import torch
-import numpy as np
-from PIL import Image
-import gradio as gr
-
-# ─── Path setup ─────────────────────────────────────────────────────────────
-# The Space repo includes the unispace/ and patch-reparameterization/ packages
-# The unispace/ package uses absolute imports like `from modeling.X` so we add
-# the unispace/ directory itself to sys.path
-SPACE_ROOT = Path(__file__).parent
-sys.path.insert(0, str(SPACE_ROOT / "unispace"))
-sys.path.insert(0, str(SPACE_ROOT / "patch-reparameterization" / "src"))
-
-# ─── HF model paths ────────────────────────────────────────────────────────
-UNISPACE_CKPT_ID = "yjb6/UniSpace"
-UNISPACE_SFT_SUBDIR = "unispace-sft-0012000"
-QWEN3_8B_ID = "Qwen/Qwen3-8B"          # for config + tokenizer only (no weights download)
-QWEN3_VL_8B_ID = "Qwen/Qwen3-VL-8B-Instruct"  # for VAE encoder backbone
-VAE_CONFIG_PATH = str(SPACE_ROOT / "patch-reparameterization" / "configs" / "vae" / "qwen3-temp-run-mar1024_eval.yaml")
-
-# ─── Inference config (from eval_sft_0012000.example.json) ─────────────────
-CFG_SCALE = 10.0
-INFERENCE_STEPS = 50
-TIMESTEP_SHIFT = 0.112
-MAX_LATENT_SIZE = 96
-RESOLUTION = 1024
-SHARE_UNIFIED2LLM = False
-USE_SPATIAL_MERGE_UND = True
-
-print(f"[app.py] SPACE_ROOT={SPACE_ROOT}", flush=True)
-print(f"[app.py] VAE_CONFIG_PATH={VAE_CONFIG_PATH}", flush=True)
-
-
-# ─── Model loading ──────────────────────────────────────────────────────────
-def _download_vae_assets():
-    """Download the PR-Qwen-ViT tokenizer checkpoint and normalization stats."""
-    from huggingface_hub import hf_hub_download
-    vae_ckpt = hf_hub_download(UNISPACE_CKPT_ID, "encoders/pr-qwen-vit-tokenizer.pt", repo_type="model")
-    stats_ckpt = hf_hub_download(UNISPACE_CKPT_ID, "stats/pr-qwen-vit-normalization-stats.pt", repo_type="model")
-    return vae_ckpt, stats_ckpt
-
-
-def _download_safetensors_index():
-    """Download the safetensors index json for the SFT checkpoint."""
-    from huggingface_hub import hf_hub_download
-    return hf_hub_download(
-        UNISPACE_CKPT_ID,
-        f"{UNISPACE_SFT_SUBDIR}/model.safetensors.index.json",
-        repo_type="model",
-    )
+import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)
 
+import gc  # noqa: E402
+import json  # noqa: E402
+import random  # noqa: E402
+import sys  # noqa: E402
+import time  # noqa: E402
 
-def _make_vae_config_yaml(stats_path, ckpt_path):
-    """Create a local VAE config yaml with resolved paths."""
-    import yaml
-    config = {
-        "stage_1": {
-            "target": "stage1.PatchReparam",
-            "params": {
-                "encoder_cls": "Qwen3Unified",
-                "encoder_config_path": QWEN3_VL_8B_ID,
-                "encoder_input_size": 256,
-                "encoder_params": {
-                    "model_name": QWEN3_VL_8B_ID,
-                    "recon_proj_dim": 128,
-                    "pixel_format": "c_major",
-                },
-                "decoder_config_path": "configs/decoder/ViTXL_nockpt",
-                "decoder_cls": "ViTDecoder",
-                "noise_tau": 0.8,
-                "reshape_to_2d": True,
-                "normalization_stat_path": stats_path,
-            },
-            "checkpoint": {
-                "path": ckpt_path,
-                "use_ema": True,
-            },
-        },
-    }
-    tmp_path = "/tmp/vae_config_resolved.yaml"
-    with open(tmp_path, "w") as f:
-        yaml.dump(config, f, default_flow_style=False)
-    return tmp_path
-
-
-def load_unispace_model():
-    """Load the UniSpace model for T2I generation and image editing."""
-    from transformers import AutoTokenizer, AutoConfig
-
-    # Import model code
-    from modeling.qwen3.configuration_qwen3 import Qwen3Config
-    from modeling.unimm.unimm_mot import UnimmMoT, UnimmConfig
-    from modeling.unimm.qwen3_mot import Qwen3MoTForConditionalGeneration
-    from modeling.autoencoder_rae import load_rae
-
-    import torch._dynamo
-    torch._dynamo.config.disable = True
-
-    device = "cuda"
-
-    # 1. Download VAE assets
-    print("[load] Downloading VAE assets...", flush=True)
-    vae_ckpt_path, stats_path = _download_vae_assets()
-    print(f"[load] VAE checkpoint: {vae_ckpt_path}", flush=True)
-    print(f"[load] Stats: {stats_path}", flush=True)
-
-    # 2. Create resolved VAE config
-    vae_config_path = _make_vae_config_yaml(stats_path, vae_ckpt_path)
-
-    # 3. Load Qwen3-8B config (no weights — we use SFT checkpoint)
-    print("[load] Loading Qwen3-8B config...", flush=True)
-    llm_config = Qwen3Config.from_pretrained(QWEN3_8B_ID)
-    llm_config.qk_norm = True
-    llm_config.tie_word_embeddings = False
-    llm_config._attn_implementation = "sdpa"
-    # Ensure pad_token_id exists (transformers 5.x may not set it)
-    if not hasattr(llm_config, "pad_token_id") or llm_config.pad_token_id is None:
-        llm_config.pad_token_id = llm_config.eos_token_id
-    # Fix rope_scaling for transformers 5.x: 'default' was removed, use 'linear' with factor=1.0
-    if hasattr(llm_config, "rope_scaling") and isinstance(llm_config.rope_scaling, dict):
-        if llm_config.rope_scaling.get("rope_type") == "default":
-            llm_config.rope_scaling = {
-                "rope_type": "linear",
-                "factor": 1.0,
-                "rope_theta": llm_config.rope_scaling.get("rope_theta", 1000000.0),
-            }
-
-    # 4. Load VAE (RAE / PatchReparam)
-    print("[load] Loading RAE VAE...", flush=True)
-    vae_model, vae_config = load_rae(
-        local_path=vae_config_path,
-        rae_accelerate_root=str(SPACE_ROOT / "patch-reparameterization"),
-    )
+import gradio as gr  # noqa: E402
+import torch  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+from omegaconf import OmegaConf  # noqa: E402
+from safetensors import safe_open  # noqa: E402
 
-    # 5. Build UnimmConfig
-    config = UnimmConfig(
-        visual_gen=True,
-        visual_und=True,  # needed for editing
-        llm_config=llm_config,
-        vit_config=None,
-        vae_config=vae_config,
-        vit_max_num_patch_per_side=MAX_LATENT_SIZE,
-        connector_act="gelu_pytorch_tanh",
-        latent_patch_size=1,
-        max_latent_size=MAX_LATENT_SIZE,
-        use_qwen_vit=False,
-        use_moe=True,
-        use_qwen3_unified=True,
-        share_unified2llm=SHARE_UNIFIED2LLM,
-        use_spatial_merge=False,
-        use_spatial_merge_und=USE_SPATIAL_MERGE_UND,
-        use_mrope=False,
-    )
+ROOT = os.path.dirname(os.path.abspath(__file__))
+PR_ROOT = os.path.join(ROOT, "patch-reparameterization")
+sys.path.insert(0, os.path.join(PR_ROOT, "src"))
+sys.path.insert(0, os.path.join(ROOT, "unispace"))
 
-    # 6. Instantiate model from config (empty weights — no Qwen3-8B download)
-    print("[load] Instantiating Qwen3MoTForConditionalGeneration (no weight download)...", flush=True)
-    language_model = Qwen3MoTForConditionalGeneration.__new__(Qwen3MoTForConditionalGeneration)
-    Qwen3MoTForConditionalGeneration.__init__(language_model, llm_config)
-
-    # 7. Load tokenizer
-    print("[load] Loading tokenizer...", flush=True)
-    tokenizer = AutoTokenizer.from_pretrained(QWEN3_8B_ID)
-
-    # 8. Build UnimmMoT
-    print("[load] Building UnimmMoT...", flush=True)
-    model = UnimmMoT(language_model, None, tokenizer, config)
-
-    new_token_ids = {
-        "bos_token_id": tokenizer.bos_token_id if tokenizer.bos_token_id else tokenizer.encode("<|endoftext|>")[0],
-        "eos_token_id": tokenizer.eos_token_id if tokenizer.eos_token_id else tokenizer.encode("<|endoftext|>")[0],
-        "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
-        "start_of_image": tokenizer.encode("<|vision_start|>")[0],
-        "end_of_image": tokenizer.encode("<|vision_end|>")[0],
-    }
-
-    # 9. Load SFT checkpoint
-    print("[load] Loading SFT checkpoint...", flush=True)
-    from safetensors.torch import load_file
-    from huggingface_hub import snapshot_download
-
-    # Download all SFT shards
-    sft_dir = snapshot_download(
-        UNISPACE_CKPT_ID,
-        allow_patterns=[f"{UNISPACE_SFT_SUBDIR}/*.safetensors", f"{UNISPACE_SFT_SUBDIR}/*.json"],
-        repo_type="model",
-    )
-    sft_dir = os.path.join(sft_dir, UNISPACE_SFT_SUBDIR)
-
-    # Load all shards into a single state dict
-    index_path = os.path.join(sft_dir, "model.safetensors.index.json")
-    import json
-    with open(index_path) as f:
-        index = json.load(f)
-
-    state_dict = {}
-    shard_files = sorted(set(index["weight_map"].values()))
-    print(f"[load] Loading {len(shard_files)} shards...", flush=True)
-    for shard_file in shard_files:
-        shard_path = os.path.join(sft_dir, shard_file)
-        print(f"  Loading {shard_file}...", flush=True)
-        shard_sd = load_file(shard_path, device="cpu")
-        state_dict.update(shard_sd)
-        del shard_sd
-        gc.collect()
+MODEL_REPO = "yjb6/UniSpace"
+CKPT_DIR = "unispace-sft-0012000"
+QWEN_LLM = "Qwen/Qwen3-8B"
+QWEN_VL = "Qwen/Qwen3-VL-8B-Instruct"
 
-    # Handle ExactMergerConnector format
-    if "unified2llm_und.norm.weight" in state_dict and hasattr(model, "unified2llm_und"):
-        fc1_w = state_dict["unified2llm_und.fc1.weight"]
-        fc2_w = state_dict["unified2llm_und.fc2.weight"]
-        nm_w = state_dict["unified2llm_und.norm.weight"]
-        model.rebuild_unified2llm_und_as_exact_merger(
-            total_in=fc1_w.shape[1],
-            merger_hidden=fc1_w.shape[0],
-            out_dim=fc2_w.shape[0],
-            sem_per_patch=nm_w.shape[0],
+# Official inference settings (unispace/scripts/eval/run_release_generation.sh +
+# unispace/eval/gen/eval_sft_0012000.example.json).
+MAX_LATENT_SIZE = 96
+TIMESTEP_SHIFT = 0.112
+DEFAULT_STEPS = 50
+DEFAULT_CFG = 10.0
+GEN_CFG_RENORM_MIN = 0.9   # gen_images_qwen3_unified_mot.py main()
+EDIT_CFG_RENORM_MIN = 0.0  # gen_images_imgedit_qwen3_unified_mot.py argparse default
+EDIT_RESOLUTION = 1024
+
+DEVICE = "cuda"
+DEVICE_TYPE = "cuda"
+
+
+def _log(msg):
+    print(f"[unispace] {msg}", flush=True)
+
+
+def _free_cached_file(path):
+    """Delete a hub-cached file (symlink + blob) to keep the Space disk small."""
+    try:
+        real = os.path.realpath(path)
+        if os.path.islink(path):
+            os.unlink(path)
+        if os.path.exists(real):
+            os.remove(real)
+    except OSError as exc:  # pragma: no cover
+        _log(f"could not free {path}: {exc}")
+
+
+# ---------------------------------------------------------------------------
+# 1. Vision-tower-only loader for the PatchReparam encoder.
+#
+# stage1/encoders/qwen3_unified.py loads the *full* Qwen3-VL-8B-Instruct
+# (17.5 GB) only to keep `.visual`, and does so with device_map="cuda:N",
+# which is incompatible with the ZeroGPU hijack. We swap in a loader that
+# builds the vision tower on CPU and pulls its weights from the single shard
+# that holds them (2.7 GB).
+# ---------------------------------------------------------------------------
+import transformers.models.qwen3_vl.modeling_qwen3_vl as _q3vl_mod  # noqa: E402
+from transformers import AutoConfig, AutoTokenizer  # noqa: E402
+from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel  # noqa: E402
+
+try:
+    from transformers.modeling_utils import no_init_weights
+except ImportError:  # pragma: no cover
+    from contextlib import nullcontext as no_init_weights
+
+
+class _VisionOnlyQwen3VL:
+    """Minimal stand-in exposing the two attributes Qwen3Unified touches."""
+
+    def __init__(self, visual):
+        self.visual = visual
+        self.model = None
+
+    @classmethod
+    def from_pretrained(cls, model_name, torch_dtype=None, device_map=None, **kwargs):
+        cfg = AutoConfig.from_pretrained(model_name)
+        vis_cfg = cfg.vision_config
+        vis_cfg._attn_implementation = "sdpa"
+        t0 = time.time()
+        with no_init_weights():
+            visual = Qwen3VLVisionModel(vis_cfg)
+        visual = visual.to(torch.float32)
+        shard = hf_hub_download(model_name, "model-00004-of-00004.safetensors")
+        prefix = "model.visual."
+        state = {}
+        with safe_open(shard, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                if key.startswith(prefix):
+                    state[key[len(prefix):]] = f.get_tensor(key).to(torch.float32)
+        missing, unexpected = visual.load_state_dict(state, strict=False)
+        _log(
+            f"vision tower: {len(state)} tensors, missing={len(missing)}, "
+            f"unexpected={len(unexpected)} ({time.time() - t0:.1f}s)"
         )
-        print(f"[load] Rebuilt unified2llm_und as ExactMergerConnector", flush=True)
+        if missing:
+            _log(f"vision tower missing keys (first 10): {missing[:10]}")
+        del state
+        _free_cached_file(shard)
+        gc.collect()
+        return cls(visual)
+
+
+_q3vl_mod.Qwen3VLForConditionalGeneration = _VisionOnlyQwen3VL
+
+
+# ---------------------------------------------------------------------------
+# 2. PatchReparam ("RAE") tokenizer / decoder
+# ---------------------------------------------------------------------------
+from modeling.autoencoder_unified import load_unified_vae  # noqa: E402
+from modeling.qwen3.configuration_qwen3 import Qwen3Config  # noqa: E402
+from modeling.unimm.qwen3_mot import Qwen3MoTForConditionalGeneration  # noqa: E402
+from modeling.unimm.unimm_mot import UnimmConfig, UnimmMoT  # noqa: E402
 
-    state_dict.pop("latent_pos_embed.pos_embed", None)
-    state_dict.pop("vit_pos_embed.pos_embed", None)
+import unispace_infer as ui  # noqa: E402
 
-    msg = model.load_state_dict(state_dict, strict=False)
-    print(f"[load] Checkpoint loaded: missing={len(msg.missing_keys)}, unexpected={len(msg.unexpected_keys)}", flush=True)
-    if msg.missing_keys:
-        print(f"  Missing: {msg.missing_keys[:10]}...", flush=True)
-    if msg.unexpected_keys:
-        print(f"  Unexpected: {msg.unexpected_keys[:10]}...", flush=True)
 
-    del state_dict
+def _build_vae_config_yaml():
+    src = os.path.join(PR_ROOT, "configs", "vae", "qwen3-temp-run-mar1024_eval.yaml")
+    cfg = OmegaConf.load(src)
+    cfg.stage_1.params.encoder_config_path = QWEN_VL
+    cfg.stage_1.params.encoder_params.model_name = QWEN_VL
+    cfg.stage_1.params.normalization_stat_path = hf_hub_download(
+        MODEL_REPO, "stats/pr-qwen-vit-normalization-stats.pt"
+    )
+    cfg.stage_1.checkpoint.path = hf_hub_download(
+        MODEL_REPO, "encoders/pr-qwen-vit-tokenizer.pt"
+    )
+    out = os.path.join(ROOT, "_vae_eval_resolved.yaml")
+    OmegaConf.save(cfg, out, resolve=False)
+    return out, str(cfg.stage_1.checkpoint.path)
+
+
+_log("loading PatchReparam tokenizer/decoder ...")
+_t0 = time.time()
+_vae_yaml, _vae_ckpt = _build_vae_config_yaml()
+vae_model, vae_config = load_unified_vae(local_path=_vae_yaml, patch_reparam_root=PR_ROOT)
+vae_model = vae_model.eval()
+_free_cached_file(_vae_ckpt)
+gc.collect()
+_log(
+    f"PatchReparam ready in {time.time() - _t0:.1f}s "
+    f"(downsample={vae_config.downsample}, z_channels={vae_config.z_channels})"
+)
+
+
+# ---------------------------------------------------------------------------
+# 3. Qwen3-8B Mixture-of-Transformers backbone
+#
+# Every `language_model.*` tensor lives in the UniSpace checkpoint, so the base
+# Qwen3-8B weights are never needed — only its config + tokenizer. The 17-shard
+# fp32 checkpoint (61 GB) is streamed one shard at a time, cast to bf16 and
+# released, so peak disk stays around 4 GB.
+# ---------------------------------------------------------------------------
+_log("building Qwen3-8B MoT backbone from config ...")
+_t0 = time.time()
+llm_config = Qwen3Config.from_pretrained(QWEN_LLM)
+llm_config.qk_norm = True
+llm_config.tie_word_embeddings = False
+llm_config._attn_implementation = "sdpa"
+
+_orig_linear_reset = torch.nn.Linear.reset_parameters
+_orig_embedding_reset = torch.nn.Embedding.reset_parameters
+torch.nn.Linear.reset_parameters = lambda self: None
+torch.nn.Embedding.reset_parameters = lambda self: None
+torch.set_default_dtype(torch.bfloat16)
+try:
+    with no_init_weights():
+        language_model = Qwen3MoTForConditionalGeneration(llm_config)
+finally:
+    torch.set_default_dtype(torch.float32)
+    torch.nn.Linear.reset_parameters = _orig_linear_reset
+    torch.nn.Embedding.reset_parameters = _orig_embedding_reset
+_log(f"backbone skeleton built in {time.time() - _t0:.1f}s")
+
+unimm_config = UnimmConfig(
+    visual_gen=True,
+    visual_und=True,
+    llm_config=llm_config,
+    vit_config=None,
+    vae_config=vae_config,
+    vit_max_num_patch_per_side=MAX_LATENT_SIZE,
+    connector_act="gelu_pytorch_tanh",
+    latent_patch_size=1,
+    max_latent_size=MAX_LATENT_SIZE,
+    use_qwen_vit=False,
+    use_moe=True,
+    use_qwen3_unified=True,
+    share_unified2llm=False,
+    use_spatial_merge=False,
+    use_spatial_merge_gen=False,
+    use_spatial_merge_und=True,
+    use_mrope=False,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(QWEN_LLM)
+model = UnimmMoT(language_model, None, tokenizer, unimm_config)
+del language_model
+gc.collect()
+
+new_token_ids = {
+    "bos_token_id": tokenizer.encode("<|im_start|>")[0],
+    "eos_token_id": tokenizer.encode("<|im_end|>")[0],
+    "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
+    "start_of_image": tokenizer.encode("<|vision_start|>")[0],
+    "end_of_image": tokenizer.encode("<|vision_end|>")[0],
+}
+
+# The two positional embeddings are frozen 2-D sincos tables rebuilt at
+# construction time — the official edit script drops them from the checkpoint.
+SKIP_KEYS = {"latent_pos_embed.pos_embed", "vit_pos_embed.pos_embed"}
+
+_log("streaming UniSpace SFT checkpoint ...")
+_t0 = time.time()
+_index = json.load(
+    open(hf_hub_download(MODEL_REPO, f"{CKPT_DIR}/model.safetensors.index.json"))
+)
+_shards = sorted(set(_index["weight_map"].values()))
+_loaded = set()
+for _i, _shard in enumerate(_shards, 1):
+    _path = hf_hub_download(MODEL_REPO, f"{CKPT_DIR}/{_shard}")
+    _sd = {}
+    with safe_open(_path, framework="pt", device="cpu") as f:
+        for _k in f.keys():
+            if _k in SKIP_KEYS:
+                continue
+            _sd[_k] = f.get_tensor(_k).to(torch.bfloat16)
+    _msg = model.load_state_dict(_sd, strict=False)
+    if _msg.unexpected_keys:
+        _log(f"  shard {_i}: unexpected {_msg.unexpected_keys[:5]}")
+    _loaded.update(_sd.keys())
+    del _sd, _msg
     gc.collect()
+    _free_cached_file(_path)
+    _log(f"  shard {_i}/{len(_shards)} loaded ({time.time() - _t0:.0f}s elapsed)")
 
-    # 10. Move to device
-    print("[load] Moving model to cuda (bf16)...", flush=True)
-    model = model.to(device).to(dtype=torch.bfloat16).eval()
-    vae_model = vae_model.to(device).eval()
-    print("[load] Model loaded successfully!", flush=True)
-
-    return model, tokenizer, new_token_ids, vae_model
-
-
-# ─── Load model at module scope ────────────────────────────────────────────
-print("=" * 60, flush=True)
-print("[app.py] Loading UniSpace model...", flush=True)
-print("=" * 60, flush=True)
-MODEL, TOKENIZER, NEW_TOKEN_IDS, VAE_MODEL = load_unispace_model()
-print("[app.py] Model loaded.", flush=True)
-
-# Import NaiveCache for use in inference functions
-from modeling.unimm.qwen3 import NaiveCache
-
-
-# ─── Helper functions ──────────────────────────────────────────────────────
-def _move_to_device(d, device):
-    for k, v in d.items():
-        if isinstance(v, torch.Tensor):
-            d[k] = v.to(device)
-    return d
-
-
-def _decode_latent(latent, h, w, latent_ch, vae_model, use_spatial_merge_gen, device_type="cuda"):
-    if use_spatial_merge_gen:
-        from modeling.unimm.unimm_mot import UnimmMoT
-        latent = UnimmMoT.spatial_unshuffle(latent, h, w)
-    latent = latent.reshape(1, h, w, latent_ch).permute(0, 3, 1, 2)
-    with torch.amp.autocast(device_type, dtype=torch.bfloat16):
-        image = vae_model.decode(latent)
-    return ((image * 0.5 + 0.5).clamp(0, 1)[0].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
-
-
-def _encode_ref_image(pil_image, vae_model, mot_model, image_size, device,
-                      ref_max_size=None, ref_min_size=256, ref_max_pixels=1_048_576,
-                      ref_use_mar=False, ref_mar_resolution=1024):
-    """Encode a reference image via the RAE und branch for editing."""
-    import torchvision.transforms.functional as TF
-    from data.transforms import MaxLongEdgeMinShortEdgeResize
-    from data.data_utils import get_flattened_position_ids_extrapolate
-
-    img = pil_image.convert("RGB")
-
-    if ref_use_mar:
-        target_h, target_w = _compute_target_size_from_ref(pil_image, ref_mar_resolution)
-        img = img.resize((target_w, target_h), Image.BICUBIC)
-        h_img, w_img = target_h, target_w
-    elif ref_max_size is not None:
-        resizer = MaxLongEdgeMinShortEdgeResize(
-            max_size=ref_max_size, min_size=ref_min_size,
-            stride=32, max_pixels=ref_max_pixels,
-        )
-        img = resizer(img)
-        w_img, h_img = img.size
-    else:
-        img = img.resize((image_size, image_size), Image.BICUBIC)
-        h_img = w_img = image_size
-
-    img_tensor = TF.to_tensor(img) * 2.0 - 1.0
-    img_tensor = img_tensor.unsqueeze(0).to(device)
-
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        z = vae_model.encode(img_tensor)
-
-    h_grid = h_img // 16
-    w_grid = w_img // 16
-    z_tokens = z[0].permute(1, 2, 0).reshape(h_grid * w_grid, -1).to(dtype=torch.bfloat16)
-
-    if mot_model.config.use_spatial_merge_und:
-        from modeling.unimm.unimm_mot import UnimmMoT
-        merged_tokens = UnimmMoT.spatial_merge(z_tokens, h_grid, w_grid)
-        patch_size = 32
-    else:
-        merged_tokens = z_tokens
-        patch_size = 16
-
-    pos_ids = get_flattened_position_ids_extrapolate(
-        h_img, w_img, patch_size=patch_size,
-        max_num_patches_per_side=mot_model.max_latent_size,
-    ).to(device)
-
-    return merged_tokens, pos_ids
-
-
-def _compute_target_size_from_ref(ref_image, resolution, stride=16):
-    from data.data_utils import MULTI_RESOLUTION_MAP
-    bucket_key = f"mar_{resolution}"
-    if bucket_key in MULTI_RESOLUTION_MAP:
-        bucket = MULTI_RESOLUTION_MAP[bucket_key]
-        pil_w, pil_h = ref_image.size
-        ref_ratio = pil_h / pil_w
-        best_key = min(bucket.keys(), key=lambda k: abs(float(k) - ref_ratio))
-        H, W = bucket[best_key]
-        return (int(H), int(W))
-    else:
-        pil_w, pil_h = ref_image.size
-        if pil_w >= pil_h:
-            target_w = resolution
-            target_h = resolution * pil_h / pil_w
-        else:
-            target_h = resolution
-            target_w = resolution * pil_w / pil_h
-        target_h = max(stride, round(target_h / stride) * stride)
-        target_w = max(stride, round(target_w / stride) * stride)
-        return (int(target_h), int(target_w))
-
-
-def _prepare_unified_vit_input(curr_kvlens, curr_rope, vit_tokens, pos_ids, new_token_ids, device):
-    packed_text_ids, packed_text_indexes = [], []
-    packed_vit_token_indexes = []
-    packed_position_ids_list, packed_seqlens, packed_indexes = [], [], []
-    packed_key_value_indexes = []
-
-    _curr = curr = 0
-    newlens, new_rope_out = [], []
-    for curr_kvlen, curr_position_id in zip(curr_kvlens, curr_rope):
-        packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
-        curr += curr_kvlen
-
-        packed_text_ids.append(new_token_ids["start_of_image"])
-        packed_text_indexes.append(_curr)
-        packed_indexes.append(curr)
-        curr += 1
-        _curr += 1
-
-        num_img_tokens = vit_tokens.shape[0]
-        packed_vit_token_indexes.extend(range(_curr, _curr + num_img_tokens))
-        packed_indexes.extend(range(curr, curr + num_img_tokens))
-        curr += num_img_tokens
-        _curr += num_img_tokens
-
-        packed_text_ids.append(new_token_ids["end_of_image"])
-        packed_text_indexes.append(_curr)
-        packed_indexes.append(curr)
-        curr += 1
-        _curr += 1
-
-        pos_1d = torch.full((num_img_tokens + 2,), curr_position_id, dtype=torch.long)
-        packed_position_ids_list.append(pos_1d)
-        packed_seqlens.append(num_img_tokens + 2)
-        newlens.append(curr_kvlen + num_img_tokens + 2)
-        new_rope_out.append(curr_position_id + num_img_tokens + 2)
-
-    return {
-        "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long).to(device),
-        "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long).to(device),
-        "packed_vit_tokens": vit_tokens,
-        "packed_vit_token_indexes": torch.tensor(packed_vit_token_indexes, dtype=torch.long).to(device),
-        "packed_vit_position_ids": pos_ids,
-        "packed_position_ids": torch.cat(packed_position_ids_list).to(device),
-        "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int).to(device),
-        "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long).to(device),
-        "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long).to(device),
-        "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int).to(device),
-    }, newlens, new_rope_out
-
-
-# ─── Inference functions ───────────────────────────────────────────────────
-
-@spaces.GPU(duration=240, size="xlarge")
-def generate_image(
+_expected = set(model.state_dict().keys())
+_missing = sorted(_expected - _loaded - SKIP_KEYS)
+_log(f"checkpoint loaded: {len(_loaded)} tensors, missing={len(_missing)}")
+if _missing:
+    _log(f"MISSING KEYS (first 20): {_missing[:20]}")
+
+model = model.to(torch.bfloat16).eval()
+for _p in model.parameters():
+    _p.requires_grad_(False)
+for _p in vae_model.parameters():
+    _p.requires_grad_(False)
+gc.collect()
+
+_log("moving weights to GPU ...")
+model = model.to(DEVICE)
+vae_model = vae_model.to(DEVICE)
+_log("model ready")
+
+
+# ---------------------------------------------------------------------------
+# 4. Gradio handlers
+# ---------------------------------------------------------------------------
+MAX_SEED = 2**31 - 1
+
+
+def _duration_generate(prompt="", steps=DEFAULT_STEPS, *args, **kwargs):
+    try:
+        steps = int(steps)
+    except (TypeError, ValueError):
+        steps = DEFAULT_STEPS
+    return min(240, 40 + int(steps * 1.6))
+
+
+def _duration_edit(image=None, instruction="", steps=DEFAULT_STEPS, *args, **kwargs):
+    try:
+        steps = int(steps)
+    except (TypeError, ValueError):
+        steps = DEFAULT_STEPS
+    return min(280, 55 + int(steps * 2.0))
+
+
+@spaces.GPU(duration=_duration_generate)
+def generate(
     prompt: str,
-    steps: int = 50,
-    cfg_scale: float = 10.0,
-    seed: int = 42,
+    steps: int = DEFAULT_STEPS,
+    cfg_scale: float = DEFAULT_CFG,
     resolution: int = 1024,
+    seed: int = 42,
+    randomize_seed: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate an image from a text prompt using UniSpace flow-matching.
+    """Generate a square image from a text prompt with UniSpace.
 
     Args:
         prompt: Text description of the image to generate.
-        steps: Number of denoising steps (default 50).
-        cfg_scale: Classifier-free guidance scale (default 10.0).
-        seed: Random seed for reproducibility (default 42).
-        resolution: Output image resolution (default 1024).
+        steps: Number of flow-matching sampling steps.
+        cfg_scale: Classifier-free guidance strength.
+        resolution: Output width/height in pixels.
+        seed: Random seed.
+        randomize_seed: Draw a fresh random seed instead of using `seed`.
 
     Returns:
-        Generated PIL Image.
+        The generated PIL image, the seed that was used, and a timing note.
     """
-    device = "cuda"
-    gen_model = MODEL
-    tokenizer = TOKENIZER
-    new_token_ids = NEW_TOKEN_IDS
-    vae_model = VAE_MODEL
-
-    generator = torch.Generator("cpu").manual_seed(seed) if seed is not None else None
-
-    _llm_cfg = gen_model.config.llm_config
-    num_layers = (
-        _llm_cfg.text_config.num_hidden_layers
-        if hasattr(_llm_cfg, "text_config") and _llm_cfg.text_config is not _llm_cfg
-        else _llm_cfg.num_hidden_layers
-    )
-    past_key_values = NaiveCache(num_layers)
-    newlens = [0]
-    new_rope = [0]
-
-    # 1. Encode prompt
-    generation_input_text, newlens, new_rope = gen_model.prepare_prompts(
-        curr_kvlens=newlens, curr_rope=new_rope,
-        prompts=[prompt], tokenizer=tokenizer,
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    t0 = time.perf_counter()
+    images = ui.generate_image(
+        prompt=prompt.strip(),
+        num_timesteps=int(steps),
+        cfg_scale=float(cfg_scale),
+        cfg_interval=[0, 1.0],
+        cfg_renorm_min=GEN_CFG_RENORM_MIN,
+        timestep_shift=TIMESTEP_SHIFT,
+        max_t=1.0,
+        num_images=1,
+        resolution=int(resolution),
+        inference_mode="flash",
+        device=DEVICE,
+        device_type=DEVICE_TYPE,
+        gen_model=model,
+        tokenizer=tokenizer,
         new_token_ids=new_token_ids,
+        vae_model=vae_model,
+        seed=seed,
     )
-    generation_input_text = _move_to_device(generation_input_text, device)
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        past_key_values = gen_model.forward_cache_update_text(past_key_values, **generation_input_text)
-
-    # 2. Prepare VAE latent slots
-    generation_input = gen_model.prepare_vae_latent(
-        curr_kvlens=newlens, curr_rope=new_rope,
-        image_sizes=[(resolution, resolution)],
-        new_token_ids=new_token_ids, generator=generator,
-    )
-    generation_input = _move_to_device(generation_input, device)
-
-    # 3. CFG prefix (text-free)
-    cfg_past_key_values = NaiveCache(num_layers)
-    cfg_newlens = [0]
-    cfg_new_rope = [0]
-    cfg_prompt_input, cfg_newlens, cfg_new_rope = gen_model.prepare_cfg_prompts(
-        curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope,
-        num_images=1, tokenizer=tokenizer, new_token_ids=new_token_ids,
-    )
-    cfg_prompt_input = _move_to_device(cfg_prompt_input, device)
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        cfg_past_key_values = gen_model.forward_cache_update_text(cfg_past_key_values, **cfg_prompt_input)
-
-    generation_input_cfg = gen_model.prepare_vae_latent_cfg(
-        curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope,
-        image_sizes=[(resolution, resolution)],
-    )
-    generation_input_cfg = _move_to_device(generation_input_cfg, device)
-
-    # 4. Flow matching denoising
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        unpacked_latent = gen_model.generate_image(
-            past_key_values=past_key_values,
-            num_timesteps=steps,
-            cfg_text_scale=cfg_scale,
-            cfg_interval=[0, 1.0],
-            cfg_renorm_min=0.9,
-            timestep_shift=TIMESTEP_SHIFT,
-            cfg_text_past_key_values=cfg_past_key_values,
-            cfg_text_packed_position_ids=generation_input_cfg["cfg_packed_position_ids"],
-            cfg_text_key_values_lens=generation_input_cfg["cfg_key_values_lens"],
-            cfg_text_packed_query_indexes=generation_input_cfg["cfg_packed_query_indexes"],
-            cfg_text_packed_key_value_indexes=generation_input_cfg["cfg_packed_key_value_indexes"],
-            inference_mode="flash",
-            **generation_input,
-        )
+    elapsed = time.perf_counter() - t0
+    return images[0], seed, f"{resolution}×{resolution} · {steps} steps · {elapsed:.1f}s"
 
-    # 5. VAE decode
-    latent_ch = gen_model.config.vae_config.z_channels
-    latent_ds = gen_model.config.vae_config.downsample
-    h = w = resolution // latent_ds
-    use_merge_gen = gen_model.config.use_spatial_merge_gen
-    arr = _decode_latent(unpacked_latent[0].float(), h, w, latent_ch, vae_model, use_merge_gen)
 
-    return Image.fromarray(arr)
-
-
-@spaces.GPU(duration=240, size="xlarge")
-def edit_image(
-    input_image: Image.Image,
+@spaces.GPU(duration=_duration_edit)
+def edit(
+    image,
     instruction: str,
-    steps: int = 50,
-    cfg_scale: float = 4.0,
+    steps: int = DEFAULT_STEPS,
+    cfg_scale: float = DEFAULT_CFG,
     seed: int = 42,
-    resolution: int = 1024,
+    randomize_seed: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Edit an image using an instruction with UniSpace.
+    """Edit an image by following a natural-language instruction.
 
     Args:
-        input_image: The source image to edit.
-        instruction: Editing instruction (e.g., "Change the background to a forest").
-        steps: Number of denoising steps (default 50).
-        cfg_scale: Classifier-free guidance scale (default 4.0).
-        seed: Random seed for reproducibility (default 42).
-        resolution: Target resolution (default 1024).
+        image: The source image to edit.
+        instruction: What to change, e.g. "Change the background to a forest.".
+        steps: Number of flow-matching sampling steps.
+        cfg_scale: Classifier-free guidance strength on the instruction.
+        seed: Random seed.
+        randomize_seed: Draw a fresh random seed instead of using `seed`.
 
     Returns:
-        Edited PIL Image.
+        The edited PIL image, the seed that was used, and a timing note.
     """
-    device = "cuda"
-    gen_model = MODEL
-    tokenizer = TOKENIZER
-    new_token_ids = NEW_TOKEN_IDS
-    vae_model = VAE_MODEL
-
-    generator = torch.Generator("cpu").manual_seed(seed) if seed is not None else None
-
-    _llm_cfg = gen_model.config.llm_config
-    num_layers = (
-        _llm_cfg.text_config.num_hidden_layers
-        if hasattr(_llm_cfg, "text_config") and _llm_cfg.text_config is not _llm_cfg
-        else _llm_cfg.num_hidden_layers
-    )
-
-    # Compute target size from input image aspect ratio
-    from data.data_utils import pil_img2rgb
-    src_image = pil_img2rgb(input_image)
-    target_h, target_w = _compute_target_size_from_ref(src_image, resolution)
-    h, w = target_h, target_w
-
-    # Encode ref image via RAE und branch
-    vit_tokens, pos_ids = _encode_ref_image(
-        src_image, vae_model, gen_model, 448, device,
-        ref_use_mar=True, ref_mar_resolution=resolution,
-    )
-
-    # === Conditional path ===
-    past_key_values = NaiveCache(num_layers)
-    newlens = [0]
-    new_rope = [0]
-
-    vit_inp, newlens, new_rope = _prepare_unified_vit_input(
-        newlens, new_rope, vit_tokens, pos_ids, new_token_ids, device)
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        past_key_values = gen_model.forward_cache_update_unified_vit(past_key_values, **vit_inp)
-
-    # user\n{instruction}
-    inp_text, newlens, new_rope = gen_model.prepare_prompts(
-        curr_kvlens=newlens, curr_rope=new_rope,
-        prompts=[f"user\n{instruction}"], tokenizer=tokenizer,
-        new_token_ids=new_token_ids, raw=True,
-    )
-    inp_text = _move_to_device(inp_text, device)
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        past_key_values = gen_model.forward_cache_update_text(past_key_values, **inp_text)
-
-    # assistant\n
-    inp_asst, newlens, new_rope = gen_model.prepare_prompts(
-        curr_kvlens=newlens, curr_rope=new_rope,
-        prompts=["assistant\n"], tokenizer=tokenizer,
-        new_token_ids=new_token_ids, raw=True,
-    )
-    inp_asst = _move_to_device(inp_asst, device)
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        past_key_values = gen_model.forward_cache_update_text(past_key_values, **inp_asst)
-
-    # Target image latent slot
-    generation_input = gen_model.prepare_vae_latent(
-        curr_kvlens=newlens, curr_rope=new_rope,
-        image_sizes=[(h, w)], new_token_ids=new_token_ids, generator=generator,
-    )
-    generation_input = _move_to_device(generation_input, device)
-
-    # === Unconditional path (only ref image, drop instruction) ===
-    cfg_past_key_values = NaiveCache(num_layers)
-    cfg_newlens = [0]
-    cfg_new_rope = [0]
-
-    vit_inp_cfg, cfg_newlens, cfg_new_rope = _prepare_unified_vit_input(
-        cfg_newlens, cfg_new_rope, vit_tokens, pos_ids, new_token_ids, device)
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        cfg_past_key_values = gen_model.forward_cache_update_unified_vit(cfg_past_key_values, **vit_inp_cfg)
-
-    inp_asst_cfg, cfg_newlens, cfg_new_rope = gen_model.prepare_prompts(
-        curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope,
-        prompts=["assistant\n"], tokenizer=tokenizer,
-        new_token_ids=new_token_ids, raw=True,
-    )
-    inp_asst_cfg = _move_to_device(inp_asst_cfg, device)
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        cfg_past_key_values = gen_model.forward_cache_update_text(cfg_past_key_values, **inp_asst_cfg)
-
-    generation_input_cfg = gen_model.prepare_vae_latent_cfg(
-        curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope,
-        image_sizes=[(h, w)],
+    if image is None:
+        raise gr.Error("Please upload an image to edit.")
+    if not instruction or not instruction.strip():
+        raise gr.Error("Please enter an editing instruction.")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    ref = image.convert("RGB")
+    target_size = ui.compute_target_size_from_ref(ref, EDIT_RESOLUTION)
+    t0 = time.perf_counter()
+    out = ui.editing_image(
+        gen_model=model,
+        tokenizer=tokenizer,
+        new_token_ids=new_token_ids,
+        vae_model=vae_model,
+        ref_image=ref,
+        prompt=instruction.strip(),
+        ref_image_size=448,
+        ref_max_size=None,
+        ref_min_size=256,
+        ref_use_mar=True,
+        ref_mar_resolution=EDIT_RESOLUTION,
+        target_size=target_size,
+        num_timesteps=int(steps),
+        cfg_text_scale=float(cfg_scale),
+        cfg_interval=[0, 1.0],
+        cfg_renorm_min=EDIT_CFG_RENORM_MIN,
+        timestep_shift=TIMESTEP_SHIFT,
+        device=DEVICE,
+        device_type=DEVICE_TYPE,
+        seed=seed,
     )
-    generation_input_cfg = _move_to_device(generation_input_cfg, device)
-
-    # === Flow matching generation ===
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
-        unpacked_latent = gen_model.generate_image(
-            past_key_values=past_key_values,
-            num_timesteps=steps,
-            cfg_text_scale=cfg_scale,
-            cfg_interval=[0, 1.0],
-            cfg_renorm_min=0.0,
-            timestep_shift=TIMESTEP_SHIFT,
-            cfg_text_past_key_values=cfg_past_key_values,
-            cfg_text_packed_position_ids=generation_input_cfg["cfg_packed_position_ids"],
-            cfg_text_key_values_lens=generation_input_cfg["cfg_key_values_lens"],
-            cfg_text_packed_query_indexes=generation_input_cfg["cfg_packed_query_indexes"],
-            cfg_text_packed_key_value_indexes=generation_input_cfg["cfg_packed_key_value_indexes"],
-            inference_mode="flash",
-            **generation_input,
-        )
-
-    # === VAE decode ===
-    latent_ch = gen_model.config.vae_config.z_channels
-    latent_ds = gen_model.config.vae_config.downsample
-    h_lat = h // latent_ds
-    w_lat = w // latent_ds
-    use_merge_gen = getattr(gen_model.config, "use_spatial_merge_gen", False)
-    arr = _decode_latent(unpacked_latent[0].float(), h_lat, w_lat, latent_ch, vae_model, use_merge_gen)
-
-    return Image.fromarray(arr)
-
-
-# ─── Gradio UI ─────────────────────────────────────────────────────────────
-CUSTOM_CSS = """
-#main-title { text-align: center; }
-#main-desc { text-align: center; max-width: 700px; margin: 0 auto; }
+    elapsed = time.perf_counter() - t0
+    return out, seed, f"{target_size[1]}×{target_size[0]} · {steps} steps · {elapsed:.1f}s"
+
+
+# Thin wrappers for gr.Examples. They deliberately take no `gr.Progress`
+# parameter: Gradio injects the Progress object at the index of that parameter
+# and would otherwise land it in `steps` when an example row supplies fewer
+# values than the full signature.
+def generate_example(prompt: str):
+    """Run the text-to-image example row with default settings."""
+    return generate(prompt)
+
+
+def edit_example(image, instruction: str):
+    """Run the image-editing example row with default settings."""
+    return edit(image, instruction)
+
+
+# ---------------------------------------------------------------------------
+# 5. UI
+# ---------------------------------------------------------------------------
+GEN_EXAMPLES = [
+    ["A photo of a press conference with microphones and a blurred crowd in the background, realistic news style"],
+    ["A watercolor painting of a lighthouse during a winter sunrise."],
+    ["An ink wash painting style image illustrating a Chinese poem. A solitary boat floats on a misty river near mountains. Vertical traditional Chinese text on the right reads: '孤舟蓑笠翁，独钓寒江雪'. The mood is serene and melancholic. A red seal stamp (chop) is visible in the bottom left corner."],
+    ["A realistic photo of a modern airport terminal sign hanging from the ceiling. The sign has a yellow background with black text. It shows directions with arrows. The text is in three languages: English 'Departures', Chinese '出发', and Japanese '出発'. An icon of an airplane taking off is next to the text. The background shows a blurred terminal hall with travelers."],
+    ["A professional product photo of a sneaker on a white background, studio lighting, advertising style"],
+]
+
+EDIT_EXAMPLES = [
+    ["examples/background_change.png", "Change the background to a forest."],
+    ["examples/color_alter.png", "Alter the color of the mirror frame to orange."],
+    ["examples/material_alter.png", "Change the cup in hand to ceramic."],
+    ["examples/subject-add.png", "Add a balloon decoration strip below the airplane."],
+]
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
 """
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=CUSTOM_CSS) as demo:
-    gr.Markdown(
-        """
-        # 🎨 UniSpace: Unified Visual Representation and Scalable Multimodal Modeling
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="UniSpace") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # 🛰️ UniSpace
+            **Unified Visual Representation and Scalable Multimodal Modeling** — an 8B
+            Mixture-of-Transformers that generates and edits images on top of a single
+            patch-reparameterized visual tokenizer.
+
+            [Paper](https://huggingface.co/papers/2608.08676) ·
+            [Model](https://huggingface.co/yjb6/UniSpace) ·
+            [Code](https://github.com/yjb6/UniSpace)
+            """
+        )
 
-        UniSpace is a unified vision model supporting text-to-image generation and instruction-based image editing,
-        built on a Qwen3-8B Mixture-of-Transformers backbone with a patch-reparameterized visual tokenizer.
-        """
+        with gr.Tabs():
+            with gr.Tab("Text to image"):
+                with gr.Row():
+                    with gr.Column():
+                        g_prompt = gr.Textbox(
+                            label="Prompt",
+                            lines=3,
+                            placeholder="A watercolor painting of a lighthouse during a winter sunrise.",
+                        )
+                        g_run = gr.Button("Generate", variant="primary")
+                        with gr.Accordion("Advanced settings", open=False):
+                            g_resolution = gr.Radio(
+                                choices=[512, 768, 1024],
+                                value=1024,
+                                label="Resolution",
+                            )
+                            g_steps = gr.Slider(
+                                8, 60, value=DEFAULT_STEPS, step=1, label="Sampling steps"
+                            )
+                            g_cfg = gr.Slider(
+                                1.0, 15.0, value=DEFAULT_CFG, step=0.5,
+                                label="Guidance scale (CFG)",
+                            )
+                            g_seed = gr.Slider(
+                                0, MAX_SEED, value=42, step=1, label="Seed"
+                            )
+                            g_rand = gr.Checkbox(value=False, label="Randomize seed")
+                    with gr.Column():
+                        g_out = gr.Image(label="Result", type="pil", height=520)
+                        g_used_seed = gr.Number(label="Seed used", interactive=False)
+                        g_info = gr.Textbox(label="Run info", interactive=False)
+
+                gr.Examples(
+                    examples=GEN_EXAMPLES,
+                    inputs=[g_prompt],
+                    outputs=[g_out, g_used_seed, g_info],
+                    fn=generate_example,
+                    cache_examples=True,
+                    cache_mode="lazy",
+                    label="Prompts from the UniSpace release",
+                )
+
+            with gr.Tab("Instruction editing"):
+                with gr.Row():
+                    with gr.Column():
+                        e_image = gr.Image(label="Input image", type="pil", height=340)
+                        e_prompt = gr.Textbox(
+                            label="Editing instruction",
+                            lines=2,
+                            placeholder="Change the background to a forest.",
+                        )
+                        e_run = gr.Button("Edit", variant="primary")
+                        with gr.Accordion("Advanced settings", open=False):
+                            e_steps = gr.Slider(
+                                8, 60, value=DEFAULT_STEPS, step=1, label="Sampling steps"
+                            )
+                            e_cfg = gr.Slider(
+                                1.0, 15.0, value=DEFAULT_CFG, step=0.5,
+                                label="Guidance scale (CFG)",
+                            )
+                            e_seed = gr.Slider(
+                                0, MAX_SEED, value=42, step=1, label="Seed"
+                            )
+                            e_rand = gr.Checkbox(value=False, label="Randomize seed")
+                    with gr.Column():
+                        e_out = gr.Image(label="Result", type="pil", height=520)
+                        e_used_seed = gr.Number(label="Seed used", interactive=False)
+                        e_info = gr.Textbox(label="Run info", interactive=False)
+
+                gr.Examples(
+                    examples=EDIT_EXAMPLES,
+                    inputs=[e_image, e_prompt],
+                    outputs=[e_out, e_used_seed, e_info],
+                    fn=edit_example,
+                    cache_examples=True,
+                    cache_mode="lazy",
+                    label="Examples from the UniSpace project page (GEdit-Bench, MIT)",
+                )
+
+    g_run.click(
+        fn=generate,
+        inputs=[g_prompt, g_steps, g_cfg, g_resolution, g_seed, g_rand],
+        outputs=[g_out, g_used_seed, g_info],
+        api_name="generate",
     )
-
-    with gr.Tabs():
-        # ─── Text-to-Image Tab ────────────────────────────────────────────
-        with gr.Tab("✏️ Text-to-Image"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    t2i_prompt = gr.Textbox(
-                        label="Prompt",
-                        placeholder="Describe the image you want to generate...",
-                        lines=3,
-                        value="A festive red and gold promotional poster for Chinese New Year. In the center, a 3D rendered gift box is open, revealing cosmetics products. Large, bold golden Chinese characters at the top read '新春大促' (New Year Sale). Below it, smaller white text says '全场五折起' (50% off all items). The background features traditional cloud patterns and lanterns. A button at the bottom reads '立即抢购' (Buy Now).",
-                    )
-                    t2i_btn = gr.Button("Generate", variant="primary", size="lg")
-                    with gr.Accordion("Advanced Settings", open=False):
-                        t2i_steps = gr.Slider(10, 100, value=50, step=1, label="Steps")
-                        t2i_cfg = gr.Slider(1.0, 20.0, value=10.0, step=0.5, label="CFG Scale")
-                        t2i_seed = gr.Number(value=42, label="Seed", precision=0)
-                        t2i_res = gr.Slider(256, 1024, value=1024, step=128, label="Resolution")
-
-                with gr.Column(scale=1):
-                    t2i_output = gr.Image(label="Generated Image", type="pil")
-
-            t2i_btn.click(
-                fn=generate_image,
-                inputs=[t2i_prompt, t2i_steps, t2i_cfg, t2i_seed, t2i_res],
-                outputs=t2i_output,
-            )
-
-            gr.Examples(
-                examples=[
-                    ["An expansive field, blanketed by the soft light of morning, cradles a collection of eight cabbages, their green heads round and plump. These vegetables are nestled among rows of rich soil, dotted with glistening droplets of dew that cling to their crinkled leaves. As wisps of mist begin to lift, the cabbages lie poised, ready for the day's impending harvest."],
-                    ["A photo of a press conference with microphones and a blurred crowd in the background, realistic news style"],
-                    ["A screenshot of a modern landing page for a tech startup, clean UI, blue and white color scheme"],
-                    ["A candid street photography shot of people walking in busy Tokyo, daytime, realistic lighting"],
-                ],
-                inputs=t2i_prompt,
-                fn=generate_image,
-                cache_examples=True,
-                cache_mode="lazy",
-                outputs=t2i_output,
-            )
-
-        # ─── Image Editing Tab ────────────────────────────────────────────
-        with gr.Tab("🖼️ Image Editing"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    edit_input = gr.Image(label="Input Image", type="pil")
-                    edit_instruction = gr.Textbox(
-                        label="Edit Instruction",
-                        placeholder="Describe the edit you want to apply...",
-                        lines=2,
-                        value="Change the background to a forest.",
-                    )
-                    edit_btn = gr.Button("Apply Edit", variant="primary", size="lg")
-                    with gr.Accordion("Advanced Settings", open=False):
-                        edit_steps = gr.Slider(10, 100, value=50, step=1, label="Steps")
-                        edit_cfg = gr.Slider(1.0, 20.0, value=4.0, step=0.5, label="CFG Scale")
-                        edit_seed = gr.Number(value=42, label="Seed", precision=0)
-                        edit_res = gr.Slider(256, 1024, value=1024, step=128, label="Resolution")
-
-                with gr.Column(scale=1):
-                    edit_output = gr.Image(label="Edited Image", type="pil")
-
-            edit_btn.click(
-                fn=edit_image,
-                inputs=[edit_input, edit_instruction, edit_steps, edit_cfg, edit_seed, edit_res],
-                outputs=edit_output,
-            )
-
-            gr.Examples(
-                examples=[
-                    ["examples/background_change-input.png", "Change the background to a forest."],
-                    ["examples/color_alter-input.png", "Alter the color of the mirror frame to orange."],
-                    ["examples/material_alter-input.png", "Change the cup in hand to ceramic."],
-                    ["examples/motion_change-input.png", "Make the expression more sorrowful."],
-                ],
-                inputs=[edit_input, edit_instruction],
-                fn=edit_image,
-                cache_examples=True,
-                cache_mode="lazy",
-                outputs=edit_output,
-            )
-
-    gr.Markdown(
-        """
-        ---
-        **Model**: [yjb6/UniSpace](https://huggingface.co/yjb6/UniSpace) |
-        **Paper**: [2608.08676](https://arxiv.org/abs/2608.08676) |
-        **Code**: [GitHub](https://github.com/yjb6/UniSpace) |
-        **Project Page**: [yjb6.github.io/UniSpace](https://yjb6.github.io/UniSpace/)
-        """
+    e_run.click(
+        fn=edit,
+        inputs=[e_image, e_prompt, e_steps, e_cfg, e_seed, e_rand],
+        outputs=[e_out, e_used_seed, e_info],
+        api_name="edit",
     )
 
-demo.launch(mcp_server=True)
+if __name__ == "__main__":
+    demo.queue().launch(mcp_server=True)

patch-reparameterization/configs/decoder/ViTXL_nockpt/config.json

@@ -21,7 +21,7 @@
   "num_attention_heads": 12,
   "num_channels": 3,
   "num_hidden_layers": 12,
-  "patch_size": 16,
+  "patch_size": "SHOULD BE RELOADED",
   "qkv_bias": true,
   "torch_dtype": "float32",
   "transformers_version": "4.42.3",

patch-reparameterization/configs/release/pr-dinov2-imagenet256-generation.yaml

@@ -0,0 +1,55 @@
+# Release generation config for PR-DINOv2.
+stage_1:
+  target: stage1.PatchReparam
+  params:
+    encoder_cls: DINOv2Unified
+    encoder_config_path: ${oc.env:MODEL_ROOT}/huggingface.co/facebook/dinov2-with-registers-base
+    encoder_input_size: 224
+    encoder_params:
+      model_name: ${oc.env:MODEL_ROOT}/huggingface.co/facebook/dinov2-with-registers-base
+      recon_proj_dim: 128
+    decoder_config_path: configs/decoder/ViTXL_nockpt
+    noise_tau: 0.8
+    decode_noise: true
+    reshape_to_2d: true
+    normalization_stat_path: ${oc.env:PR_STATS_ROOT}/pr-dinov2-normalization-stats.pt
+    norm_feature: merged_tokens
+  checkpoint:
+    path: ${oc.env:PR_CHECKPOINT_ROOT}/pr-dinov2-tokenizer.pt
+    use_ema: true
+stage_2:
+  target: stage2.models.DDT.DiTwDDTHead
+  params:
+    input_size: 16
+    patch_size: 1
+    in_channels: 896
+    hidden_size: [1152, 2048]
+    depth: [28, 2]
+    num_heads: [16, 16]
+    mlp_ratio: 4.0
+    class_dropout_prob: 0.1
+    num_classes: 1000
+    use_qknorm: false
+    use_swiglu: true
+    use_rope: true
+    use_rmsnorm: true
+    wo_shift: false
+    use_pos_embed: true
+  ckpt: ${oc.env:PR_CHECKPOINT_ROOT}/pr-dinov2-dit.pt
+transport:
+  params: {path_type: Linear, prediction: velocity, loss_weight: null, time_dist_type: logit-normal_0_1}
+sampler:
+  mode: ODE
+  params: {sampling_method: euler, num_steps: 50, atol: 1.0e-6, rtol: 1.0e-3, reverse: false}
+guidance: {method: cfg, scale: 1.0, t_min: 0.0, t_max: 1.0}
+misc:
+  save_latents: false
+  latent_size: [896, 16, 16]
+  num_classes: 1000
+  time_dist_shift_dim: 229376
+  time_dist_shift_base: 4096
+training:
+  encoder_output_key: merged_tokens_norm
+  sem_dim: 768
+  recon_loss_weight: 18.0
+  normalize_channel_weight: true

patch-reparameterization/configs/release/pr-dinov2-imagenet256.yaml

@@ -0,0 +1,66 @@
+# Release evaluation config for PR-DINOv2.
+# Set MODEL_ROOT, PR_CHECKPOINT_ROOT, and PR_STATS_ROOT before running.
+stage_1:
+  target: stage1.PatchReparam
+  params:
+    encoder_cls: DINOv2Unified
+    encoder_config_path: ${oc.env:MODEL_ROOT}/huggingface.co/facebook/dinov2-with-registers-base
+    encoder_input_size: 224
+    encoder_params:
+      model_name: ${oc.env:MODEL_ROOT}/huggingface.co/facebook/dinov2-with-registers-base
+      recon_proj_dim: 128
+    decoder_config_path: configs/decoder/ViTXL_nockpt
+    noise_tau: 0.8
+    decode_noise: false
+    reshape_to_2d: true
+    normalization_stat_path: ${oc.env:PR_STATS_ROOT}/pr-dinov2-normalization-stats.pt
+  checkpoint:
+    path: ${oc.env:PR_CHECKPOINT_ROOT}/pr-dinov2-tokenizer.pt
+    use_ema: true
+
+stage_2:
+  target: stage2.models.DDT.DiTwDDTHead
+  params:
+    input_size: 16
+    patch_size: 1
+    in_channels: 896
+    hidden_size: [1152, 2048]
+    depth: [28, 2]
+    num_heads: [16, 16]
+    mlp_ratio: 4.0
+    class_dropout_prob: 0.1
+    num_classes: 1000
+    use_qknorm: false
+    use_swiglu: true
+    use_rope: true
+    use_rmsnorm: true
+    wo_shift: false
+    use_pos_embed: true
+  ckpt: ${oc.env:PR_CHECKPOINT_ROOT}/pr-dinov2-dit.pt
+
+transport:
+  params:
+    path_type: Linear
+    prediction: velocity
+    loss_weight: null
+    time_dist_type: logit-normal_0_1
+sampler:
+  mode: ODE
+  params: {sampling_method: euler, num_steps: 50, atol: 1.0e-6, rtol: 1.0e-3, reverse: false}
+guidance: {method: cfg, scale: 1.0, t_min: 0.0, t_max: 1.0}
+misc:
+  save_latents: false
+  latent_size: [896, 16, 16]
+  num_classes: 1000
+  time_dist_shift_dim: 229376
+  time_dist_shift_base: 4096
+training:
+  encoder_output_key: merged_tokens_norm
+  sem_dim: 768
+  recon_loss_weight: 18.0
+  normalize_channel_weight: true
+eval:
+  data_path: ${oc.env:IMAGENET_ROOT}/val
+  reference_npz_path: ${oc.env:IMAGENET_VAL_NPZ}
+  metrics: [psnr, ssim, rfid]
+  eval_weights: ema_only

patch-reparameterization/configs/release/pr-qwen-vit-imagenet256.yaml

@@ -0,0 +1,26 @@
+# Release reconstruction config for PR-Qwen-ViT (Table 3).
+# Set MODEL_ROOT, PR_CHECKPOINT_ROOT, and PR_STATS_ROOT before running.
+stage_1:
+  target: stage1.PatchReparam
+  params:
+    encoder_cls: Qwen3Unified
+    encoder_config_path: ${oc.env:QWEN_MODEL_ROOT}/Qwen/Qwen3-VL-8B-Instruct
+    encoder_input_size: 256
+    encoder_params:
+      model_name: ${oc.env:QWEN_MODEL_ROOT}/Qwen/Qwen3-VL-8B-Instruct
+      recon_proj_dim: 128
+      pixel_format: c_major
+    decoder_config_path: configs/decoder/ViTXL_nockpt
+    decoder_cls: ViTDecoder
+    noise_tau: 0.8
+    reshape_to_2d: true
+    normalization_stat_path: ${oc.env:PR_STATS_ROOT}/pr-qwen-vit-normalization-stats.pt
+  checkpoint:
+    path: ${oc.env:PR_CHECKPOINT_ROOT}/pr-qwen-vit-tokenizer.pt
+    use_ema: true
+
+eval:
+  data_path: ${oc.env:IMAGENET_ROOT}/val
+  reference_npz_path: ${oc.env:IMAGENET_VAL_NPZ}
+  metrics: [psnr, ssim, rfid]
+  eval_weights: ema_only

patch-reparameterization/configs/release/pr-siglip2-imagenet256-generation.yaml

@@ -0,0 +1,55 @@
+# Release generation config for PR-SigLIP2.
+stage_1:
+  target: stage1.PatchReparam
+  params:
+    encoder_cls: SigLIP2Unified
+    encoder_config_path: ${oc.env:MODEL_ROOT}/siglip2-base-patch16-256/main
+    encoder_input_size: 256
+    encoder_params:
+      model_name: ${oc.env:MODEL_ROOT}/siglip2-base-patch16-256/main
+      recon_proj_dim: 128
+    decoder_config_path: configs/decoder/ViTXL_nockpt
+    noise_tau: 0.8
+    decode_noise: true
+    reshape_to_2d: true
+    normalization_stat_path: ${oc.env:PR_STATS_ROOT}/pr-siglip2-normalization-stats.pt
+    norm_feature: merged_tokens
+  checkpoint:
+    path: ${oc.env:PR_CHECKPOINT_ROOT}/pr-siglip2-tokenizer.pt
+    use_ema: true
+stage_2:
+  target: stage2.models.DDT.DiTwDDTHead
+  params:
+    input_size: 16
+    patch_size: 1
+    in_channels: 896
+    hidden_size: [1152, 2048]
+    depth: [28, 2]
+    num_heads: [16, 16]
+    mlp_ratio: 4.0
+    class_dropout_prob: 0.1
+    num_classes: 1000
+    use_qknorm: false
+    use_swiglu: true
+    use_rope: true
+    use_rmsnorm: true
+    wo_shift: false
+    use_pos_embed: true
+  ckpt: ${oc.env:PR_CHECKPOINT_ROOT}/pr-siglip2-dit.pt
+transport:
+  params: {path_type: Linear, prediction: velocity, loss_weight: null, time_dist_type: logit-normal_0_1}
+sampler:
+  mode: ODE
+  params: {sampling_method: euler, num_steps: 50, atol: 1.0e-6, rtol: 1.0e-3, reverse: false}
+guidance: {method: cfg, scale: 1.0, t_min: 0.0, t_max: 1.0}
+misc:
+  save_latents: false
+  latent_size: [896, 16, 16]
+  num_classes: 1000
+  time_dist_shift_dim: 229376
+  time_dist_shift_base: 4096
+training:
+  encoder_output_key: merged_tokens_norm
+  sem_dim: 768
+  recon_loss_weight: 18.0
+  normalize_channel_weight: true

patch-reparameterization/configs/release/pr-siglip2-imagenet256.yaml

@@ -0,0 +1,73 @@
+# Release evaluation config for PR-SigLIP2.
+# Set MODEL_ROOT, PR_CHECKPOINT_ROOT, and PR_STATS_ROOT before running.
+stage_1:
+  target: stage1.PatchReparam
+  params:
+    encoder_cls: SigLIP2wTwoBackbone
+    encoder_config_path: ${oc.env:MODEL_ROOT}/siglip2-base-patch16-256/main
+    encoder_input_size: 256
+    encoder_params:
+      sementic_model_name: ${oc.env:MODEL_ROOT}/siglip2-base-patch16-256/main
+      recon_model_name: ${oc.env:MODEL_ROOT}/siglip2-base-patch16-256/main
+      num_tokens: 256
+      recon_post_norm_type: no_norm
+      sementic_post_norm_type: no_norm
+      merger_cls: ConcatMerger
+      merger_kwargs:
+        recon_proj_dim: 128
+        recon_proj_cls: LinearAdapter
+        recon_proj_kwargs: {}
+    decoder_config_path: configs/decoder/ViTXL_nockpt
+    noise_tau: 0.8
+    decode_noise: false
+    reshape_to_2d: true
+  checkpoint:
+    path: ${oc.env:PR_SIGLIP2_TOKENIZER_CKPT}
+    use_ema: true
+
+stage_2:
+  target: stage2.models.DDT.DiTwDDTHead
+  params:
+    input_size: 16
+    patch_size: 1
+    in_channels: 896
+    hidden_size: [1152, 2048]
+    depth: [28, 2]
+    num_heads: [16, 16]
+    mlp_ratio: 4.0
+    class_dropout_prob: 0.1
+    num_classes: 1000
+    use_qknorm: false
+    use_swiglu: true
+    use_rope: true
+    use_rmsnorm: true
+    wo_shift: false
+    use_pos_embed: true
+  ckpt: ${oc.env:PR_CHECKPOINT_ROOT}/pr-siglip2-dit.pt
+
+transport:
+  params:
+    path_type: Linear
+    prediction: velocity
+    loss_weight: null
+    time_dist_type: logit-normal_0_1
+sampler:
+  mode: ODE
+  params: {sampling_method: euler, num_steps: 50, atol: 1.0e-6, rtol: 1.0e-3, reverse: false}
+guidance: {method: cfg, scale: 1.0, t_min: 0.0, t_max: 1.0}
+misc:
+  save_latents: false
+  latent_size: [896, 16, 16]
+  num_classes: 1000
+  time_dist_shift_dim: 229376
+  time_dist_shift_base: 4096
+training:
+  encoder_output_key: merged_tokens_norm
+  sem_dim: 768
+  recon_loss_weight: 18.0
+  normalize_channel_weight: true
+eval:
+  data_path: ${oc.env:IMAGENET_ROOT}/val
+  reference_npz_path: ${oc.env:IMAGENET_VAL_NPZ}
+  metrics: [psnr, ssim, rfid]
+  eval_weights: ema_only

patch-reparameterization/src/eval/__init__.py

@@ -0,0 +1,580 @@
+from .ref_iqa import calculate_psnr, calculate_lpips, calculate_ssim
+from .fid import calculate_rfid, calculate_gfid
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torch.cuda.amp import autocast
+from typing import Union, Dict, Optional
+try:
+    from accelerate import Accelerator
+except ImportError:
+    Accelerator = None
+from torch.utils.data import DataLoader, Subset
+from tqdm import tqdm
+import os
+import sys
+from functools import lru_cache
+@lru_cache(maxsize=4)
+def _load_encoder_cached(model_path: str, device_str: str):
+    """
+    懒加载并缓存 encoder（按 model_path + device 缓存，进程内只加载一次）。
+    使用 transformers.AutoModel + AutoProcessor。
+    """
+    from transformers import AutoModel, AutoProcessor
+    print(f"[CosSim] Loading encoder from {model_path} onto {device_str} (cached)...")
+    processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
+    model = model.to(device_str).eval()
+    return model, processor
+
+
+@torch.no_grad()
+def calculate_encoder_cos_sim(
+    ref_arr: np.ndarray,
+    rec_arr: np.ndarray,
+    encoder_model: str,
+    device: torch.device,
+    batch_size: int = 64,
+) -> float:
+    """
+    用指定 encoder 对 GT 图像和重建图像编码，计算 image embedding 的 cos 相似度均值。
+
+    Args:
+        ref_arr: GT 图像 [N, H, W, C] uint8
+        rec_arr: 重建图像 [N, H, W, C] uint8
+        encoder_model: HuggingFace 模型路径或名称
+        device: 计算设备
+        batch_size: 每批编码的图像数
+
+    Returns:
+        所有样本的 cos 相似度均值 (float)
+    """
+    device_str = "cuda" if device.type == "cuda" else "cpu"
+    model, processor = _load_encoder_cached(encoder_model, device_str)
+
+    N = len(ref_arr)
+    cos_sims = []
+
+    for start in range(0, N, batch_size):
+        end = min(start + batch_size, N)
+        ref_batch = [Image.fromarray(img) for img in ref_arr[start:end]]
+        rec_batch = [Image.fromarray(img) for img in rec_arr[start:end]]
+
+        # 预处理：使用 processor 处理两组图像
+        ref_inputs = processor(images=ref_batch, return_tensors="pt").to(device_str)
+        rec_inputs = processor(images=rec_batch, return_tensors="pt").to(device_str)
+
+        # 提取 image embedding
+        ref_feats = model.get_image_features(**ref_inputs)  # [B, D]
+        rec_feats = model.get_image_features(**rec_inputs)  # [B, D]
+
+        # L2 归一化后计算 cos 相似度
+        ref_feats = F.normalize(ref_feats.float(), dim=-1)
+        rec_feats = F.normalize(rec_feats.float(), dim=-1)
+        batch_cos = (ref_feats * rec_feats).sum(dim=-1)  # [B]
+        cos_sims.append(batch_cos.cpu())
+
+    cos_sim_mean = torch.cat(cos_sims).mean().item()
+    return cos_sim_mean
+
+
+def compute_reconstruction_metrics(
+    ref_arr: np.ndarray,
+    rec_arr: np.ndarray,
+    device: torch.device,
+    batch_size: int = 128,
+    metrics_to_compute = ("psnr", "ssim", "rfid"),
+    disable_bar: bool = True,
+    encoder_model: Optional[str] = None,
+    psnr_resize: Optional[int] = None,
+    l1_resize: Optional[int] = None,
+) -> Dict[str, float]:
+    """
+    Compute reconstruction metrics between reference and reconstructed images.
+
+    Args:
+        ref_arr: Reference images [N, H, W, C] uint8
+        rec_arr: Reconstructed images [N, H, W, C] uint8
+        device: Device for computation
+        batch_size: Batch size for metric computation
+
+    Returns:
+        Dictionary with metrics: eval/psnr, eval/ssim, eval/rfid
+        Note: LPIPS is not computed here since it's already tracked during training
+    """
+    device_str = "cuda" if device.type == "cuda" else "cpu"
+    # For paired metrics (psnr, ssim, l1, cos_sim), align ref/rec count
+    n_paired = min(len(ref_arr), len(rec_arr))
+    ref_paired = ref_arr[:n_paired]
+    rec_paired = rec_arr[:n_paired]
+    results_dict = {}
+    if 'psnr' in metrics_to_compute:
+        psnr = calculate_psnr(ref_paired, rec_paired, batch_size, device_str, disable_bar=disable_bar, psnr_resize=psnr_resize)
+        results_dict["psnr"] = psnr
+    if 'ssim' in metrics_to_compute:
+        ssim = calculate_ssim(ref_paired, rec_paired, batch_size, device_str, disable_bar=disable_bar)
+        results_dict["ssim"] = ssim
+    if 'rfid' in metrics_to_compute:
+        rfid = calculate_rfid(ref_arr, rec_arr, batch_size, device_str)
+        results_dict["rfid"] = rfid
+    if 'l1' in metrics_to_compute:
+        # L1 = mean(|rec - ref|), same as train: (recon - images).abs().mean()
+        ref_f = ref_paired.astype(np.float32) / 255.0
+        rec_f = rec_paired.astype(np.float32) / 255.0
+        if l1_resize is not None:
+            # resize via torch interpolate to keep consistent with psnr_resize style
+            # Process in small batches to avoid CUDA OOM on large datasets
+            device_str = "cuda" if device.type == "cuda" else "cpu"
+            resize_batch = 64
+            n = ref_f.shape[0]
+            l1_sum = 0.0
+            l1_count = 0
+            for i in range(0, n, resize_batch):
+                ref_b = torch.from_numpy(ref_f[i:i+resize_batch]).permute(0, 3, 1, 2).to(device_str)
+                rec_b = torch.from_numpy(rec_f[i:i+resize_batch]).permute(0, 3, 1, 2).to(device_str)
+                ref_b = F.interpolate(ref_b, size=(l1_resize, l1_resize), mode="bilinear", align_corners=False)
+                rec_b = F.interpolate(rec_b, size=(l1_resize, l1_resize), mode="bilinear", align_corners=False)
+                l1_sum += (ref_b - rec_b).abs().sum().item()
+                l1_count += ref_b.numel()
+            l1 = l1_sum / l1_count
+        else:
+            l1 = np.abs(rec_f - ref_f).mean()
+        results_dict["l1"] = float(l1)
+    if 'cos_sim' in metrics_to_compute:
+        assert encoder_model is not None, \
+            "'cos_sim' metric requires encoder_model (HuggingFace model path) to be specified."
+        print(f"[CosSim] Computing encoder cos_sim with model: {encoder_model}")
+        cos_sim = calculate_encoder_cos_sim(
+            ref_paired, rec_paired,
+            encoder_model=encoder_model,
+            device=device,
+            batch_size=batch_size,
+        )
+        results_dict["cos_sim"] = cos_sim
+    assert len(results_dict) > 0, "No metrics were computed."
+    return results_dict
+def compute_generation_metrics(
+    ref_arr: np.ndarray,
+    rec_arr: np.ndarray,
+    device: torch.device,
+    batch_size: int = 128,
+):
+    device_str = "cuda" if device.type == "cuda" else "cpu"
+    # only eval FID
+    fid = calculate_gfid(rec_arr, ref_arr, batch_size, device_str)
+    return {
+        'fid': fid
+    }
+@torch.no_grad()
+def evaluate_generation_distributed(
+    model_fn,
+    sample_fn,
+    latent_size, # for noise
+    additional_model_kwargs,
+    use_guidance: bool,
+    pr,
+    val_dataset,
+    num_samples: int,
+    batch_size: int,
+    rank: Union[int, None] = None,
+    world_size: Union[int, None] = None,
+    device: torch.device = None,
+    experiment_dir: str = None,
+    global_step: int = None,
+    autocast_kwargs: dict = None,
+    metric_batch_size: int = 128,
+    reference_npz_path: Optional[str] = None,
+    accelerator: Union[Accelerator, None] = None,
+) -> Optional[Dict[str, float]]:
+    """
+    Evaluate generation metrics using all GPUs in a distributed manner.
+
+    Args:
+        val_dataset: Validation dataset
+        batch_size: Batch size per GPU for generation
+        rank: Current GPU rank (deprecated, use accelerator instead)
+        world_size: Total number of GPUs (deprecated, use accelerator instead)
+        device: Device to use (deprecated, use accelerator instead)
+        experiment_dir: Experiment directory
+        global_step: Current training step
+        autocast_kwargs: Autocast configuration
+        metric_batch_size: Batch size for metric computation (on rank 0)
+        reference_npz_path: Optional path to existing reference NPZ file
+        accelerator: Accelerator instance (preferred over rank/world_size/device)
+
+    Returns:
+        Dictionary of metrics (only on rank 0, None on other ranks)
+    """
+    # Support both Accelerator and legacy DDP interface
+    if accelerator is not None:
+        rank = accelerator.process_index
+        world_size = accelerator.num_processes
+        device = accelerator.device
+        is_main_process = accelerator.is_main_process
+        wait_for_everyone = accelerator.wait_for_everyone
+    else:
+        # Legacy DDP mode
+        import torch.distributed as dist
+        if rank is None or world_size is None or device is None:
+            raise ValueError("Either accelerator must be provided, or rank/world_size/device must be provided")
+        is_main_process = (rank == 0)
+        wait_for_everyone = dist.barrier
+
+    # model.eval()
+    # Save shard NPZ
+    temp_dir = os.path.join(experiment_dir, "eval_npzs")
+    if is_main_process:
+        print(f"\n[Eval] Starting distributed sampling evaluation at step {global_step}")
+        os.makedirs(temp_dir, exist_ok=True)
+
+    # Wait for rank 0 to create the directory before other ranks try to save
+    wait_for_everyone()
+    # print(f"[Rank {rank}] Starting sampling...")
+    # Each rank processes its shard
+    N = min(len(val_dataset), num_samples)
+    chunk = N // world_size
+
+    if rank < world_size - 1:
+        start = rank * chunk
+        end   = (rank + 1) * chunk
+    else:
+        # Last rank takes the remainder (and handles N < world_size gracefully)
+        start = rank * chunk
+        end   = N
+
+    rank_indices = list(range(start, end))
+    subset = Subset(val_dataset, rank_indices)
+    loader = DataLoader(
+        subset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    # Reconstruct images on this rank
+    generations = []
+    iterator = tqdm(loader, desc=f"[Rank {rank}] Sampling", file=sys.stdout) if is_main_process else loader
+
+    with torch.inference_mode():
+        for _, label in iterator: # don't actually need images at sampling time
+            n = label.size(0)
+            z = torch.randn(n, *latent_size, device = device)
+            y = label.to(device)
+            if use_guidance:
+                z = torch.cat([z, z], dim=0)
+                y_null = torch.full((n,), null_label, device=device)
+                y = torch.cat([y, y_null], dim=0)
+            model_kwargs = dict(y=y, **additional_model_kwargs)
+            with autocast(**autocast_kwargs):
+                samples = sample_fn(z, model_fn, **model_kwargs)[-1]
+                if use_guidance:
+                    samples = samples.chunk(2, dim = 0)
+                if callable(pr) and not hasattr(pr, 'decode'):
+                    samples = pr(samples)
+                else:
+                    samples = pr.decode(samples)
+                if hasattr(samples, 'x_rec'):
+                    samples = samples.x_rec
+                samples = samples.clamp(0, 1)
+            gen_np = samples.mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+            for img in gen_np:
+                generations.append(img)
+
+    generations = np.stack(generations)
+    shard_path = os.path.join(temp_dir, f"gen_{global_step:07d}_{rank:02d}.npz")
+    np.savez(shard_path, arr_0=generations)
+
+    if is_main_process:
+        print(f"[Rank {rank}] Saved {len(generations)} generation to {shard_path}")
+
+    # Wait for all ranks to finish generation
+    wait_for_everyone()
+
+    # Rank 0 computes metrics
+    metrics = None
+    if is_main_process:
+        # Combine all generation shards
+        all_gens = []
+        for r in range(world_size):
+            shard_file = os.path.join(temp_dir, f"gen_{global_step:07d}_{r:02d}.npz")
+            shard_data = np.load(shard_file)["arr_0"]
+            all_gens.append(shard_data)
+
+        combined_recons = np.concatenate(all_gens, axis=0)[:num_samples]
+        print(f"[Eval] Combined generation NPZ shape: {combined_recons.shape}")
+
+        # Load reference NPZ
+        ref_npz_path = reference_npz_path
+
+        if not os.path.exists(ref_npz_path):
+            raise FileNotFoundError(f"Reference NPZ not found at {ref_npz_path}")
+
+        ref_stats = np.load(ref_npz_path)
+        print(f"[Eval] Loaded reference NPZ from {ref_npz_path}")
+
+        # Compute metrics
+        print("[Eval] Computing metrics...")
+        metrics = compute_generation_metrics(
+            ref_stats,
+            combined_recons,
+            device,
+            metric_batch_size,
+        )
+
+        # Print results
+        print(f"[Eval] Step {global_step} Metrics:")
+        for key, value in metrics.items():
+            print(f"  {key}: {value:.6f}")
+
+        # Cleanup reconstruction shards
+        for r in range(world_size):
+            shard_file = os.path.join(temp_dir, f"gen_{global_step:07d}_{r:02d}.npz")
+            if os.path.exists(shard_file):
+                os.remove(shard_file)
+
+    wait_for_everyone()
+    return metrics
+@torch.no_grad()
+def evaluate_reconstruction_distributed(
+    model,
+    val_dataset,
+    num_samples: int,
+    batch_size: int,
+    rank: Union[int, None] = None,
+    world_size: Union[int, None] = None,
+    device: torch.device = None,
+    experiment_dir: str = None,
+    global_step: int = None,
+    autocast_kwargs: dict = None,
+    metric_batch_size: int = 128,
+    reference_npz_path: Optional[str] = None,
+    metrics_to_compute: Optional[list] = ("psnr", "ssim", "rfid"),
+    accelerator: Union[Accelerator, None] = None,
+    save_recon_dir: Optional[str] = None,
+    image_features_mode: Optional[str] = None,
+    encoder_cos_sim_model: Optional[str] = None,
+    psnr_resize: Optional[int] = None,
+    l1_resize: Optional[int] = None,
+) -> Optional[Dict[str, float]]:
+    """
+    Evaluate reconstruction metrics using all GPUs in a distributed manner.
+
+    Args:
+        model: Model to evaluate (should be in eval mode)
+        val_dataset: Validation dataset
+        batch_size: Batch size per GPU for reconstruction
+        rank: Current GPU rank (deprecated, use accelerator instead)
+        world_size: Total number of GPUs (deprecated, use accelerator instead)
+        device: Device to use (deprecated, use accelerator instead)
+        experiment_dir: Experiment directory
+        global_step: Current training step
+        autocast_kwargs: Autocast configuration
+        metric_batch_size: Batch size for metric computation (on rank 0)
+        reference_npz_path: Optional path to existing reference NPZ file
+        accelerator: Accelerator instance (preferred over rank/world_size/device)
+
+    Returns:
+        Dictionary of metrics (only on rank 0, None on other ranks)
+    """
+    # Support both Accelerator and legacy DDP interface
+    if accelerator is not None:
+        rank = accelerator.process_index
+        world_size = accelerator.num_processes
+        device = accelerator.device
+        is_main_process = accelerator.is_main_process
+        wait_for_everyone = accelerator.wait_for_everyone
+    else:
+        # Legacy DDP mode
+        import torch.distributed as dist
+        if rank is None or world_size is None or device is None:
+            raise ValueError("Either accelerator must be provided, or rank/world_size/device must be provided")
+        is_main_process = (rank == 0)
+        wait_for_everyone = dist.barrier
+
+    model.eval()
+    # Save shard NPZ
+    temp_dir = os.path.join(experiment_dir, "eval_npzs")
+    if is_main_process:
+        print(f"\n[Eval] Starting distributed reconstruction evaluation at step {global_step}")
+        os.makedirs(temp_dir, exist_ok=True)
+    # Wait for rank 0 to create the directory before other ranks try to save
+    wait_for_everyone()
+    # print(f"[Rank {rank}] Starting reconstruction...")
+    # Each rank processes its shard
+    N = min(len(val_dataset), num_samples)
+    chunk = N // world_size
+
+    if rank < world_size - 1:
+        start = rank * chunk
+        end   = (rank + 1) * chunk
+    else:
+        # Last rank takes the remainder (and handles N < world_size gracefully)
+        start = rank * chunk
+        end   = N
+
+    rank_indices = list(range(start, end))
+    subset = Subset(val_dataset, rank_indices)
+    loader = DataLoader(
+        subset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    # Reconstruct images on this rank, keep originals for distributed metrics
+    reconstructions = []
+    originals = []
+    iterator = tqdm(loader, desc=f"[Rank {rank}] Reconstructing", file=sys.stdout) if is_main_process else loader
+
+    with torch.inference_mode():
+        for images, _ in iterator:
+            images = images.to(device, non_blocking=True)
+            # Save originals before model forward
+            orig_np = images.clamp(0, 1).mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+            with autocast(**autocast_kwargs):
+                if image_features_mode is not None:
+                    recon = model(images, image_features_mode=image_features_mode)
+                else:
+                    recon = model(images)
+
+            # Handle DecodeOutput from split_adaln decoder
+            if hasattr(recon, 'x_rec'):
+                recon = recon.x_rec
+            # Convert to numpy uint8 [H, W, C]
+            recon = recon.clamp(0, 1)
+            recon_np = recon.mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+
+            for img in recon_np:
+                reconstructions.append(img)
+            for img in orig_np:
+                originals.append(img)
+
+    reconstructions = np.stack(reconstructions)
+    originals = np.stack(originals)
+
+    # Save shard for backward compat
+    shard_path = os.path.join(temp_dir, f"recon_{global_step:07d}_{rank:02d}.npz")
+    np.savez(shard_path, arr_0=reconstructions)
+    if is_main_process:
+        print(f"[Rank {rank}] Saved {len(reconstructions)} reconstructions to {shard_path}")
+
+    # Wait for all ranks to finish reconstruction
+    wait_for_everyone()
+
+    # ── Distributed metrics: each rank computes on its own subset ──────────
+    import torch.distributed as dist
+    from .fid import compute_inception_features, calculate_rfid_from_features
+    local_n = len(reconstructions)
+
+    # PSNR / SSIM: compute locally then all_reduce weighted sum
+    scalar_metrics_to_compute = [m for m in (metrics_to_compute or []) if m in ('psnr', 'ssim', 'lpips')]
+    local_scalar = {}
+    if scalar_metrics_to_compute:
+        if is_main_process:
+            print(f"[Eval] Rank {rank}: computing {scalar_metrics_to_compute} on {local_n} images...")
+        local_scalar = compute_reconstruction_metrics(
+            originals, reconstructions, device, metric_batch_size,
+            metrics_to_compute=scalar_metrics_to_compute,
+            disable_bar=is_main_process,
+            encoder_model=encoder_cos_sim_model,
+            psnr_resize=psnr_resize,
+            l1_resize=l1_resize,
+        )
+
+    # Weighted all_reduce: sum(metric * count) / total_count
+    global_scalar = {}
+    count_t = torch.tensor(float(local_n), device=device)
+    distributed = world_size > 1 and dist.is_available() and dist.is_initialized()
+    if distributed:
+        dist.all_reduce(count_t, op=dist.ReduceOp.SUM)
+    for k, v in local_scalar.items():
+        val_t = torch.tensor(v * local_n, dtype=torch.float64, device=device)
+        if distributed:
+            dist.all_reduce(val_t, op=dist.ReduceOp.SUM)
+        global_scalar[k] = (val_t / count_t).item()
+
+    # rFID: each rank computes inception features, all_gather, rank 0 computes FID
+    global_rfid = {}
+    if 'rfid' in (metrics_to_compute or []):
+        if is_main_process:
+            print(f"[Eval] Rank {rank}: computing inception features on {local_n} images...")
+        local_recon_feats = compute_inception_features(reconstructions, metric_batch_size, device)
+        local_ref_feats   = compute_inception_features(originals,       metric_batch_size, device)
+
+        recon_t = torch.from_numpy(local_recon_feats).to(device)
+        ref_t   = torch.from_numpy(local_ref_feats).to(device)
+
+        if distributed:
+            all_recon = [torch.zeros_like(recon_t) for _ in range(world_size)]
+            all_ref   = [torch.zeros_like(ref_t)   for _ in range(world_size)]
+            dist.all_gather(all_recon, recon_t)
+            dist.all_gather(all_ref, ref_t)
+        else:
+            all_recon = [recon_t]
+            all_ref = [ref_t]
+
+        if is_main_process:
+            combined_recon_feats = torch.cat(all_recon, dim=0).cpu().double().numpy()
+            combined_ref_feats   = torch.cat(all_ref,   dim=0).cpu().double().numpy()
+            rfid = calculate_rfid_from_features(combined_recon_feats, combined_ref_feats)
+            global_rfid = {'rfid': rfid}
+
+    metrics = None
+    if is_main_process:
+        metrics = {**global_scalar, **global_rfid}
+        # Persist the all-reduced count so release validation proves how many
+        # examples were actually processed rather than trusting CLI intent.
+        metrics["_evaluated_samples"] = int(count_t.item())
+        print(f"[Eval] Step {global_step} Metrics:")
+        for key, value in metrics.items():
+            print(f"  {key}: {value:.6f}")
+
+        # Optionally save reconstructed images
+        if save_recon_dir:
+            import threading
+            recons_copy = np.copy(reconstructions)
+            save_dir = save_recon_dir
+            def _save_recon_images():
+                os.makedirs(save_dir, exist_ok=True)
+                for i, img in enumerate(recons_copy):
+                    Image.fromarray(img).save(os.path.join(save_dir, f"recon_{i:06d}.png"))
+                print(f"[Eval] Saved {len(recons_copy)} reconstruction images to {save_dir}")
+            threading.Thread(target=_save_recon_images, daemon=False).start()
+
+        # Cleanup shards
+        for r in range(world_size):
+            shard_file = os.path.join(temp_dir, f"recon_{global_step:07d}_{r:02d}.npz")
+            if os.path.exists(shard_file):
+                os.remove(shard_file)
+
+    wait_for_everyone()
+    model.train()
+
+    return metrics
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ref-img", type=str, default="samples/imagenet-256-val.npz")
+    parser.add_argument("--rec-img", type=str, default="samples/sdvae-ft-mse-f8d4.npz")
+    parser.add_argument("--bs", type=int, default=128)
+    args = parser.parse_args()
+
+    # Load images
+    device = "cuda"
+    ref_img = np.load(args.ref_img)["arr_0"]
+    rec_img = np.load(args.rec_img)["arr_0"]
+    print(f"Loaded images: ref: {ref_img.shape}, rec: {rec_img.shape}")
+
+    psnr = calculate_psnr(ref_img, rec_img, args.bs, device)
+    print(f"PSNR: {psnr:.6f}")
+    lpips = calculate_lpips(ref_img, rec_img, args.bs, device)
+    print(f"LPIPS: {lpips:.6f}")
+    ssim_val = calculate_ssim(ref_img, rec_img, args.bs, device)
+    print(f"SSIM: {ssim_val:.6f}")
+    rfid = calculate_rfid(ref_img, rec_img, args.bs, device)
+    print(f"rFID: {rfid:.6f}")

patch-reparameterization/src/eval/clip_t.py

@@ -0,0 +1,363 @@
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from transformers import SiglipModel, SiglipProcessor
+from tqdm import tqdm
+import os
+
+# =================配置部分=================
+# 本地模型路径 (包含 model.safetensors, config.json, preprocessor_config.json 等)
+# MODEL_PATH = HF_HOME 下 siglip2-base-patch16-256
+# MODEL_PATH = CKPT_ROOT 下 siglip2/test/ep-last/
+MODEL_PATH = os.environ.get("SIGLIP2_MODEL_PATH", "google/siglip2-base-patch16-256")
+# ImageNet 验证集根目录 (即包含 n01440764 等文件夹的上一级目录)
+IMAGENET_VAL_PATH = os.path.join(os.environ.get("CLIP_BENCH_ROOT", "/path/to/clip_benchmark"), "imagenet-1k", "val")
+
+BATCH_SIZE = 256
+NUM_WORKERS = 8
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# =================ImageNet 类别标签=================
+# 标准 ImageNet-1k 的 1000 个类别名称 (对应 WordNet ID 排序)
+# 为了代码简洁，这里列出前几个作为示例，完整列表请在代码下方查看说明或从外部加载
+# 这里我们使用一个简单的函数来获取标签，通常实际评测会使用 OpenAI 的 1000 个类名
+try:
+    # 尝试从网络加载标准标签，如果无法联网请手动替换 IMAGENET_CLASSES 列表
+    import requests
+    url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
+    IMAGENET_CLASSES = requests.get(url).text.strip().split('\n')
+    print(f"成功加载 {len(IMAGENET_CLASSES)} 个类别标签。")
+except Exception as e:
+    print("无法下载标签，请确保联网或手动填入 IMAGENET_CLASSES。")
+    # 示例: IMAGENET_CLASSES = ["tench", "goldfish", ...]
+    # raise e
+    IMAGENET_CLASSES = (
+    "tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", "electric ray",
+    "stingray", "rooster", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco",
+    "indigo bunting", "American robin", "bulbul", "jay", "magpie", "chickadee", "American dipper",
+    "kite (bird of prey)", "bald eagle", "vulture", "great grey owl", "fire salamander",
+    "smooth newt", "newt", "spotted salamander", "axolotl", "American bullfrog", "tree frog",
+    "tailed frog", "loggerhead sea turtle", "leatherback sea turtle", "mud turtle", "terrapin",
+    "box turtle", "banded gecko", "green iguana", "Carolina anole",
+    "desert grassland whiptail lizard", "agama", "frilled-necked lizard", "alligator lizard",
+    "Gila monster", "European green lizard", "chameleon", "Komodo dragon", "Nile crocodile",
+    "American alligator", "triceratops", "worm snake", "ring-necked snake",
+    "eastern hog-nosed snake", "smooth green snake", "kingsnake", "garter snake", "water snake",
+    "vine snake", "night snake", "boa constrictor", "African rock python", "Indian cobra",
+    "green mamba", "sea snake", "Saharan horned viper", "eastern diamondback rattlesnake",
+    "sidewinder rattlesnake", "trilobite", "harvestman", "scorpion", "yellow garden spider",
+    "barn spider", "European garden spider", "southern black widow", "tarantula", "wolf spider",
+    "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie grouse", "peafowl",
+    "quail", "partridge", "african grey parrot", "macaw", "sulphur-crested cockatoo", "lorikeet",
+    "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "duck",
+    "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby",
+    "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch",
+    "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab",
+    "fiddler crab", "red king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab",
+    "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron",
+    "great egret", "bittern bird", "crane bird", "limpkin", "common gallinule", "American coot",
+    "bustard", "ruddy turnstone", "dunlin", "common redshank", "dowitcher", "oystercatcher",
+    "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion",
+    "Chihuahua", "Japanese Chin", "Maltese", "Pekingese", "Shih Tzu", "King Charles Spaniel",
+    "Papillon", "toy terrier", "Rhodesian Ridgeback", "Afghan Hound", "Basset Hound", "Beagle",
+    "Bloodhound", "Bluetick Coonhound", "Black and Tan Coonhound", "Treeing Walker Coonhound",
+    "English foxhound", "Redbone Coonhound", "borzoi", "Irish Wolfhound", "Italian Greyhound",
+    "Whippet", "Ibizan Hound", "Norwegian Elkhound", "Otterhound", "Saluki", "Scottish Deerhound",
+    "Weimaraner", "Staffordshire Bull Terrier", "American Staffordshire Terrier",
+    "Bedlington Terrier", "Border Terrier", "Kerry Blue Terrier", "Irish Terrier",
+    "Norfolk Terrier", "Norwich Terrier", "Yorkshire Terrier", "Wire Fox Terrier",
+    "Lakeland Terrier", "Sealyham Terrier", "Airedale Terrier", "Cairn Terrier",
+    "Australian Terrier", "Dandie Dinmont Terrier", "Boston Terrier", "Miniature Schnauzer",
+    "Giant Schnauzer", "Standard Schnauzer", "Scottish Terrier", "Tibetan Terrier",
+    "Australian Silky Terrier", "Soft-coated Wheaten Terrier", "West Highland White Terrier",
+    "Lhasa Apso", "Flat-Coated Retriever", "Curly-coated Retriever", "Golden Retriever",
+    "Labrador Retriever", "Chesapeake Bay Retriever", "German Shorthaired Pointer", "Vizsla",
+    "English Setter", "Irish Setter", "Gordon Setter", "Brittany dog", "Clumber Spaniel",
+    "English Springer Spaniel", "Welsh Springer Spaniel", "Cocker Spaniel", "Sussex Spaniel",
+    "Irish Water Spaniel", "Kuvasz", "Schipperke", "Groenendael dog", "Malinois", "Briard",
+    "Australian Kelpie", "Komondor", "Old English Sheepdog", "Shetland Sheepdog", "collie",
+    "Border Collie", "Bouvier des Flandres dog", "Rottweiler", "German Shepherd Dog", "Dobermann",
+    "Miniature Pinscher", "Greater Swiss Mountain Dog", "Bernese Mountain Dog",
+    "Appenzeller Sennenhund", "Entlebucher Sennenhund", "Boxer", "Bullmastiff", "Tibetan Mastiff",
+    "French Bulldog", "Great Dane", "St. Bernard", "husky", "Alaskan Malamute", "Siberian Husky",
+    "Dalmatian", "Affenpinscher", "Basenji", "pug", "Leonberger", "Newfoundland dog",
+    "Great Pyrenees dog", "Samoyed", "Pomeranian", "Chow Chow", "Keeshond", "brussels griffon",
+    "Pembroke Welsh Corgi", "Cardigan Welsh Corgi", "Toy Poodle", "Miniature Poodle",
+    "Standard Poodle", "Mexican hairless dog (xoloitzcuintli)", "grey wolf", "Alaskan tundra wolf",
+    "red wolf or maned wolf", "coyote", "dingo", "dhole", "African wild dog", "hyena", "red fox",
+    "kit fox", "Arctic fox", "grey fox", "tabby cat", "tiger cat", "Persian cat", "Siamese cat",
+    "Egyptian Mau", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger",
+    "cheetah", "brown bear", "American black bear", "polar bear", "sloth bear", "mongoose",
+    "meerkat", "tiger beetle", "ladybug", "ground beetle", "longhorn beetle", "leaf beetle",
+    "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper",
+    "cricket insect", "stick insect", "cockroach", "praying mantis", "cicada", "leafhopper",
+    "lacewing", "dragonfly", "damselfly", "red admiral butterfly", "ringlet butterfly",
+    "monarch butterfly", "small white butterfly", "sulphur butterfly", "gossamer-winged butterfly",
+    "starfish", "sea urchin", "sea cucumber", "cottontail rabbit", "hare", "Angora rabbit",
+    "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "common sorrel horse",
+    "zebra", "pig", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison",
+    "ram (adult male sheep)", "bighorn sheep", "Alpine ibex", "hartebeest", "impala (antelope)",
+    "gazelle", "arabian camel", "llama", "weasel", "mink", "European polecat",
+    "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan",
+    "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas monkey", "baboon", "macaque",
+    "langur", "black-and-white colobus", "proboscis monkey", "marmoset", "white-headed capuchin",
+    "howler monkey", "titi monkey", "Geoffroy's spider monkey", "common squirrel monkey",
+    "ring-tailed lemur", "indri", "Asian elephant", "African bush elephant", "red panda",
+    "giant panda", "snoek fish", "eel", "silver salmon", "rock beauty fish", "clownfish",
+    "sturgeon", "gar fish", "lionfish", "pufferfish", "abacus", "abaya", "academic gown",
+    "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance",
+    "amphibious vehicle", "analog clock", "apiary", "apron", "trash can", "assault rifle",
+    "backpack", "bakery", "balance beam", "balloon", "ballpoint pen", "Band-Aid", "banjo",
+    "baluster / handrail", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel",
+    "wheelbarrow", "baseball", "basketball", "bassinet", "bassoon", "swimming cap", "bath towel",
+    "bathtub", "station wagon", "lighthouse", "beaker", "military hat (bearskin or shako)",
+    "beer bottle", "beer glass", "bell tower", "baby bib", "tandem bicycle", "bikini",
+    "ring binder", "binoculars", "birdhouse", "boathouse", "bobsleigh", "bolo tie", "poke bonnet",
+    "bookcase", "bookstore", "bottle cap", "hunting bow", "bow tie", "brass memorial plaque", "bra",
+    "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest",
+    "high-speed train", "butcher shop", "taxicab", "cauldron", "candle", "cannon", "canoe",
+    "can opener", "cardigan", "car mirror", "carousel", "tool kit", "cardboard box / carton",
+    "car wheel", "automated teller machine", "cassette", "cassette player", "castle", "catamaran",
+    "CD player", "cello", "mobile phone", "chain", "chain-link fence", "chain mail", "chainsaw",
+    "storage chest", "chiffonier", "bell or wind chime", "china cabinet", "Christmas stocking",
+    "church", "movie theater", "cleaver", "cliff dwelling", "cloak", "clogs", "cocktail shaker",
+    "coffee mug", "coffeemaker", "spiral or coil", "combination lock", "computer keyboard",
+    "candy store", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot",
+    "cowboy hat", "cradle", "construction crane", "crash helmet", "crate", "infant bed",
+    "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer",
+    "rotary dial telephone", "diaper", "digital clock", "digital watch", "dining table",
+    "dishcloth", "dishwasher", "disc brake", "dock", "dog sled", "dome", "doormat", "drilling rig",
+    "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar",
+    "electric locomotive", "entertainment center", "envelope", "espresso machine", "face powder",
+    "feather boa", "filing cabinet", "fireboat", "fire truck", "fire screen", "flagpole", "flute",
+    "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster bed",
+    "freight car", "French horn", "frying pan", "fur coat", "garbage truck",
+    "gas mask or respirator", "gas pump", "goblet", "go-kart", "golf ball", "golf cart", "gondola",
+    "gong", "gown", "grand piano", "greenhouse", "radiator grille", "grocery store", "guillotine",
+    "hair clip", "hair spray", "half-track", "hammer", "hamper", "hair dryer", "hand-held computer",
+    "handkerchief", "hard disk drive", "harmonica", "harp", "combine harvester", "hatchet",
+    "holster", "home theater", "honeycomb", "hook", "hoop skirt", "gymnastic horizontal bar",
+    "horse-drawn vehicle", "hourglass", "iPod", "clothes iron", "carved pumpkin", "jeans", "jeep",
+    "T-shirt", "jigsaw puzzle", "rickshaw", "joystick", "kimono", "knee pad", "knot", "lab coat",
+    "ladle", "lampshade", "laptop computer", "lawn mower", "lens cap", "letter opener", "library",
+    "lifeboat", "lighter", "limousine", "ocean liner", "lipstick", "slip-on shoe", "lotion",
+    "music speaker", "loupe magnifying glass", "sawmill", "magnetic compass", "messenger bag",
+    "mailbox", "tights", "one-piece bathing suit", "manhole cover", "maraca", "marimba", "mask",
+    "matchstick", "maypole", "maze", "measuring cup", "medicine cabinet", "megalith", "microphone",
+    "microwave oven", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile",
+    "mitten", "mixing bowl", "mobile home", "ford model t", "modem", "monastery", "monitor",
+    "moped", "mortar and pestle", "graduation cap", "mosque", "mosquito net", "vespa",
+    "mountain bike", "tent", "computer mouse", "mousetrap", "moving van", "muzzle", "metal nail",
+    "neck brace", "necklace", "baby pacifier", "notebook computer", "obelisk", "oboe", "ocarina",
+    "odometer", "oil filter", "pipe organ", "oscilloscope", "overskirt", "bullock cart",
+    "oxygen mask", "product packet / packaging", "paddle", "paddle wheel", "padlock", "paintbrush",
+    "pajamas", "palace", "pan flute", "paper towel", "parachute", "parallel bars", "park bench",
+    "parking meter", "railroad car", "patio", "payphone", "pedestal", "pencil case",
+    "pencil sharpener", "perfume", "Petri dish", "photocopier", "plectrum", "Pickelhaube",
+    "picket fence", "pickup truck", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball",
+    "pinwheel", "pirate ship", "drink pitcher", "block plane", "planetarium", "plastic bag",
+    "plate rack", "farm plow", "plunger", "Polaroid camera", "pole", "police van", "poncho",
+    "pool table", "soda bottle", "plant pot", "potter's wheel", "power drill", "prayer rug",
+    "printer", "prison", "missile", "projector", "hockey puck", "punching bag", "purse", "quill",
+    "quilt", "race car", "racket", "radiator", "radio", "radio telescope", "rain barrel",
+    "recreational vehicle", "fishing casting reel", "reflex camera", "refrigerator",
+    "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "eraser",
+    "rugby ball", "ruler measuring stick", "sneaker", "safe", "safety pin", "salt shaker", "sandal",
+    "sarong", "saxophone", "scabbard", "weighing scale", "school bus", "schooner", "scoreboard",
+    "CRT monitor", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe store",
+    "shoji screen / room divider", "shopping basket", "shopping cart", "shovel", "shower cap",
+    "shower curtain", "ski", "balaclava ski mask", "sleeping bag", "slide rule", "sliding door",
+    "slot machine", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock",
+    "solar thermal collector", "sombrero", "soup bowl", "keyboard space bar", "space heater",
+    "space shuttle", "spatula", "motorboat", "spider web", "spindle", "sports car", "spotlight",
+    "stage", "steam locomotive", "through arch bridge", "steel drum", "stethoscope", "scarf",
+    "stone wall", "stopwatch", "stove", "strainer", "tram", "stretcher", "couch", "stupa",
+    "submarine", "suit", "sundial", "sunglasses", "sunglasses", "sunscreen", "suspension bridge",
+    "mop", "sweatshirt", "swim trunks / shorts", "swing", "electrical switch", "syringe",
+    "table lamp", "tank", "tape player", "teapot", "teddy bear", "television", "tennis ball",
+    "thatched roof", "front curtain", "thimble", "threshing machine", "throne", "tile roof",
+    "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toy store",
+    "tractor", "semi-trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod",
+    "triumphal arch", "trolleybus", "trombone", "hot tub", "turnstile", "typewriter keyboard",
+    "umbrella", "unicycle", "upright piano", "vacuum cleaner", "vase", "vaulted or arched ceiling",
+    "velvet fabric", "vending machine", "vestment", "viaduct", "violin", "volleyball",
+    "waffle iron", "wall clock", "wallet", "wardrobe", "military aircraft", "sink",
+    "washing machine", "water bottle", "water jug", "water tower", "whiskey jug", "whistle",
+    "hair wig", "window screen", "window shade", "Windsor tie", "wine bottle", "airplane wing",
+    "wok", "wooden spoon", "wool", "split-rail fence", "shipwreck", "sailboat", "yurt", "website",
+    "comic book", "crossword", "traffic or street sign", "traffic light", "dust jacket", "menu",
+    "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "popsicle", "baguette",
+    "bagel", "pretzel", "cheeseburger", "hot dog", "mashed potatoes", "cabbage", "broccoli",
+    "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber",
+    "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith apple", "strawberry", "orange",
+    "lemon", "fig", "pineapple", "banana", "jackfruit", "cherimoya (custard apple)", "pomegranate",
+    "hay", "carbonara", "chocolate syrup", "dough", "meatloaf", "pizza", "pot pie", "burrito",
+    "red wine", "espresso", "tea cup", "eggnog", "mountain", "bubble", "cliff", "coral reef",
+    "geyser", "lakeshore", "promontory", "sandbar", "beach", "valley", "volcano", "baseball player",
+    "bridegroom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn",
+    "rose hip", "horse chestnut seed", "coral fungus", "agaric", "gyromitra", "stinkhorn mushroom",
+    "earth star fungus", "hen of the woods mushroom", "bolete", "corn cob", "toilet paper"
+)
+
+IMAGENET_CLASSES = tuple(label.lower() for label in IMAGENET_CLASSES)
+
+# SigLIP 2 推荐的 Prompt 模板 (也可以使用 ensemble)
+TEMPLATE = "a photo of a {}"
+
+def load_model(model_path):
+    """加载本地 SigLIP 模型和处理器"""
+    print(f"正在加载模型: {model_path} ...")
+    # SigLIP 2 在 HF 中通常复用 SiglipModel 结构
+    model = SiglipModel.from_pretrained(model_path, torch_dtype=torch.float16)
+    processor = SiglipProcessor.from_pretrained(model_path)
+    model.to(DEVICE)
+    model.eval()
+    return model, processor
+
+def build_zero_shot_classifier(model, processor, class_names, template):
+    """
+    预计算文本嵌入 (Text Embeddings) 作为分类器权重
+    """
+    print("构建 Zero-Shot 分类器 (文本编码)...")
+
+    # 构建 Prompts: ["a photo of a tench", "a photo of a goldfish", ...]
+    texts = [template.format(label) for label in class_names]
+
+    # 分批处理文本以防 OOM
+    text_features_list = []
+    batch_size = 256
+
+    # 修复1: 使用 torch.amp.autocast('cuda') 替代过期的 torch.cuda.amp.autocast()
+    with torch.no_grad(), torch.amp.autocast('cuda'):
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i : i + batch_size]
+            inputs = processor(
+                text=batch_texts,
+                padding="max_length",
+                max_length=64, # 修复2: 显式指定 max_length (SigLIP 通常为 64)，解决 padding 失败导致的维度错误
+                return_tensors="pt",
+                truncation=True
+            ).to(DEVICE)
+
+            # 获取文本特征
+            text_outputs = model.get_text_features(**inputs)
+
+            # 归一化 (SigLIP/CLIP 必须步骤)
+            text_outputs = text_outputs / text_outputs.norm(p=2, dim=-1, keepdim=True)
+            text_features_list.append(text_outputs)
+
+    # 合并所有类别的特征 [1000, Hidden_Dim]
+    text_classifier = torch.cat(text_features_list, dim=0)
+    return text_classifier
+
+class ImageNetDataset(torch.utils.data.Dataset):
+    """
+    自定义 Dataset 包装器，利用 SiglipProcessor 处理图像
+    """
+    def __init__(self, root_dir, processor):
+        # 使用 ImageFolder 自动处理文件夹结构和标签索引
+        self.dataset = datasets.ImageFolder(root_dir)
+        self.processor = processor
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        image, label = self.dataset[idx]
+        # 这里的 image 是 PIL Image
+        return image, label
+
+def collate_fn(batch, processor):
+    """
+    自定义 collate_fn，将 PIL 图片列表转换为 Tensor
+    """
+    images, labels = zip(*batch)
+    # processor 处理图片: Resize, Rescale, Normalize
+    # 修复: zip(*batch) 返回的是 tuple，必须转为 list 才能被 processor 接受
+    inputs = processor(images=list(images), return_tensors="pt")
+    return inputs["pixel_values"], torch.tensor(labels)
+
+def evaluate(model, classifier, dataloader):
+    """执行评测循环"""
+    print("开始评测 ImageNet-1k...")
+    top1_correct = 0
+    top5_correct = 0
+    total_samples = 0
+
+    # SigLIP 的 Logit Scale (Temperature)
+    logit_scale = model.logit_scale.exp()
+    # 如果是 SigLIP 2，bias 也是需要的
+    logit_bias = model.logit_bias if hasattr(model, 'logit_bias') else 0
+
+    # 修复3: 同步更新 evaluate 中的 autocast 写法
+    with torch.no_grad(), torch.amp.autocast('cuda'):
+        for pixel_values, targets in tqdm(dataloader):
+            pixel_values = pixel_values.to(DEVICE).to(torch.float16)
+            targets = targets.to(DEVICE)
+
+            # 1. 图像编码
+            # SigLIP image encoder 输出
+            image_features = model.get_image_features(pixel_values=pixel_values)
+
+            # 2. 归一化图像特征
+            image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True)
+
+            # 3. 计算 Logits
+            # Logits = (Image_Feat @ Text_Feat.T) * Scale + Bias
+            logits = (torch.matmul(image_features, classifier.t()) * logit_scale) + logit_bias
+
+            # 4. 计算准确率
+            # Top-1
+            pred_top1 = logits.argmax(dim=1)
+            top1_correct += (pred_top1 == targets).sum().item()
+
+            # Top-5
+            _, pred_top5 = logits.topk(5, dim=1)
+            top5_correct += (pred_top5 == targets.unsqueeze(1)).sum().item()
+
+            total_samples += targets.size(0)
+
+    acc1 = top1_correct / total_samples
+    acc5 = top5_correct / total_samples
+
+    print(f"\n====== 评测完成 ======")
+    print(f"Total Samples: {total_samples}")
+    print(f"Top-1 Accuracy: {acc1 * 100:.2f}%")
+    print(f"Top-5 Accuracy: {acc5 * 100:.2f}%")
+
+    return acc1, acc5
+
+# =================主执行逻辑=================
+if __name__ == "__main__":
+    from functools import partial
+
+    # 1. 加载模型
+    # 确保你的本地路径正确
+    if not os.path.exists(MODEL_PATH):
+        raise FileNotFoundError(f"找不到模型路径: {MODEL_PATH}")
+
+    model, processor = load_model(MODEL_PATH)
+
+    # 2. 构建 Zero-Shot 文本分类器
+    # 注意：IMAGENET_CLASSES 必须与 datasets.ImageFolder 的类别顺序一致
+    # ImageFolder 默认按照文件夹名称(n01440764...)的字母顺序排序
+    # 这里我们假设下载的 IMAGENET_CLASSES 是已经按照 WordNet ID 排序好的
+    classifier = build_zero_shot_classifier(model, processor, IMAGENET_CLASSES, TEMPLATE)
+
+    # 3. 准备数据加载器
+    # 使用 partial 将 processor 传入 collate_fn
+    collate_func = partial(collate_fn, processor=processor)
+
+    dataset = ImageNetDataset(IMAGENET_VAL_PATH, processor)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=BATCH_SIZE,
+        shuffle=False,
+        num_workers=NUM_WORKERS,
+        collate_fn=collate_func,
+        pin_memory=True
+    )
+
+    # 4. 开始评测
+    evaluate(model, classifier, dataloader)

patch-reparameterization/src/eval/clip_test.py

@@ -0,0 +1,567 @@
+#!/usr/bin/env python3
+"""
+Zero-shot evaluation script for VTP.
+
+Modified from OpenCLIP.
+
+Usage:
+    python tools/test_zero_shot_hf.py --model_path /path/to/vtp_hf_model --data_path /path/to/imagenet/val
+
+    # With custom settings
+    python tools/test_zero_shot_hf.py \
+        --model_path /path/to/vtp_hf_model \
+        --data_path /path/to/imagenet/val \
+        --batch_size 256 \
+        --device cuda:0 \
+        --precision bf16
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from functools import partial
+from itertools import islice
+from pathlib import Path
+from typing import Callable, List, Sequence, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from torchvision.datasets import ImageFolder
+from tqdm import tqdm
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from vtp.models.vtp_hf import VTPConfig, VTPModel
+from vtp.tokenizers import get_tokenizer
+
+
+# ============================================================================
+# ImageNet classnames and templates (copied from tools/zero_shot/templates.py)
+# ============================================================================
+
+IMAGENET_CLASSNAMES = (
+    "tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", "electric ray",
+    "stingray", "rooster", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco",
+    "indigo bunting", "American robin", "bulbul", "jay", "magpie", "chickadee", "American dipper",
+    "kite (bird of prey)", "bald eagle", "vulture", "great grey owl", "fire salamander",
+    "smooth newt", "newt", "spotted salamander", "axolotl", "American bullfrog", "tree frog",
+    "tailed frog", "loggerhead sea turtle", "leatherback sea turtle", "mud turtle", "terrapin",
+    "box turtle", "banded gecko", "green iguana", "Carolina anole",
+    "desert grassland whiptail lizard", "agama", "frilled-necked lizard", "alligator lizard",
+    "Gila monster", "European green lizard", "chameleon", "Komodo dragon", "Nile crocodile",
+    "American alligator", "triceratops", "worm snake", "ring-necked snake",
+    "eastern hog-nosed snake", "smooth green snake", "kingsnake", "garter snake", "water snake",
+    "vine snake", "night snake", "boa constrictor", "African rock python", "Indian cobra",
+    "green mamba", "sea snake", "Saharan horned viper", "eastern diamondback rattlesnake",
+    "sidewinder rattlesnake", "trilobite", "harvestman", "scorpion", "yellow garden spider",
+    "barn spider", "European garden spider", "southern black widow", "tarantula", "wolf spider",
+    "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie grouse", "peafowl",
+    "quail", "partridge", "african grey parrot", "macaw", "sulphur-crested cockatoo", "lorikeet",
+    "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "duck",
+    "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby",
+    "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch",
+    "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab",
+    "fiddler crab", "red king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab",
+    "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron",
+    "great egret", "bittern bird", "crane bird", "limpkin", "common gallinule", "American coot",
+    "bustard", "ruddy turnstone", "dunlin", "common redshank", "dowitcher", "oystercatcher",
+    "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion",
+    "Chihuahua", "Japanese Chin", "Maltese", "Pekingese", "Shih Tzu", "King Charles Spaniel",
+    "Papillon", "toy terrier", "Rhodesian Ridgeback", "Afghan Hound", "Basset Hound", "Beagle",
+    "Bloodhound", "Bluetick Coonhound", "Black and Tan Coonhound", "Treeing Walker Coonhound",
+    "English foxhound", "Redbone Coonhound", "borzoi", "Irish Wolfhound", "Italian Greyhound",
+    "Whippet", "Ibizan Hound", "Norwegian Elkhound", "Otterhound", "Saluki", "Scottish Deerhound",
+    "Weimaraner", "Staffordshire Bull Terrier", "American Staffordshire Terrier",
+    "Bedlington Terrier", "Border Terrier", "Kerry Blue Terrier", "Irish Terrier",
+    "Norfolk Terrier", "Norwich Terrier", "Yorkshire Terrier", "Wire Fox Terrier",
+    "Lakeland Terrier", "Sealyham Terrier", "Airedale Terrier", "Cairn Terrier",
+    "Australian Terrier", "Dandie Dinmont Terrier", "Boston Terrier", "Miniature Schnauzer",
+    "Giant Schnauzer", "Standard Schnauzer", "Scottish Terrier", "Tibetan Terrier",
+    "Australian Silky Terrier", "Soft-coated Wheaten Terrier", "West Highland White Terrier",
+    "Lhasa Apso", "Flat-Coated Retriever", "Curly-coated Retriever", "Golden Retriever",
+    "Labrador Retriever", "Chesapeake Bay Retriever", "German Shorthaired Pointer", "Vizsla",
+    "English Setter", "Irish Setter", "Gordon Setter", "Brittany dog", "Clumber Spaniel",
+    "English Springer Spaniel", "Welsh Springer Spaniel", "Cocker Spaniel", "Sussex Spaniel",
+    "Irish Water Spaniel", "Kuvasz", "Schipperke", "Groenendael dog", "Malinois", "Briard",
+    "Australian Kelpie", "Komondor", "Old English Sheepdog", "Shetland Sheepdog", "collie",
+    "Border Collie", "Bouvier des Flandres dog", "Rottweiler", "German Shepherd Dog", "Dobermann",
+    "Miniature Pinscher", "Greater Swiss Mountain Dog", "Bernese Mountain Dog",
+    "Appenzeller Sennenhund", "Entlebucher Sennenhund", "Boxer", "Bullmastiff", "Tibetan Mastiff",
+    "French Bulldog", "Great Dane", "St. Bernard", "husky", "Alaskan Malamute", "Siberian Husky",
+    "Dalmatian", "Affenpinscher", "Basenji", "pug", "Leonberger", "Newfoundland dog",
+    "Great Pyrenees dog", "Samoyed", "Pomeranian", "Chow Chow", "Keeshond", "brussels griffon",
+    "Pembroke Welsh Corgi", "Cardigan Welsh Corgi", "Toy Poodle", "Miniature Poodle",
+    "Standard Poodle", "Mexican hairless dog (xoloitzcuintli)", "grey wolf", "Alaskan tundra wolf",
+    "red wolf or maned wolf", "coyote", "dingo", "dhole", "African wild dog", "hyena", "red fox",
+    "kit fox", "Arctic fox", "grey fox", "tabby cat", "tiger cat", "Persian cat", "Siamese cat",
+    "Egyptian Mau", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger",
+    "cheetah", "brown bear", "American black bear", "polar bear", "sloth bear", "mongoose",
+    "meerkat", "tiger beetle", "ladybug", "ground beetle", "longhorn beetle", "leaf beetle",
+    "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper",
+    "cricket insect", "stick insect", "cockroach", "praying mantis", "cicada", "leafhopper",
+    "lacewing", "dragonfly", "damselfly", "red admiral butterfly", "ringlet butterfly",
+    "monarch butterfly", "small white butterfly", "sulphur butterfly", "gossamer-winged butterfly",
+    "starfish", "sea urchin", "sea cucumber", "cottontail rabbit", "hare", "Angora rabbit",
+    "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "common sorrel horse",
+    "zebra", "pig", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison",
+    "ram (adult male sheep)", "bighorn sheep", "Alpine ibex", "hartebeest", "impala (antelope)",
+    "gazelle", "arabian camel", "llama", "weasel", "mink", "European polecat",
+    "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan",
+    "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas monkey", "baboon", "macaque",
+    "langur", "black-and-white colobus", "proboscis monkey", "marmoset", "white-headed capuchin",
+    "howler monkey", "titi monkey", "Geoffroy's spider monkey", "common squirrel monkey",
+    "ring-tailed lemur", "indri", "Asian elephant", "African bush elephant", "red panda",
+    "giant panda", "snoek fish", "eel", "silver salmon", "rock beauty fish", "clownfish",
+    "sturgeon", "gar fish", "lionfish", "pufferfish", "abacus", "abaya", "academic gown",
+    "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance",
+    "amphibious vehicle", "analog clock", "apiary", "apron", "trash can", "assault rifle",
+    "backpack", "bakery", "balance beam", "balloon", "ballpoint pen", "Band-Aid", "banjo",
+    "baluster / handrail", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel",
+    "wheelbarrow", "baseball", "basketball", "bassinet", "bassoon", "swimming cap", "bath towel",
+    "bathtub", "station wagon", "lighthouse", "beaker", "military hat (bearskin or shako)",
+    "beer bottle", "beer glass", "bell tower", "baby bib", "tandem bicycle", "bikini",
+    "ring binder", "binoculars", "birdhouse", "boathouse", "bobsleigh", "bolo tie", "poke bonnet",
+    "bookcase", "bookstore", "bottle cap", "hunting bow", "bow tie", "brass memorial plaque", "bra",
+    "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest",
+    "high-speed train", "butcher shop", "taxicab", "cauldron", "candle", "cannon", "canoe",
+    "can opener", "cardigan", "car mirror", "carousel", "tool kit", "cardboard box / carton",
+    "car wheel", "automated teller machine", "cassette", "cassette player", "castle", "catamaran",
+    "CD player", "cello", "mobile phone", "chain", "chain-link fence", "chain mail", "chainsaw",
+    "storage chest", "chiffonier", "bell or wind chime", "china cabinet", "Christmas stocking",
+    "church", "movie theater", "cleaver", "cliff dwelling", "cloak", "clogs", "cocktail shaker",
+    "coffee mug", "coffeemaker", "spiral or coil", "combination lock", "computer keyboard",
+    "candy store", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot",
+    "cowboy hat", "cradle", "construction crane", "crash helmet", "crate", "infant bed",
+    "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer",
+    "rotary dial telephone", "diaper", "digital clock", "digital watch", "dining table",
+    "dishcloth", "dishwasher", "disc brake", "dock", "dog sled", "dome", "doormat", "drilling rig",
+    "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar",
+    "electric locomotive", "entertainment center", "envelope", "espresso machine", "face powder",
+    "feather boa", "filing cabinet", "fireboat", "fire truck", "fire screen", "flagpole", "flute",
+    "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster bed",
+    "freight car", "French horn", "frying pan", "fur coat", "garbage truck",
+    "gas mask or respirator", "gas pump", "goblet", "go-kart", "golf ball", "golf cart", "gondola",
+    "gong", "gown", "grand piano", "greenhouse", "radiator grille", "grocery store", "guillotine",
+    "hair clip", "hair spray", "half-track", "hammer", "hamper", "hair dryer", "hand-held computer",
+    "handkerchief", "hard disk drive", "harmonica", "harp", "combine harvester", "hatchet",
+    "holster", "home theater", "honeycomb", "hook", "hoop skirt", "gymnastic horizontal bar",
+    "horse-drawn vehicle", "hourglass", "iPod", "clothes iron", "carved pumpkin", "jeans", "jeep",
+    "T-shirt", "jigsaw puzzle", "rickshaw", "joystick", "kimono", "knee pad", "knot", "lab coat",
+    "ladle", "lampshade", "laptop computer", "lawn mower", "lens cap", "letter opener", "library",
+    "lifeboat", "lighter", "limousine", "ocean liner", "lipstick", "slip-on shoe", "lotion",
+    "music speaker", "loupe magnifying glass", "sawmill", "magnetic compass", "messenger bag",
+    "mailbox", "tights", "one-piece bathing suit", "manhole cover", "maraca", "marimba", "mask",
+    "matchstick", "maypole", "maze", "measuring cup", "medicine cabinet", "megalith", "microphone",
+    "microwave oven", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile",
+    "mitten", "mixing bowl", "mobile home", "ford model t", "modem", "monastery", "monitor",
+    "moped", "mortar and pestle", "graduation cap", "mosque", "mosquito net", "vespa",
+    "mountain bike", "tent", "computer mouse", "mousetrap", "moving van", "muzzle", "metal nail",
+    "neck brace", "necklace", "baby pacifier", "notebook computer", "obelisk", "oboe", "ocarina",
+    "odometer", "oil filter", "pipe organ", "oscilloscope", "overskirt", "bullock cart",
+    "oxygen mask", "product packet / packaging", "paddle", "paddle wheel", "padlock", "paintbrush",
+    "pajamas", "palace", "pan flute", "paper towel", "parachute", "parallel bars", "park bench",
+    "parking meter", "railroad car", "patio", "payphone", "pedestal", "pencil case",
+    "pencil sharpener", "perfume", "Petri dish", "photocopier", "plectrum", "Pickelhaube",
+    "picket fence", "pickup truck", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball",
+    "pinwheel", "pirate ship", "drink pitcher", "block plane", "planetarium", "plastic bag",
+    "plate rack", "farm plow", "plunger", "Polaroid camera", "pole", "police van", "poncho",
+    "pool table", "soda bottle", "plant pot", "potter's wheel", "power drill", "prayer rug",
+    "printer", "prison", "missile", "projector", "hockey puck", "punching bag", "purse", "quill",
+    "quilt", "race car", "racket", "radiator", "radio", "radio telescope", "rain barrel",
+    "recreational vehicle", "fishing casting reel", "reflex camera", "refrigerator",
+    "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "eraser",
+    "rugby ball", "ruler measuring stick", "sneaker", "safe", "safety pin", "salt shaker", "sandal",
+    "sarong", "saxophone", "scabbard", "weighing scale", "school bus", "schooner", "scoreboard",
+    "CRT monitor", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe store",
+    "shoji screen / room divider", "shopping basket", "shopping cart", "shovel", "shower cap",
+    "shower curtain", "ski", "balaclava ski mask", "sleeping bag", "slide rule", "sliding door",
+    "slot machine", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock",
+    "solar thermal collector", "sombrero", "soup bowl", "keyboard space bar", "space heater",
+    "space shuttle", "spatula", "motorboat", "spider web", "spindle", "sports car", "spotlight",
+    "stage", "steam locomotive", "through arch bridge", "steel drum", "stethoscope", "scarf",
+    "stone wall", "stopwatch", "stove", "strainer", "tram", "stretcher", "couch", "stupa",
+    "submarine", "suit", "sundial", "sunglasses", "sunglasses", "sunscreen", "suspension bridge",
+    "mop", "sweatshirt", "swim trunks / shorts", "swing", "electrical switch", "syringe",
+    "table lamp", "tank", "tape player", "teapot", "teddy bear", "television", "tennis ball",
+    "thatched roof", "front curtain", "thimble", "threshing machine", "throne", "tile roof",
+    "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toy store",
+    "tractor", "semi-trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod",
+    "triumphal arch", "trolleybus", "trombone", "hot tub", "turnstile", "typewriter keyboard",
+    "umbrella", "unicycle", "upright piano", "vacuum cleaner", "vase", "vaulted or arched ceiling",
+    "velvet fabric", "vending machine", "vestment", "viaduct", "violin", "volleyball",
+    "waffle iron", "wall clock", "wallet", "wardrobe", "military aircraft", "sink",
+    "washing machine", "water bottle", "water jug", "water tower", "whiskey jug", "whistle",
+    "hair wig", "window screen", "window shade", "Windsor tie", "wine bottle", "airplane wing",
+    "wok", "wooden spoon", "wool", "split-rail fence", "shipwreck", "sailboat", "yurt", "website",
+    "comic book", "crossword", "traffic or street sign", "traffic light", "dust jacket", "menu",
+    "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "popsicle", "baguette",
+    "bagel", "pretzel", "cheeseburger", "hot dog", "mashed potatoes", "cabbage", "broccoli",
+    "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber",
+    "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith apple", "strawberry", "orange",
+    "lemon", "fig", "pineapple", "banana", "jackfruit", "cherimoya (custard apple)", "pomegranate",
+    "hay", "carbonara", "chocolate syrup", "dough", "meatloaf", "pizza", "pot pie", "burrito",
+    "red wine", "espresso", "tea cup", "eggnog", "mountain", "bubble", "cliff", "coral reef",
+    "geyser", "lakeshore", "promontory", "sandbar", "beach", "valley", "volcano", "baseball player",
+    "bridegroom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn",
+    "rose hip", "horse chestnut seed", "coral fungus", "agaric", "gyromitra", "stinkhorn mushroom",
+    "earth star fungus", "hen of the woods mushroom", "bolete", "corn cob", "toilet paper"
+)
+
+OPENAI_IMAGENET_TEMPLATES = (
+    lambda c: f'a bad photo of a {c}.',
+    lambda c: f'a photo of many {c}.',
+    lambda c: f'a sculpture of a {c}.',
+    lambda c: f'a photo of the hard to see {c}.',
+    lambda c: f'a low resolution photo of the {c}.',
+    lambda c: f'a rendering of a {c}.',
+    lambda c: f'graffiti of a {c}.',
+    lambda c: f'a bad photo of the {c}.',
+    lambda c: f'a cropped photo of the {c}.',
+    lambda c: f'a tattoo of a {c}.',
+    lambda c: f'the embroidered {c}.',
+    lambda c: f'a photo of a hard to see {c}.',
+    lambda c: f'a bright photo of a {c}.',
+    lambda c: f'a photo of a clean {c}.',
+    lambda c: f'a photo of a dirty {c}.',
+    lambda c: f'a dark photo of the {c}.',
+    lambda c: f'a drawing of a {c}.',
+    lambda c: f'a photo of my {c}.',
+    lambda c: f'the plastic {c}.',
+    lambda c: f'a photo of the cool {c}.',
+    lambda c: f'a close-up photo of a {c}.',
+    lambda c: f'a black and white photo of the {c}.',
+    lambda c: f'a painting of the {c}.',
+    lambda c: f'a painting of a {c}.',
+    lambda c: f'a pixelated photo of the {c}.',
+    lambda c: f'a sculpture of the {c}.',
+    lambda c: f'a bright photo of the {c}.',
+    lambda c: f'a cropped photo of a {c}.',
+    lambda c: f'a plastic {c}.',
+    lambda c: f'a photo of the dirty {c}.',
+    lambda c: f'a jpeg corrupted photo of a {c}.',
+    lambda c: f'a blurry photo of the {c}.',
+    lambda c: f'a photo of the {c}.',
+    lambda c: f'a good photo of the {c}.',
+    lambda c: f'a rendering of the {c}.',
+    lambda c: f'a {c} in a video game.',
+    lambda c: f'a photo of one {c}.',
+    lambda c: f'a doodle of a {c}.',
+    lambda c: f'a close-up photo of the {c}.',
+    lambda c: f'a photo of a {c}.',
+    lambda c: f'the origami {c}.',
+    lambda c: f'the {c} in a video game.',
+    lambda c: f'a sketch of a {c}.',
+    lambda c: f'a doodle of the {c}.',
+    lambda c: f'a origami {c}.',
+    lambda c: f'a low resolution photo of a {c}.',
+    lambda c: f'the toy {c}.',
+    lambda c: f'a rendition of the {c}.',
+    lambda c: f'a photo of the clean {c}.',
+    lambda c: f'a photo of a large {c}.',
+    lambda c: f'a rendition of a {c}.',
+    lambda c: f'a photo of a nice {c}.',
+    lambda c: f'a photo of a weird {c}.',
+    lambda c: f'a blurry photo of a {c}.',
+    lambda c: f'a cartoon {c}.',
+    lambda c: f'art of a {c}.',
+    lambda c: f'a sketch of the {c}.',
+    lambda c: f'a embroidered {c}.',
+    lambda c: f'a pixelated photo of a {c}.',
+    lambda c: f'itap of the {c}.',
+    lambda c: f'a jpeg corrupted photo of the {c}.',
+    lambda c: f'a good photo of a {c}.',
+    lambda c: f'a plushie {c}.',
+    lambda c: f'a photo of the nice {c}.',
+    lambda c: f'a photo of the small {c}.',
+    lambda c: f'a photo of the weird {c}.',
+    lambda c: f'the cartoon {c}.',
+    lambda c: f'art of the {c}.',
+    lambda c: f'a drawing of the {c}.',
+    lambda c: f'a photo of the large {c}.',
+    lambda c: f'a black and white photo of a {c}.',
+    lambda c: f'the plushie {c}.',
+    lambda c: f'a dark photo of a {c}.',
+    lambda c: f'itap of a {c}.',
+    lambda c: f'graffiti of the {c}.',
+    lambda c: f'a toy {c}.',
+    lambda c: f'itap of my {c}.',
+    lambda c: f'a photo of a cool {c}.',
+    lambda c: f'a photo of a small {c}.',
+    lambda c: f'a tattoo of the {c}.',
+)
+
+
+# ============================================================================
+# Utility functions
+# ============================================================================
+
+def batched(iterable, n: int):
+    """Batch data into lists of length n. The last batch may be shorter."""
+    it = iter(iterable)
+    while True:
+        batch = list(islice(it, n))
+        if not batch:
+            break
+        yield batch
+
+
+def accuracy(output: torch.Tensor, target: torch.Tensor, topk: Tuple[int, ...] = (1,)) -> List[float]:
+    """Calculate top-k accuracy."""
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) for k in topk]
+
+
+def get_input_dtype(precision: str) -> torch.dtype:
+    """Get input dtype based on precision string."""
+    if precision in ('bf16', 'bfloat16'):
+        return torch.bfloat16
+    elif precision in ('fp16', 'float16'):
+        return torch.float16
+    return torch.float32
+
+
+def get_autocast_context(precision: str, device_type: str = 'cuda'):
+    """Get autocast context manager based on precision."""
+    from contextlib import nullcontext
+    if precision in ('bf16', 'bfloat16'):
+        return partial(torch.amp.autocast, device_type=device_type, dtype=torch.bfloat16)
+    elif precision in ('fp16', 'float16'):
+        return partial(torch.amp.autocast, device_type=device_type, dtype=torch.float16)
+    return nullcontext
+
+
+# ============================================================================
+# Zero-shot classifier
+# ============================================================================
+
+def build_zero_shot_classifier(
+    model: VTPModel,
+    tokenizer,
+    classnames: Sequence[str],
+    templates: Sequence[Callable],
+    num_classes_per_batch: int = 10,
+    device: torch.device = None,
+    use_tqdm: bool = True,
+) -> torch.Tensor:
+    """Build zero-shot classifier weights.
+
+    Args:
+        model: VTPModel instance with get_clip_text_feature method.
+        tokenizer: Text tokenizer.
+        classnames: Sequence of class names.
+        templates: Sequence of template functions.
+        num_classes_per_batch: Number of classes to process per batch.
+        device: Device to use.
+        use_tqdm: Whether to show progress bar.
+
+    Returns:
+        Zero-shot classifier weights, shape (embed_dim, num_classes).
+    """
+    num_templates = len(templates)
+    num_classes = len(classnames)
+
+    if use_tqdm:
+        num_iter = (num_classes - 1) // num_classes_per_batch + 1
+        iter_wrap = partial(tqdm, total=num_iter, desc="Building classifier", unit_scale=num_classes_per_batch)
+    else:
+        iter_wrap = iter
+
+    def _process_batch(batch_classnames: List[str]) -> torch.Tensor:
+        num_batch_classes = len(batch_classnames)
+        texts = [template(c) for c in batch_classnames for template in templates]
+        tokens = tokenizer(texts).to(device)
+
+        # Use VTPModel's get_clip_text_feature method
+        text_features = model.get_clip_text_feature(tokens, normalize=True)
+
+        # Reshape and average over templates
+        text_features = text_features.reshape(num_batch_classes, num_templates, -1).mean(dim=1)
+        text_features = F.normalize(text_features, dim=1)
+        return text_features.T  # (embed_dim, num_batch_classes)
+
+    with torch.no_grad():
+        batched_embeds = [
+            _process_batch(batch)
+            for batch in iter_wrap(batched(classnames, num_classes_per_batch))
+        ]
+        zeroshot_weights = torch.cat(batched_embeds, dim=1)
+
+    return zeroshot_weights
+
+
+# ============================================================================
+# Evaluation
+# ============================================================================
+
+def evaluate(
+    model: VTPModel,
+    classifier: torch.Tensor,
+    dataloader: DataLoader,
+    device: torch.device,
+    precision: str = 'fp32',
+) -> Tuple[float, float]:
+    """Run zero-shot evaluation.
+
+    Args:
+        model: VTPModel instance.
+        classifier: Zero-shot classifier weights, shape (embed_dim, num_classes).
+        dataloader: ImageNet validation dataloader.
+        device: Device to use.
+        precision: Precision for inference.
+
+    Returns:
+        Tuple of (top1_accuracy, top5_accuracy).
+    """
+    autocast_ctx = get_autocast_context(precision, device_type=device.type)
+    input_dtype = get_input_dtype(precision)
+
+    top1, top5, n = 0.0, 0.0, 0
+
+    with torch.inference_mode():
+        for images, targets in tqdm(dataloader, desc="Evaluating"):
+            images = images.to(device=device, dtype=input_dtype)
+            targets = targets.to(device)
+
+            with autocast_ctx():
+                image_features = model.get_clip_image_feature(images, normalize=True)
+                logits = 100.0 * image_features @ classifier
+
+            acc1, acc5 = accuracy(logits, targets, topk=(1, 5))
+            top1 += acc1
+            top5 += acc5
+            n += images.size(0)
+
+    top1 = top1 / n * 100
+    top5 = top5 / n * 100
+    return top1, top5
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def create_imagenet_dataloader(
+    data_path: str,
+    image_size: int = 256,
+    batch_size: int = 128,
+    num_workers: int = 8,
+) -> DataLoader:
+    """Create ImageNet validation dataloader."""
+    preprocess = transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+    dataset = ImageFolder(root=data_path, transform=preprocess)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+    return dataloader
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Zero-shot ImageNet evaluation for VTP HuggingFace model")
+    parser.add_argument("--model_path", type=str, required=True,
+                        help="Path to VTP HuggingFace model directory")
+    parser.add_argument("--data_path", type=str, required=True,
+                        help="Path to ImageNet validation dataset")
+    parser.add_argument("--batch_size", type=int, default=128,
+                        help="Batch size for evaluation")
+    parser.add_argument("--num_workers", type=int, default=8,
+                        help="Number of dataloader workers")
+    parser.add_argument("--device", type=str, default="cuda:0",
+                        help="Device to use (e.g., cuda:0, cpu)")
+    parser.add_argument("--precision", type=str, default="fp32",
+                        choices=["fp32", "fp16", "bf16"],
+                        help="Precision for inference")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+
+    print("=" * 60)
+    print("Zero-Shot ImageNet Evaluation (VTP HuggingFace)")
+    print("=" * 60)
+    print(f"Model path: {args.model_path}")
+    print(f"Data path: {args.data_path}")
+    print(f"Device: {device}")
+    print(f"Precision: {args.precision}")
+    print(f"Batch size: {args.batch_size}")
+    print()
+
+    # Load model
+    print("Loading model...")
+    model = VTPModel.from_pretrained(args.model_path)
+    model = model.to(device)
+    model.eval()
+
+    # Get image size from config
+    image_size = model.config.image_size
+    context_length = model.config.text_context_length
+    print(f"Image size: {image_size}")
+    print(f"Context length: {context_length}")
+
+    # Load tokenizer
+    print("Loading tokenizer...")
+    tokenizer = get_tokenizer('ViT-B-32', context_length=context_length)
+
+    # Create dataloader
+    print(f"Loading ImageNet validation: {args.data_path}")
+    dataloader = create_imagenet_dataloader(
+        args.data_path,
+        image_size=image_size,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+    )
+    num_samples = len(dataloader.dataset)
+    print(f"Number of samples: {num_samples}")
+    if num_samples != 50000:
+        print(f"  [Warning] Expected 50000 samples for ImageNet validation, got {num_samples}")
+    print()
+
+    # Build classifier
+    print("Building zero-shot classifier...")
+    autocast_ctx = get_autocast_context(args.precision, device_type=device.type)
+    with autocast_ctx():
+        classifier = build_zero_shot_classifier(
+            model,
+            tokenizer,
+            classnames=IMAGENET_CLASSNAMES,
+            templates=OPENAI_IMAGENET_TEMPLATES,
+            num_classes_per_batch=10,
+            device=device,
+            use_tqdm=True,
+        )
+    print()
+
+    # Evaluate
+    print("Running evaluation...")
+    top1, top5 = evaluate(
+        model,
+        classifier,
+        dataloader,
+        device=device,
+        precision=args.precision,
+    )
+
+    print()
+    print("=" * 60)
+    print("Results:")
+    print(f"  ImageNet Zero-Shot Top-1 Accuracy: {top1:.2f}%")
+    print(f"  ImageNet Zero-Shot Top-5 Accuracy: {top5:.2f}%")
+    print("=" * 60)
+
+    return {"top1": top1, "top5": top5}
+
+
+if __name__ == "__main__":
+    main()

patch-reparameterization/src/eval/clip_zero_shot.py

@@ -0,0 +1,714 @@
+#!/usr/bin/env python3
+"""
+Zero-shot evaluation script for SigLIP2.
+
+Modified from OpenCLIP.
+
+Usage:
+    python clip_zero_shot.py --model_path /path/to/siglip2_model --data_path /path/to/imagenet/val
+
+    # With custom settings
+    python clip_zero_shot.py \
+        --model_path /path/to/siglip2_model \
+        --data_path /path/to/imagenet/val \
+        --batch_size 256 \
+        --device cuda:0 \
+        --precision bf16
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+import string
+import re
+
+from functools import partial
+from itertools import islice
+from pathlib import Path
+from typing import Callable, List, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from torchvision.datasets import ImageFolder
+from PIL import Image
+from tqdm import tqdm
+from transformers import AutoProcessor, SiglipModel
+
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# ImageNet classnames and templates (copied from tools/zero_shot/templates.py)
+# ============================================================================
+
+IMAGENET_CLASSNAMES = (
+    "tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", "electric ray",
+    "stingray", "rooster", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco",
+    "indigo bunting", "American robin", "bulbul", "jay", "magpie", "chickadee", "American dipper",
+    "kite (bird of prey)", "bald eagle", "vulture", "great grey owl", "fire salamander",
+    "smooth newt", "newt", "spotted salamander", "axolotl", "American bullfrog", "tree frog",
+    "tailed frog", "loggerhead sea turtle", "leatherback sea turtle", "mud turtle", "terrapin",
+    "box turtle", "banded gecko", "green iguana", "Carolina anole",
+    "desert grassland whiptail lizard", "agama", "frilled-necked lizard", "alligator lizard",
+    "Gila monster", "European green lizard", "chameleon", "Komodo dragon", "Nile crocodile",
+    "American alligator", "triceratops", "worm snake", "ring-necked snake",
+    "eastern hog-nosed snake", "smooth green snake", "kingsnake", "garter snake", "water snake",
+    "vine snake", "night snake", "boa constrictor", "African rock python", "Indian cobra",
+    "green mamba", "sea snake", "Saharan horned viper", "eastern diamondback rattlesnake",
+    "sidewinder rattlesnake", "trilobite", "harvestman", "scorpion", "yellow garden spider",
+    "barn spider", "European garden spider", "southern black widow", "tarantula", "wolf spider",
+    "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie grouse", "peafowl",
+    "quail", "partridge", "african grey parrot", "macaw", "sulphur-crested cockatoo", "lorikeet",
+    "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "duck",
+    "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby",
+    "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch",
+    "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab",
+    "fiddler crab", "red king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab",
+    "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron",
+    "great egret", "bittern bird", "crane bird", "limpkin", "common gallinule", "American coot",
+    "bustard", "ruddy turnstone", "dunlin", "common redshank", "dowitcher", "oystercatcher",
+    "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion",
+    "Chihuahua", "Japanese Chin", "Maltese", "Pekingese", "Shih Tzu", "King Charles Spaniel",
+    "Papillon", "toy terrier", "Rhodesian Ridgeback", "Afghan Hound", "Basset Hound", "Beagle",
+    "Bloodhound", "Bluetick Coonhound", "Black and Tan Coonhound", "Treeing Walker Coonhound",
+    "English foxhound", "Redbone Coonhound", "borzoi", "Irish Wolfhound", "Italian Greyhound",
+    "Whippet", "Ibizan Hound", "Norwegian Elkhound", "Otterhound", "Saluki", "Scottish Deerhound",
+    "Weimaraner", "Staffordshire Bull Terrier", "American Staffordshire Terrier",
+    "Bedlington Terrier", "Border Terrier", "Kerry Blue Terrier", "Irish Terrier",
+    "Norfolk Terrier", "Norwich Terrier", "Yorkshire Terrier", "Wire Fox Terrier",
+    "Lakeland Terrier", "Sealyham Terrier", "Airedale Terrier", "Cairn Terrier",
+    "Australian Terrier", "Dandie Dinmont Terrier", "Boston Terrier", "Miniature Schnauzer",
+    "Giant Schnauzer", "Standard Schnauzer", "Scottish Terrier", "Tibetan Terrier",
+    "Australian Silky Terrier", "Soft-coated Wheaten Terrier", "West Highland White Terrier",
+    "Lhasa Apso", "Flat-Coated Retriever", "Curly-coated Retriever", "Golden Retriever",
+    "Labrador Retriever", "Chesapeake Bay Retriever", "German Shorthaired Pointer", "Vizsla",
+    "English Setter", "Irish Setter", "Gordon Setter", "Brittany dog", "Clumber Spaniel",
+    "English Springer Spaniel", "Welsh Springer Spaniel", "Cocker Spaniel", "Sussex Spaniel",
+    "Irish Water Spaniel", "Kuvasz", "Schipperke", "Groenendael dog", "Malinois", "Briard",
+    "Australian Kelpie", "Komondor", "Old English Sheepdog", "Shetland Sheepdog", "collie",
+    "Border Collie", "Bouvier des Flandres dog", "Rottweiler", "German Shepherd Dog", "Dobermann",
+    "Miniature Pinscher", "Greater Swiss Mountain Dog", "Bernese Mountain Dog",
+    "Appenzeller Sennenhund", "Entlebucher Sennenhund", "Boxer", "Bullmastiff", "Tibetan Mastiff",
+    "French Bulldog", "Great Dane", "St. Bernard", "husky", "Alaskan Malamute", "Siberian Husky",
+    "Dalmatian", "Affenpinscher", "Basenji", "pug", "Leonberger", "Newfoundland dog",
+    "Great Pyrenees dog", "Samoyed", "Pomeranian", "Chow Chow", "Keeshond", "brussels griffon",
+    "Pembroke Welsh Corgi", "Cardigan Welsh Corgi", "Toy Poodle", "Miniature Poodle",
+    "Standard Poodle", "Mexican hairless dog (xoloitzcuintli)", "grey wolf", "Alaskan tundra wolf",
+    "red wolf or maned wolf", "coyote", "dingo", "dhole", "African wild dog", "hyena", "red fox",
+    "kit fox", "Arctic fox", "grey fox", "tabby cat", "tiger cat", "Persian cat", "Siamese cat",
+    "Egyptian Mau", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger",
+    "cheetah", "brown bear", "American black bear", "polar bear", "sloth bear", "mongoose",
+    "meerkat", "tiger beetle", "ladybug", "ground beetle", "longhorn beetle", "leaf beetle",
+    "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper",
+    "cricket insect", "stick insect", "cockroach", "praying mantis", "cicada", "leafhopper",
+    "lacewing", "dragonfly", "damselfly", "red admiral butterfly", "ringlet butterfly",
+    "monarch butterfly", "small white butterfly", "sulphur butterfly", "gossamer-winged butterfly",
+    "starfish", "sea urchin", "sea cucumber", "cottontail rabbit", "hare", "Angora rabbit",
+    "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "common sorrel horse",
+    "zebra", "pig", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison",
+    "ram (adult male sheep)", "bighorn sheep", "Alpine ibex", "hartebeest", "impala (antelope)",
+    "gazelle", "arabian camel", "llama", "weasel", "mink", "European polecat",
+    "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan",
+    "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas monkey", "baboon", "macaque",
+    "langur", "black-and-white colobus", "proboscis monkey", "marmoset", "white-headed capuchin",
+    "howler monkey", "titi monkey", "Geoffroy's spider monkey", "common squirrel monkey",
+    "ring-tailed lemur", "indri", "Asian elephant", "African bush elephant", "red panda",
+    "giant panda", "snoek fish", "eel", "silver salmon", "rock beauty fish", "clownfish",
+    "sturgeon", "gar fish", "lionfish", "pufferfish", "abacus", "abaya", "academic gown",
+    "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance",
+    "amphibious vehicle", "analog clock", "apiary", "apron", "trash can", "assault rifle",
+    "backpack", "bakery", "balance beam", "balloon", "ballpoint pen", "Band-Aid", "banjo",
+    "baluster / handrail", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel",
+    "wheelbarrow", "baseball", "basketball", "bassinet", "bassoon", "swimming cap", "bath towel",
+    "bathtub", "station wagon", "lighthouse", "beaker", "military hat (bearskin or shako)",
+    "beer bottle", "beer glass", "bell tower", "baby bib", "tandem bicycle", "bikini",
+    "ring binder", "binoculars", "birdhouse", "boathouse", "bobsleigh", "bolo tie", "poke bonnet",
+    "bookcase", "bookstore", "bottle cap", "hunting bow", "bow tie", "brass memorial plaque", "bra",
+    "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest",
+    "high-speed train", "butcher shop", "taxicab", "cauldron", "candle", "cannon", "canoe",
+    "can opener", "cardigan", "car mirror", "carousel", "tool kit", "cardboard box / carton",
+    "car wheel", "automated teller machine", "cassette", "cassette player", "castle", "catamaran",
+    "CD player", "cello", "mobile phone", "chain", "chain-link fence", "chain mail", "chainsaw",
+    "storage chest", "chiffonier", "bell or wind chime", "china cabinet", "Christmas stocking",
+    "church", "movie theater", "cleaver", "cliff dwelling", "cloak", "clogs", "cocktail shaker",
+    "coffee mug", "coffeemaker", "spiral or coil", "combination lock", "computer keyboard",
+    "candy store", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot",
+    "cowboy hat", "cradle", "construction crane", "crash helmet", "crate", "infant bed",
+    "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer",
+    "rotary dial telephone", "diaper", "digital clock", "digital watch", "dining table",
+    "dishcloth", "dishwasher", "disc brake", "dock", "dog sled", "dome", "doormat", "drilling rig",
+    "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar",
+    "electric locomotive", "entertainment center", "envelope", "espresso machine", "face powder",
+    "feather boa", "filing cabinet", "fireboat", "fire truck", "fire screen", "flagpole", "flute",
+    "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster bed",
+    "freight car", "French horn", "frying pan", "fur coat", "garbage truck",
+    "gas mask or respirator", "gas pump", "goblet", "go-kart", "golf ball", "golf cart", "gondola",
+    "gong", "gown", "grand piano", "greenhouse", "radiator grille", "grocery store", "guillotine",
+    "hair clip", "hair spray", "half-track", "hammer", "hamper", "hair dryer", "hand-held computer",
+    "handkerchief", "hard disk drive", "harmonica", "harp", "combine harvester", "hatchet",
+    "holster", "home theater", "honeycomb", "hook", "hoop skirt", "gymnastic horizontal bar",
+    "horse-drawn vehicle", "hourglass", "iPod", "clothes iron", "carved pumpkin", "jeans", "jeep",
+    "T-shirt", "jigsaw puzzle", "rickshaw", "joystick", "kimono", "knee pad", "knot", "lab coat",
+    "ladle", "lampshade", "laptop computer", "lawn mower", "lens cap", "letter opener", "library",
+    "lifeboat", "lighter", "limousine", "ocean liner", "lipstick", "slip-on shoe", "lotion",
+    "music speaker", "loupe magnifying glass", "sawmill", "magnetic compass", "messenger bag",
+    "mailbox", "tights", "one-piece bathing suit", "manhole cover", "maraca", "marimba", "mask",
+    "matchstick", "maypole", "maze", "measuring cup", "medicine cabinet", "megalith", "microphone",
+    "microwave oven", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile",
+    "mitten", "mixing bowl", "mobile home", "ford model t", "modem", "monastery", "monitor",
+    "moped", "mortar and pestle", "graduation cap", "mosque", "mosquito net", "vespa",
+    "mountain bike", "tent", "computer mouse", "mousetrap", "moving van", "muzzle", "metal nail",
+    "neck brace", "necklace", "baby pacifier", "notebook computer", "obelisk", "oboe", "ocarina",
+    "odometer", "oil filter", "pipe organ", "oscilloscope", "overskirt", "bullock cart",
+    "oxygen mask", "product packet / packaging", "paddle", "paddle wheel", "padlock", "paintbrush",
+    "pajamas", "palace", "pan flute", "paper towel", "parachute", "parallel bars", "park bench",
+    "parking meter", "railroad car", "patio", "payphone", "pedestal", "pencil case",
+    "pencil sharpener", "perfume", "Petri dish", "photocopier", "plectrum", "Pickelhaube",
+    "picket fence", "pickup truck", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball",
+    "pinwheel", "pirate ship", "drink pitcher", "block plane", "planetarium", "plastic bag",
+    "plate rack", "farm plow", "plunger", "Polaroid camera", "pole", "police van", "poncho",
+    "pool table", "soda bottle", "plant pot", "potter's wheel", "power drill", "prayer rug",
+    "printer", "prison", "missile", "projector", "hockey puck", "punching bag", "purse", "quill",
+    "quilt", "race car", "racket", "radiator", "radio", "radio telescope", "rain barrel",
+    "recreational vehicle", "fishing casting reel", "reflex camera", "refrigerator",
+    "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "eraser",
+    "rugby ball", "ruler measuring stick", "sneaker", "safe", "safety pin", "salt shaker", "sandal",
+    "sarong", "saxophone", "scabbard", "weighing scale", "school bus", "schooner", "scoreboard",
+    "CRT monitor", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe store",
+    "shoji screen / room divider", "shopping basket", "shopping cart", "shovel", "shower cap",
+    "shower curtain", "ski", "balaclava ski mask", "sleeping bag", "slide rule", "sliding door",
+    "slot machine", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock",
+    "solar thermal collector", "sombrero", "soup bowl", "keyboard space bar", "space heater",
+    "space shuttle", "spatula", "motorboat", "spider web", "spindle", "sports car", "spotlight",
+    "stage", "steam locomotive", "through arch bridge", "steel drum", "stethoscope", "scarf",
+    "stone wall", "stopwatch", "stove", "strainer", "tram", "stretcher", "couch", "stupa",
+    "submarine", "suit", "sundial", "sunglasses", "sunglasses", "sunscreen", "suspension bridge",
+    "mop", "sweatshirt", "swim trunks / shorts", "swing", "electrical switch", "syringe",
+    "table lamp", "tank", "tape player", "teapot", "teddy bear", "television", "tennis ball",
+    "thatched roof", "front curtain", "thimble", "threshing machine", "throne", "tile roof",
+    "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toy store",
+    "tractor", "semi-trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod",
+    "triumphal arch", "trolleybus", "trombone", "hot tub", "turnstile", "typewriter keyboard",
+    "umbrella", "unicycle", "upright piano", "vacuum cleaner", "vase", "vaulted or arched ceiling",
+    "velvet fabric", "vending machine", "vestment", "viaduct", "violin", "volleyball",
+    "waffle iron", "wall clock", "wallet", "wardrobe", "military aircraft", "sink",
+    "washing machine", "water bottle", "water jug", "water tower", "whiskey jug", "whistle",
+    "hair wig", "window screen", "window shade", "Windsor tie", "wine bottle", "airplane wing",
+    "wok", "wooden spoon", "wool", "split-rail fence", "shipwreck", "sailboat", "yurt", "website",
+    "comic book", "crossword", "traffic or street sign", "traffic light", "dust jacket", "menu",
+    "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "popsicle", "baguette",
+    "bagel", "pretzel", "cheeseburger", "hot dog", "mashed potatoes", "cabbage", "broccoli",
+    "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber",
+    "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith apple", "strawberry", "orange",
+    "lemon", "fig", "pineapple", "banana", "jackfruit", "cherimoya (custard apple)", "pomegranate",
+    "hay", "carbonara", "chocolate syrup", "dough", "meatloaf", "pizza", "pot pie", "burrito",
+    "red wine", "espresso", "tea cup", "eggnog", "mountain", "bubble", "cliff", "coral reef",
+    "geyser", "lakeshore", "promontory", "sandbar", "beach", "valley", "volcano", "baseball player",
+    "bridegroom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn",
+    "rose hip", "horse chestnut seed", "coral fungus", "agaric", "gyromitra", "stinkhorn mushroom",
+    "earth star fungus", "hen of the woods mushroom", "bolete", "corn cob", "toilet paper"
+)
+
+CLIP_PAPER_PROMPT_TEMPLATES = [
+    'a bad photo of a {}.', 'a photo of many {}.', 'a sculpture of a {}.',
+    'a photo of the hard to see {}.', 'a low resolution photo of the {}.',
+    'a rendering of a {}.', 'graffiti of a {}.', 'a bad photo of the {}.',
+    'a cropped photo of the {}.', 'a tattoo of a {}.', 'the embroidered {}.',
+    'a photo of a hard to see {}.', 'a bright photo of a {}.',
+    'a photo of a clean {}.', 'a photo of a dirty {}.', 'a dark photo of the {}.',
+    'a drawing of a {}.', 'a photo of my {}.', 'the plastic {}.',
+    'a photo of the cool {}.', 'a close-up photo of a {}.',
+    'a black and white photo of the {}.', 'a painting of the {}.',
+    'a painting of a {}.', 'a pixelated photo of the {}.', 'a sculpture of the {}.',
+    'a bright photo of the {}.', 'a cropped photo of a {}.', 'a plastic {}.',
+    'a photo of the dirty {}.', 'a jpeg corrupted photo of a {}.',
+    'a blurry photo of the {}.', 'a photo of the {}.', 'a good photo of the {}.',
+    'a rendering of the {}.', 'a {} in a video game.', 'a photo of one {}.',
+    'a doodle of a {}.', 'a close-up photo of the {}.', 'a photo of a {}.',
+    'the origami {}.', 'the {} in a video game.', 'a sketch of a {}.',
+    'a doodle of the {}.', 'a origami {}.', 'a low resolution photo of a {}.',
+    'the toy {}.', 'a rendition of the {}.', 'a photo of the clean {}.',
+    'a photo of a large {}.', 'a rendition of a {}.', 'a photo of a nice {}.',
+    'a photo of a weird {}.', 'a blurry photo of a {}.', 'a cartoon {}.',
+    'art of a {}.', 'a sketch of the {}.', 'a embroidered {}.',
+    'a pixelated photo of a {}.', 'itap of the {}.',
+    'a jpeg corrupted photo of the {}.', 'a good photo of a {}.',
+    'a plushie {}.', 'a photo of the nice {}.', 'a photo of the small {}.',
+    'a photo of the weird {}.', 'the cartoon {}.', 'art of the {}.',
+    'a drawing of the {}.', 'a photo of the large {}.',
+    'a black and white photo of a {}.', 'the plushie {}.',
+    'a dark photo of a {}.', 'itap of a {}.', 'graffiti of the {}.',
+    'a toy {}.', 'itap of my {}.', 'a photo of a cool {}.',
+    'a photo of a small {}.', 'a tattoo of the {}.', '{}',
+]
+def canonicalize_text(text: str, *, keep_punctuation_exact_string: Optional[str] = None) -> str:
+    """Canonicalize text (lowercase, remove punctuation, normalize whitespace).
+    From big_vision.evaluators.proj.image_text.prompt_engineering.
+    """
+    text = text.replace("_", " ")
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans("", "", string.punctuation))
+            for part in text.split(keep_punctuation_exact_string)
+        )
+    else:
+        text = text.translate(str.maketrans("", "", string.punctuation))
+    text = text.lower()
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+def get_canonicalized_class_names(class_names: List[str], first_only: bool = True) -> List[str]:
+    """Canonicalize class names, optionally keeping first alias only."""
+    out = []
+    for name in class_names:
+        n = canonicalize_text(name, keep_punctuation_exact_string=",")
+        if first_only:
+            n = n.split(",")[0].strip() if "," in n else n
+        out.append(n)
+    return out
+
+
+def get_canonicalized_templates(templates: List[str]) -> List[str]:
+    """Canonicalize prompt templates, preserving '{}' placeholder."""
+    return [
+        canonicalize_text(t, keep_punctuation_exact_string="{}")
+        for t in templates
+    ]
+def substitute_prompt(template: str, class_name: str) -> str:
+    """Substitute '{}' in template with class_name. From discriminative_classifier."""
+    placeholder = "{}"
+    parts = template.split(placeholder)
+    assert len(parts) == 2, f"Template must have exactly one '{placeholder}': {template}"
+    return parts[0] + class_name + parts[1]
+# ============================================================================
+# Utility functions
+# ============================================================================
+
+def batched(iterable, n: int):
+    """Batch data into lists of length n. The last batch may be shorter."""
+    it = iter(iterable)
+    while True:
+        batch = list(islice(it, n))
+        if not batch:
+            break
+        yield batch
+
+
+def accuracy(output: torch.Tensor, target: torch.Tensor, topk: Tuple[int, ...] = (1,)) -> List[float]:
+    """Calculate top-k accuracy."""
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) for k in topk]
+
+
+def get_input_dtype(precision: str) -> torch.dtype:
+    """Get input dtype based on precision string."""
+    if precision in ('bf16', 'bfloat16'):
+        return torch.bfloat16
+    elif precision in ('fp16', 'float16'):
+        return torch.float16
+    return torch.float32
+
+
+def get_autocast_context(precision: str, device_type: str = 'cuda'):
+    """Get autocast context manager based on precision."""
+    from contextlib import nullcontext
+    if precision in ('bf16', 'bfloat16'):
+        return partial(torch.amp.autocast, device_type=device_type, dtype=torch.bfloat16)
+    elif precision in ('fp16', 'float16'):
+        return partial(torch.amp.autocast, device_type=device_type, dtype=torch.float16)
+    return nullcontext
+
+
+# ============================================================================
+# Zero-shot classifier
+# ============================================================================
+
+@torch.no_grad()
+def build_zero_shot_classifier(
+    model: AutoModel,
+    processor,
+    classnames: Sequence[str],
+    templates: Sequence[Callable],
+    num_classes_per_batch: int = 10,
+    device: torch.device = None,
+    use_tqdm: bool = True,
+    canonicalize: bool = True
+
+) -> torch.Tensor:
+    """Build zero-shot classifier weights.
+
+    Args:
+        model: SigLIP2 model instance.
+        processor: SigLIP processor for text tokenization.
+        classnames: Sequence of class names.
+        templates: Sequence of template functions.
+        num_classes_per_batch: Number of classes to process per batch.
+        device: Device to use.
+        use_tqdm: Whether to show progress bar.
+
+    Returns:
+        Zero-shot classifier weights, shape (embed_dim, num_classes).
+    """
+    if canonicalize:
+        classnames = get_canonicalized_class_names(list(classnames))
+        templates = get_canonicalized_templates(list(templates))
+
+    num_templates = len(templates)
+    num_classes = len(classnames)
+    # 将classnames全部转为小写
+    classnames = [c.lower() for c in classnames]
+    if use_tqdm:
+        num_iter = (num_classes - 1) // num_classes_per_batch + 1
+        iter_wrap = partial(tqdm, total=num_iter, desc="Building classifier", unit_scale=num_classes_per_batch)
+    else:
+        iter_wrap = iter
+
+    def _process_batch(batch_classnames: List[str]) -> torch.Tensor:
+        num_batch_classes = len(batch_classnames)
+        texts = [
+            substitute_prompt(tpl, c)
+            for c in batch_classnames
+            for tpl in templates
+        ]
+        # Use processor to tokenize texts
+        inputs = processor(text=texts, return_tensors="pt", padding="max_length", max_length=64 , truncation=True)
+        # inputs = processor(text=texts, return_tensors="pt", padding="max_length", max_length=64)
+
+        input_ids = inputs["input_ids"].to(device)
+        attention_mask = inputs.get("attention_mask", None)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+
+        # Get text features from SigLIP model (with projection head)
+        with torch.no_grad():
+            text_features = model.get_text_features(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            )
+
+        # Normalize text features
+        text_features = F.normalize(text_features, dim=-1)
+
+        # Reshape and average over templates
+        text_features = text_features.reshape(num_batch_classes, num_templates, -1).mean(dim=1)
+        text_features = F.normalize(text_features, dim=1)
+        return text_features.T  # (embed_dim, num_batch_classes)
+
+    with torch.no_grad():
+        batched_embeds = [
+            _process_batch(batch)
+            for batch in iter_wrap(batched(classnames, num_classes_per_batch))
+        ]
+        zeroshot_weights = torch.cat(batched_embeds, dim=1)
+
+    return zeroshot_weights
+
+
+# ============================================================================
+# Evaluation
+# ============================================================================
+@torch.no_grad()
+def evaluate(
+    model: AutoModel,
+    processor,
+    classifier: torch.Tensor,
+    dataloader: DataLoader,
+    device: torch.device,
+    precision: str = 'fp32',
+) -> Tuple[float, float]:
+    """Run zero-shot evaluation.
+
+    Args:
+        model: SigLIP2 model instance.
+        processor: SigLIP processor for image preprocessing.
+        classifier: Zero-shot classifier weights, shape (embed_dim, num_classes).
+        dataloader: ImageNet validation dataloader.
+        device: Device to use.
+        precision: Precision for inference.
+
+    Returns:
+        Tuple of (top1_accuracy, top5_accuracy).
+    """
+    autocast_ctx = get_autocast_context(precision, device_type=device.type)
+    input_dtype = get_input_dtype(precision)
+
+    top1, top5, n = 0.0, 0.0, 0
+
+    with torch.inference_mode():
+        for images, targets in tqdm(dataloader, desc="Evaluating"):
+            images = images.to(device=device, dtype=input_dtype)
+            targets = targets.to(device)
+
+            with autocast_ctx():
+                # Get image features from SigLIP model (with projection head)
+                image_features = model.get_image_features(pixel_values=images)
+
+                # Normalize image features
+                image_features = F.normalize(image_features, dim=-1)
+                logits =  image_features @ classifier
+                logits = model.logit_scale.exp() * logits + model.logit_bias
+            acc1, acc5 = accuracy(logits, targets, topk=(1, 5))
+            top1 += acc1
+            top5 += acc5
+            n += images.size(0)
+
+    top1 = top1 / n * 100
+    top5 = top5 / n * 100
+    return top1, top5
+
+
+# ============================================================================
+# PatchReparam encoder wrapper and dataloader for training integration
+# ============================================================================
+
+
+@torch.no_grad()
+class PatchReparamZeroShotWrapper(nn.Module):
+    """Wrapper to run zero-shot eval with PatchReparam encoder (image) + full SigLIP (text + logit).
+
+    Used when integrating zero-shot into training: image features come from the
+    current PatchReparam encoder (with output_head), text features and logit scale/bias
+    from a frozen full SiglipModel.
+    """
+
+    def __init__(self, encoder: nn.Module, full_siglip: SiglipModel):
+        super().__init__()
+        self.encoder = encoder
+        self._full_siglip = full_siglip
+        # Expose logit parameters for evaluate()
+        self.logit_scale = full_siglip.logit_scale
+        self.logit_bias = full_siglip.logit_bias
+
+    def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Compute image features using PatchReparam encoder + output head (same space as SigLIP)."""
+        features = self.encoder.get_image_features(pixel_values)
+
+        return F.normalize(features, dim=-1)
+
+    def get_text_features(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Delegate to full SigLIP text tower."""
+        return self._full_siglip.get_text_features(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        )
+
+
+@torch.no_grad()
+def create_imagenet_dataloader_patch_reparam(
+    data_path: str,
+    encoder_input_size: int,
+    encoder_mean: Union[torch.Tensor, List[float]],
+    encoder_std: Union[torch.Tensor, List[float]],
+    batch_size: int = 128,
+    num_workers: int = 8,
+) -> DataLoader:
+    """Create ImageNet val dataloader with PatchReparam encoder preprocessing (resize + normalize).
+
+    Use this when evaluating PatchReparam encoder zero-shot so input matches encoder expectations.
+    """
+    if isinstance(encoder_mean, torch.Tensor):
+        mean = encoder_mean.squeeze().tolist()
+    else:
+        mean = encoder_mean
+    if isinstance(encoder_std, torch.Tensor):
+        std = encoder_std.squeeze().tolist()
+    else:
+        std = encoder_std
+    transform = transforms.Compose([
+        transforms.Resize(encoder_input_size, interpolation=transforms.InterpolationMode.BICUBIC),
+        transforms.CenterCrop(encoder_input_size),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=mean, std=std),
+    ])
+    dataset = ImageFolder(root=data_path, transform=transform)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def create_imagenet_dataloader(
+    data_path: str,
+    processor,
+    batch_size: int = 128,
+    num_workers: int = 8,
+) -> DataLoader:
+    """Create ImageNet validation dataloader using processor for preprocessing."""
+    from PIL import Image
+
+    class ImageNetDataset(ImageFolder):
+        def __getitem__(self, index):
+            path, target = self.samples[index]
+            image = Image.open(path).convert('RGB')
+            # Use processor to preprocess image
+            inputs = processor(images=image, return_tensors="pt")
+            pixel_values = inputs["pixel_values"].squeeze(0)  # Remove batch dimension
+            return pixel_values, target
+
+    dataset = ImageNetDataset(root=data_path)
+    # 验证类别顺序
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+    return dataloader
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Zero-shot ImageNet evaluation for SigLIP2")
+    parser.add_argument("--model_path", type=str, required=True,
+                        help="Path to SigLIP2 model directory")
+    parser.add_argument("--data_path", type=str, required=True,
+                        help="Path to ImageNet validation dataset")
+    parser.add_argument("--batch_size", type=int, default=128,
+                        help="Batch size for evaluation")
+    parser.add_argument("--num_workers", type=int, default=8,
+                        help="Number of dataloader workers")
+    parser.add_argument("--device", type=str, default="cuda:0",
+                        help="Device to use (e.g., cuda:0, cpu)")
+    parser.add_argument("--precision", type=str, default="fp32",
+                        choices=["fp32", "fp16", "bf16"],
+                        help="Precision for inference")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+
+    print("=" * 60)
+    print("Zero-Shot ImageNet Evaluation (SigLIP2)")
+    print("=" * 60)
+    print(f"Model path: {args.model_path}")
+    print(f"Data path: {args.data_path}")
+    print(f"Device: {device}")
+    print(f"Precision: {args.precision}")
+    print(f"Batch size: {args.batch_size}")
+    print()
+
+    # Load model and processor
+    print("Loading model and processor...")
+    # model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True)
+    model = SiglipModel.from_pretrained(args.model_path)
+    processor = AutoProcessor.from_pretrained(args.model_path)
+    model = model.to(device)
+    model.eval()
+
+    # Debug: print model structure
+    print(f"Model type: {type(model).__name__}")
+    if hasattr(model, 'vision_model'):
+        print(f"  Has vision_model: {type(model.vision_model).__name__}")
+    if hasattr(model, 'text_model'):
+        print(f"  Has text_model: {type(model.text_model).__name__}")
+    if hasattr(model, 'get_image_features'):
+        print(f"  Has get_image_features method")
+    if hasattr(model, 'get_text_features'):
+        print(f"  Has get_text_features method")
+    print()
+
+    # Get image size from processor
+    if hasattr(processor, 'image_processor'):
+        image_size = processor.image_processor.size
+        print(f"Image size: {image_size}")
+    print()
+
+    # Create dataloader
+    print(f"Loading ImageNet validation: {args.data_path}")
+    dataloader = create_imagenet_dataloader(
+        args.data_path,
+        processor=processor,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+    )
+
+    dataset = ImageFolder(root=args.data_path)
+    class_to_idx = dataset.class_to_idx  # 字典：{文件夹名: 索引}
+    sorted_classes = sorted(class_to_idx.items(), key=lambda x: x[1])  # 按索引排序
+
+    # 检查前几个类别是否匹配
+    for i, (folder_name, idx) in enumerate(sorted_classes[:10]):
+        print(f"Index {idx}: {folder_name} -> {IMAGENET_CLASSNAMES[idx]}")
+
+    num_samples = len(dataloader.dataset)
+    print(f"Number of samples: {num_samples}")
+    if num_samples != 50000:
+        print(f"  [Warning] Expected 50000 samples for ImageNet validation, got {num_samples}")
+    print()
+
+    # Build classifier
+    print("Building zero-shot classifier...")
+    autocast_ctx = get_autocast_context(args.precision, device_type=device.type)
+    with autocast_ctx():
+        classifier = build_zero_shot_classifier(
+            model,
+            processor,
+            classnames=IMAGENET_CLASSNAMES,
+            templates=CLIP_PAPER_PROMPT_TEMPLATES,
+            num_classes_per_batch=10,
+            device=device,
+            use_tqdm=True,
+            canonicalize=True,
+        )
+    print()
+
+    # Evaluate
+    print("Running evaluation...")
+    top1, top5 = evaluate(
+        model,
+        processor,
+        classifier,
+        dataloader,
+        device=device,
+        precision=args.precision,
+    )
+
+    print()
+    print("=" * 60)
+    print("Results:")
+    print(f"  ImageNet Zero-Shot Top-1 Accuracy: {top1:.2f}%")
+    print(f"  ImageNet Zero-Shot Top-5 Accuracy: {top5:.2f}%")
+    print("=" * 60)
+
+    results = {"top1": top1, "top5": top5}
+    # Write results to model_path
+    out_dir = Path(args.model_path)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_json = out_dir / "zeroshot_imagenet_results.json"
+    with open(out_json, "w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "top1": top1,
+                "top5": top5,
+                "data_path": args.data_path,
+                "precision": args.precision,
+                "batch_size": args.batch_size,
+            },
+            f,
+            indent=2,
+            ensure_ascii=False,
+        )
+    print(f"Results written to {out_json}")
+
+    return results
+
+
+if __name__ == "__main__":
+    main()

patch-reparameterization/src/eval/clip_zero_shot_bigvision.py

@@ -0,0 +1,562 @@
+#!/usr/bin/env python3
+"""
+Zero-shot ImageNet evaluation, fully aligned with big_vision discriminative_classifier.
+
+Differences from clip_zero_shot.py:
+  - Image: resize(224) | value_range(-1,1) [no mean/std]
+  - Logit: raw cosine similarity (no logit_scale, no logit_bias)
+  - Class names: canonicalize (lowercase, remove punctuation, normalize whitespace)
+  - Templates: CLIP_PAPER 77 templates with canonicalize
+  - Text: average embeddings per class, then L2 normalize
+
+Usage:
+    python clip_zero_shot_bigvision.py --model_path /path/to/model --data_path /path/to/imagenet/val
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import string
+from functools import partial
+from itertools import islice
+from pathlib import Path
+from typing import Callable, List, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from torchvision.datasets import ImageFolder
+from tqdm import tqdm
+from transformers import AutoProcessor, SiglipModel
+
+
+# ============================================================================
+# Aligned with big_vision prompt_engineering
+# ============================================================================
+
+def canonicalize_text(text: str, *, keep_punctuation_exact_string: Optional[str] = None) -> str:
+    """Canonicalize text (lowercase, remove punctuation, normalize whitespace).
+    From big_vision.evaluators.proj.image_text.prompt_engineering.
+    """
+    text = text.replace("_", " ")
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans("", "", string.punctuation))
+            for part in text.split(keep_punctuation_exact_string)
+        )
+    else:
+        text = text.translate(str.maketrans("", "", string.punctuation))
+    text = text.lower()
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+
+
+# CLIP ImageNet class names (from big_vision.datasets.imagenet.class_names)
+CLIP_IMAGENET_CLASS_NAMES = [
+    'tench', 'goldfish', 'great white shark', 'tiger shark', 'hammerhead shark',
+    'electric ray', 'stingray', 'rooster', 'hen', 'ostrich', 'brambling',
+    'goldfinch', 'house finch', 'junco', 'indigo bunting', 'American robin',
+    'bulbul', 'jay', 'magpie', 'chickadee', 'American dipper',
+    'kite (bird of prey)', 'bald eagle', 'vulture', 'great grey owl',
+    'fire salamander', 'smooth newt', 'newt', 'spotted salamander', 'axolotl',
+    'American bullfrog', 'tree frog', 'tailed frog', 'loggerhead sea turtle',
+    'leatherback sea turtle', 'mud turtle', 'terrapin', 'box turtle',
+    'banded gecko', 'green iguana', 'Carolina anole',
+    'desert grassland whiptail lizard', 'agama', 'frilled-necked lizard',
+    'alligator lizard', 'Gila monster', 'European green lizard', 'chameleon',
+    'Komodo dragon', 'Nile crocodile', 'American alligator', 'triceratops',
+    'worm snake', 'ring-necked snake', 'eastern hog-nosed snake',
+    'smooth green snake', 'kingsnake', 'garter snake', 'water snake',
+    'vine snake', 'night snake', 'boa constrictor', 'African rock python',
+    'Indian cobra', 'green mamba', 'sea snake', 'Saharan horned viper',
+    'eastern diamondback rattlesnake', 'sidewinder rattlesnake', 'trilobite',
+    'harvestman', 'scorpion', 'yellow garden spider', 'barn spider',
+    'European garden spider', 'southern black widow', 'tarantula',
+    'wolf spider', 'tick', 'centipede', 'black grouse', 'ptarmigan',
+    'ruffed grouse', 'prairie grouse', 'peafowl', 'quail', 'partridge',
+    'african grey parrot', 'macaw', 'sulphur-crested cockatoo', 'lorikeet',
+    'coucal', 'bee eater', 'hornbill', 'hummingbird', 'jacamar', 'toucan',
+    'duck', 'red-breasted merganser', 'goose', 'black swan', 'tusker',
+    'echidna', 'platypus', 'wallaby', 'koala', 'wombat', 'jellyfish',
+    'sea anemone', 'brain coral', 'flatworm', 'nematode', 'conch', 'snail',
+    'slug', 'sea slug', 'chiton', 'chambered nautilus', 'Dungeness crab',
+    'rock crab', 'fiddler crab', 'red king crab', 'American lobster',
+    'spiny lobster', 'crayfish', 'hermit crab', 'isopod', 'white stork',
+    'black stork', 'spoonbill', 'flamingo', 'little blue heron', 'great egret',
+    'bittern bird', 'crane bird', 'limpkin', 'common gallinule',
+    'American coot', 'bustard', 'ruddy turnstone', 'dunlin', 'common redshank',
+    'dowitcher', 'oystercatcher', 'pelican', 'king penguin', 'albatross',
+    'grey whale', 'killer whale', 'dugong', 'sea lion', 'Chihuahua',
+    'Japanese Chin', 'Maltese', 'Pekingese', 'Shih Tzu', 'King Charles Spaniel',
+    'Papillon', 'toy terrier', 'Rhodesian Ridgeback', 'Afghan Hound',
+    'Basset Hound', 'Beagle', 'Bloodhound', 'Bluetick Coonhound',
+    'Black and Tan Coonhound', 'Treeing Walker Coonhound', 'English foxhound',
+    'Redbone Coonhound', 'borzoi', 'Irish Wolfhound', 'Italian Greyhound',
+    'Whippet', 'Ibizan Hound', 'Norwegian Elkhound', 'Otterhound', 'Saluki',
+    'Scottish Deerhound', 'Weimaraner', 'Staffordshire Bull Terrier',
+    'American Staffordshire Terrier', 'Bedlington Terrier', 'Border Terrier',
+    'Kerry Blue Terrier', 'Irish Terrier', 'Norfolk Terrier', 'Norwich Terrier',
+    'Yorkshire Terrier', 'Wire Fox Terrier', 'Lakeland Terrier',
+    'Sealyham Terrier', 'Airedale Terrier', 'Cairn Terrier',
+    'Australian Terrier', 'Dandie Dinmont Terrier', 'Boston Terrier',
+    'Miniature Schnauzer', 'Giant Schnauzer', 'Standard Schnauzer',
+    'Scottish Terrier', 'Tibetan Terrier', 'Australian Silky Terrier',
+    'Soft-coated Wheaten Terrier', 'West Highland White Terrier', 'Lhasa Apso',
+    'Flat-Coated Retriever', 'Curly-coated Retriever', 'Golden Retriever',
+    'Labrador Retriever', 'Chesapeake Bay Retriever',
+    'German Shorthaired Pointer', 'Vizsla', 'English Setter', 'Irish Setter',
+    'Gordon Setter', 'Brittany dog', 'Clumber Spaniel',
+    'English Springer Spaniel', 'Welsh Springer Spaniel', 'Cocker Spaniel',
+    'Sussex Spaniel', 'Irish Water Spaniel', 'Kuvasz', 'Schipperke',
+    'Groenendael dog', 'Malinois', 'Briard', 'Australian Kelpie', 'Komondor',
+    'Old English Sheepdog', 'Shetland Sheepdog', 'collie', 'Border Collie',
+    'Bouvier des Flandres dog', 'Rottweiler', 'German Shepherd Dog',
+    'Dobermann', 'Miniature Pinscher', 'Greater Swiss Mountain Dog',
+    'Bernese Mountain Dog', 'Appenzeller Sennenhund', 'Entlebucher Sennenhund',
+    'Boxer', 'Bullmastiff', 'Tibetan Mastiff', 'French Bulldog', 'Great Dane',
+    'St. Bernard', 'husky', 'Alaskan Malamute', 'Siberian Husky', 'Dalmatian',
+    'Affenpinscher', 'Basenji', 'pug', 'Leonberger', 'Newfoundland dog',
+    'Great Pyrenees dog', 'Samoyed', 'Pomeranian', 'Chow Chow', 'Keeshond',
+    'brussels griffon', 'Pembroke Welsh Corgi', 'Cardigan Welsh Corgi',
+    'Toy Poodle', 'Miniature Poodle', 'Standard Poodle',
+    'Mexican hairless dog (xoloitzcuintli)', 'grey wolf', 'Alaskan tundra wolf',
+    'red wolf or maned wolf', 'coyote', 'dingo', 'dhole', 'African wild dog',
+    'hyena', 'red fox', 'kit fox', 'Arctic fox', 'grey fox', 'tabby cat',
+    'tiger cat', 'Persian cat', 'Siamese cat', 'Egyptian Mau', 'cougar', 'lynx',
+    'leopard', 'snow leopard', 'jaguar', 'lion', 'tiger', 'cheetah',
+    'brown bear', 'American black bear', 'polar bear', 'sloth bear', 'mongoose',
+    'meerkat', 'tiger beetle', 'ladybug', 'ground beetle', 'longhorn beetle',
+    'leaf beetle', 'dung beetle', 'rhinoceros beetle', 'weevil', 'fly', 'bee',
+    'ant', 'grasshopper', 'cricket insect', 'stick insect', 'cockroach',
+    'praying mantis', 'cicada', 'leafhopper', 'lacewing', 'dragonfly',
+    'damselfly', 'red admiral butterfly', 'ringlet butterfly',
+    'monarch butterfly', 'small white butterfly', 'sulphur butterfly',
+    'gossamer-winged butterfly', 'starfish', 'sea urchin', 'sea cucumber',
+    'cottontail rabbit', 'hare', 'Angora rabbit', 'hamster', 'porcupine',
+    'fox squirrel', 'marmot', 'beaver', 'guinea pig', 'common sorrel horse',
+    'zebra', 'pig', 'wild boar', 'warthog', 'hippopotamus', 'ox',
+    'water buffalo', 'bison', 'ram (adult male sheep)', 'bighorn sheep',
+    'Alpine ibex', 'hartebeest', 'impala (antelope)', 'gazelle',
+    'arabian camel', 'llama', 'weasel', 'mink', 'European polecat',
+    'black-footed ferret', 'otter', 'skunk', 'badger', 'armadillo',
+    'three-toed sloth', 'orangutan', 'gorilla', 'chimpanzee', 'gibbon',
+    'siamang', 'guenon', 'patas monkey', 'baboon', 'macaque', 'langur',
+    'black-and-white colobus', 'proboscis monkey', 'marmoset',
+    'white-headed capuchin', 'howler monkey', 'titi monkey',
+    "Geoffroy's spider monkey", 'common squirrel monkey', 'ring-tailed lemur',
+    'indri', 'Asian elephant', 'African bush elephant', 'red panda',
+    'giant panda', 'snoek fish', 'eel', 'silver salmon', 'rock beauty fish',
+    'clownfish', 'sturgeon', 'gar fish', 'lionfish', 'pufferfish', 'abacus',
+    'abaya', 'academic gown', 'accordion', 'acoustic guitar', 'aircraft carrier',
+    'airliner', 'airship', 'altar', 'ambulance', 'amphibious vehicle',
+    'analog clock', 'apiary', 'apron', 'trash can', 'assault rifle',
+    'backpack', 'bakery', 'balance beam', 'balloon', 'ballpoint pen',
+    'Band-Aid', 'banjo', 'baluster / handrail', 'barbell', 'barber chair',
+    'barbershop', 'barn', 'barometer', 'barrel', 'wheelbarrow', 'baseball',
+    'basketball', 'bassinet', 'bassoon', 'swimming cap', 'bath towel',
+    'bathtub', 'station wagon', 'lighthouse', 'beaker',
+    'military hat (bearskin or shako)', 'beer bottle', 'beer glass',
+    'bell tower', 'baby bib', 'tandem bicycle', 'bikini', 'ring binder',
+    'binoculars', 'birdhouse', 'boathouse', 'bobsleigh', 'bolo tie',
+    'poke bonnet', 'bookcase', 'bookstore', 'bottle cap', 'hunting bow',
+    'bow tie', 'brass memorial plaque', 'bra', 'breakwater', 'breastplate',
+    'broom', 'bucket', 'buckle', 'bulletproof vest', 'high-speed train',
+    'butcher shop', 'taxicab', 'cauldron', 'candle', 'cannon', 'canoe',
+    'can opener', 'cardigan', 'car mirror', 'carousel', 'tool kit',
+    'cardboard box / carton', 'car wheel', 'automated teller machine',
+    'cassette', 'cassette player', 'castle', 'catamaran', 'CD player',
+    'cello', 'mobile phone', 'chain', 'chain-link fence', 'chain mail',
+    'chainsaw', 'storage chest', 'chiffonier', 'bell or wind chime',
+    'china cabinet', 'Christmas stocking', 'church', 'movie theater',
+    'cleaver', 'cliff dwelling', 'cloak', 'clogs', 'cocktail shaker',
+    'coffee mug', 'coffeemaker', 'spiral or coil', 'combination lock',
+    'computer keyboard', 'candy store', 'container ship', 'convertible',
+    'corkscrew', 'cornet', 'cowboy boot', 'cowboy hat', 'cradle',
+    'construction crane', 'crash helmet', 'crate', 'infant bed',
+    'Crock Pot', 'croquet ball', 'crutch', 'cuirass', 'dam', 'desk',
+    'desktop computer', 'rotary dial telephone', 'diaper', 'digital clock',
+    'digital watch', 'dining table', 'dishcloth', 'dishwasher', 'disc brake',
+    'dock', 'dog sled', 'dome', 'doormat', 'drilling rig', 'drum', 'drumstick',
+    'dumbbell', 'Dutch oven', 'electric fan', 'electric guitar',
+    'electric locomotive', 'entertainment center', 'envelope',
+    'espresso machine', 'face powder', 'feather boa', 'filing cabinet',
+    'fireboat', 'fire truck', 'fire screen', 'flagpole', 'flute',
+    'folding chair', 'football helmet', 'forklift', 'fountain', 'fountain pen',
+    'four-poster bed', 'freight car', 'French horn', 'frying pan', 'fur coat',
+    'garbage truck', 'gas mask or respirator', 'gas pump', 'goblet', 'go-kart',
+    'golf ball', 'golf cart', 'gondola', 'gong', 'gown', 'grand piano',
+    'greenhouse', 'radiator grille', 'grocery store', 'guillotine',
+    'hair clip', 'hair spray', 'half-track', 'hammer', 'hamper', 'hair dryer',
+    'hand-held computer', 'handkerchief', 'hard disk drive', 'harmonica',
+    'harp', 'combine harvester', 'hatchet', 'holster', 'home theater',
+    'honeycomb', 'hook', 'hoop skirt', 'gymnastic horizontal bar',
+    'horse-drawn vehicle', 'hourglass', 'iPod', 'clothes iron',
+    'carved pumpkin', 'jeans', 'jeep', 'T-shirt', 'jigsaw puzzle', 'rickshaw',
+    'joystick', 'kimono', 'knee pad', 'knot', 'lab coat', 'ladle', 'lampshade',
+    'laptop computer', 'lawn mower', 'lens cap', 'letter opener', 'library',
+    'lifeboat', 'lighter', 'limousine', 'ocean liner', 'lipstick',
+    'slip-on shoe', 'lotion', 'music speaker', 'loupe magnifying glass',
+    'sawmill', 'magnetic compass', 'messenger bag', 'mailbox', 'tights',
+    'one-piece bathing suit', 'manhole cover', 'maraca', 'marimba', 'mask',
+    'matchstick', 'maypole', 'maze', 'measuring cup', 'medicine cabinet',
+    'megalith', 'microphone', 'microwave oven', 'military uniform', 'milk can',
+    'minibus', 'miniskirt', 'minivan', 'missile', 'mitten', 'mixing bowl',
+    'mobile home', 'ford model t', 'modem', 'monastery', 'monitor', 'moped',
+    'mortar and pestle', 'graduation cap', 'mosque', 'mosquito net', 'vespa',
+    'mountain bike', 'tent', 'computer mouse', 'mousetrap', 'moving van',
+    'muzzle', 'metal nail', 'neck brace', 'necklace', 'baby pacifier',
+    'notebook computer', 'obelisk', 'oboe', 'ocarina', 'odometer', 'oil filter',
+    'pipe organ', 'oscilloscope', 'overskirt', 'bullock cart', 'oxygen mask',
+    'product packet / packaging', 'paddle', 'paddle wheel', 'padlock',
+    'paintbrush', 'pajamas', 'palace', 'pan flute', 'paper towel', 'parachute',
+    'parallel bars', 'park bench', 'parking meter', 'railroad car', 'patio',
+    'payphone', 'pedestal', 'pencil case', 'pencil sharpener', 'perfume',
+    'Petri dish', 'photocopier', 'plectrum', 'Pickelhaube', 'picket fence',
+    'pickup truck', 'pier', 'piggy bank', 'pill bottle', 'pillow',
+    'ping-pong ball', 'pinwheel', 'pirate ship', 'drink pitcher',
+    'block plane', 'planetarium', 'plastic bag', 'plate rack', 'farm plow',
+    'plunger', 'Polaroid camera', 'pole', 'police van', 'poncho', 'pool table',
+    'soda bottle', 'plant pot', "potter's wheel", 'power drill', 'prayer rug',
+    'printer', 'prison', 'missile', 'projector', 'hockey puck', 'punching bag',
+    'purse', 'quill', 'quilt', 'race car', 'racket', 'radiator', 'radio',
+    'radio telescope', 'rain barrel', 'recreational vehicle',
+    'fishing casting reel', 'reflex camera', 'refrigerator', 'remote control',
+    'restaurant', 'revolver', 'rifle', 'rocking chair', 'rotisserie', 'eraser',
+    'rugby ball', 'ruler measuring stick', 'sneaker', 'safe', 'safety pin',
+    'salt shaker', 'sandal', 'sarong', 'saxophone', 'scabbard', 'weighing scale',
+    'school bus', 'schooner', 'scoreboard', 'CRT monitor', 'screw', 'screwdriver',
+    'seat belt', 'sewing machine', 'shield', 'shoe store',
+    'shoji screen / room divider', 'shopping basket', 'shopping cart', 'shovel',
+    'shower cap', 'shower curtain', 'ski', 'balaclava ski mask', 'sleeping bag',
+    'slide rule', 'sliding door', 'slot machine', 'snorkel', 'snowmobile',
+    'snowplow', 'soap dispenser', 'soccer ball', 'sock',
+    'solar thermal collector', 'sombrero', 'soup bowl', 'keyboard space bar',
+    'space heater', 'space shuttle', 'spatula', 'motorboat', 'spider web',
+    'spindle', 'sports car', 'spotlight', 'stage', 'steam locomotive',
+    'through arch bridge', 'steel drum', 'stethoscope', 'scarf', 'stone wall',
+    'stopwatch', 'stove', 'strainer', 'tram', 'stretcher', 'couch', 'stupa',
+    'submarine', 'suit', 'sundial', 'sunglasses', 'sunglasses', 'sunscreen',
+    'suspension bridge', 'mop', 'sweatshirt', 'swim trunks / shorts', 'swing',
+    'electrical switch', 'syringe', 'table lamp', 'tank', 'tape player',
+    'teapot', 'teddy bear', 'television', 'tennis ball', 'thatched roof',
+    'front curtain', 'thimble', 'threshing machine', 'throne', 'tile roof',
+    'toaster', 'tobacco shop', 'toilet seat', 'torch', 'totem pole', 'tow truck',
+    'toy store', 'tractor', 'semi-trailer truck', 'tray', 'trench coat',
+    'tricycle', 'trimaran', 'tripod', 'triumphal arch', 'trolleybus',
+    'trombone', 'hot tub', 'turnstile', 'typewriter keyboard', 'umbrella',
+    'unicycle', 'upright piano', 'vacuum cleaner', 'vase',
+    'vaulted or arched ceiling', 'velvet fabric', 'vending machine', 'vestment',
+    'viaduct', 'violin', 'volleyball', 'waffle iron', 'wall clock', 'wallet',
+    'wardrobe', 'military aircraft', 'sink', 'washing machine', 'water bottle',
+    'water jug', 'water tower', 'whiskey jug', 'whistle', 'hair wig',
+    'window screen', 'window shade', 'Windsor tie', 'wine bottle',
+    'airplane wing', 'wok', 'wooden spoon', 'wool', 'split-rail fence',
+    'shipwreck', 'sailboat', 'yurt', 'website', 'comic book', 'crossword',
+    'traffic or street sign', 'traffic light', 'dust jacket', 'menu', 'plate',
+    'guacamole', 'consomme', 'hot pot', 'trifle', 'ice cream', 'popsicle',
+    'baguette', 'bagel', 'pretzel', 'cheeseburger', 'hot dog',
+    'mashed potatoes', 'cabbage', 'broccoli', 'cauliflower', 'zucchini',
+    'spaghetti squash', 'acorn squash', 'butternut squash', 'cucumber',
+    'artichoke', 'bell pepper', 'cardoon', 'mushroom', 'Granny Smith apple',
+    'strawberry', 'orange', 'lemon', 'fig', 'pineapple', 'banana', 'jackfruit',
+    'cherimoya (custard apple)', 'pomegranate', 'hay', 'carbonara',
+    'chocolate syrup', 'dough', 'meatloaf', 'pizza', 'pot pie', 'burrito',
+    'red wine', 'espresso', 'tea cup', 'eggnog', 'mountain', 'bubble',
+    'cliff', 'coral reef', 'geyser', 'lakeshore', 'promontory', 'sandbar',
+    'beach', 'valley', 'volcano', 'baseball player', 'bridegroom',
+    'scuba diver', 'rapeseed', 'daisy', 'yellow lady\'s slipper', 'corn',
+    'acorn', 'rose hip', 'horse chestnut seed', 'coral fungus', 'agaric',
+    'gyromitra', 'stinkhorn mushroom', 'earth star fungus',
+    'hen of the woods mushroom', 'bolete', 'corn cob', 'toilet paper',
+]
+
+
+# CLIP_PAPER templates (from big_vision prompt_engineering_constants)
+CLIP_PAPER_PROMPT_TEMPLATES = [
+    'a bad photo of a {}.', 'a photo of many {}.', 'a sculpture of a {}.',
+    'a photo of the hard to see {}.', 'a low resolution photo of the {}.',
+    'a rendering of a {}.', 'graffiti of a {}.', 'a bad photo of the {}.',
+    'a cropped photo of the {}.', 'a tattoo of a {}.', 'the embroidered {}.',
+    'a photo of a hard to see {}.', 'a bright photo of a {}.',
+    'a photo of a clean {}.', 'a photo of a dirty {}.', 'a dark photo of the {}.',
+    'a drawing of a {}.', 'a photo of my {}.', 'the plastic {}.',
+    'a photo of the cool {}.', 'a close-up photo of a {}.',
+    'a black and white photo of the {}.', 'a painting of the {}.',
+    'a painting of a {}.', 'a pixelated photo of the {}.', 'a sculpture of the {}.',
+    'a bright photo of the {}.', 'a cropped photo of a {}.', 'a plastic {}.',
+    'a photo of the dirty {}.', 'a jpeg corrupted photo of a {}.',
+    'a blurry photo of the {}.', 'a photo of the {}.', 'a good photo of the {}.',
+    'a rendering of the {}.', 'a {} in a video game.', 'a photo of one {}.',
+    'a doodle of a {}.', 'a close-up photo of the {}.', 'a photo of a {}.',
+    'the origami {}.', 'the {} in a video game.', 'a sketch of a {}.',
+    'a doodle of the {}.', 'a origami {}.', 'a low resolution photo of a {}.',
+    'the toy {}.', 'a rendition of the {}.', 'a photo of the clean {}.',
+    'a photo of a large {}.', 'a rendition of a {}.', 'a photo of a nice {}.',
+    'a photo of a weird {}.', 'a blurry photo of a {}.', 'a cartoon {}.',
+    'art of a {}.', 'a sketch of the {}.', 'a embroidered {}.',
+    'a pixelated photo of a {}.', 'itap of the {}.',
+    'a jpeg corrupted photo of the {}.', 'a good photo of a {}.',
+    'a plushie {}.', 'a photo of the nice {}.', 'a photo of the small {}.',
+    'a photo of the weird {}.', 'the cartoon {}.', 'art of the {}.',
+    'a drawing of the {}.', 'a photo of the large {}.',
+    'a black and white photo of a {}.', 'the plushie {}.',
+    'a dark photo of a {}.', 'itap of a {}.', 'graffiti of the {}.',
+    'a toy {}.', 'itap of my {}.', 'a photo of a cool {}.',
+    'a photo of a small {}.', 'a tattoo of the {}.', '{}',
+]
+
+
+def get_canonicalized_class_names(class_names: List[str], first_only: bool = True) -> List[str]:
+    """Canonicalize class names, optionally keeping first alias only."""
+    out = []
+    for name in class_names:
+        n = canonicalize_text(name, keep_punctuation_exact_string=",")
+        if first_only:
+            n = n.split(",")[0].strip() if "," in n else n
+        out.append(n)
+    return out
+
+
+def get_canonicalized_templates(templates: List[str]) -> List[str]:
+    """Canonicalize prompt templates, preserving '{}' placeholder."""
+    return [
+        canonicalize_text(t, keep_punctuation_exact_string="{}")
+        for t in templates
+    ]
+
+
+def substitute_prompt(template: str, class_name: str) -> str:
+    """Substitute '{}' in template with class_name. From discriminative_classifier."""
+    placeholder = "{}"
+    parts = template.split(placeholder)
+    assert len(parts) == 2, f"Template must have exactly one '{placeholder}': {template}"
+    return parts[0] + class_name + parts[1]
+
+
+# ============================================================================
+# Image preprocessing: resize(224) | value_range(-1,1) [big_vision default]
+# ============================================================================
+
+def get_bigvision_imagenet_transform(image_size: int = 224):
+    """Transform matching big_vision pp_img='resize(224)|value_range(-1,1)'."""
+    return transforms.Compose([
+        transforms.Resize((image_size, image_size), interpolation=transforms.InterpolationMode.BILINEAR),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),  # [0,1] -> [-1,1]
+    ])
+
+
+# ============================================================================
+# Utility
+# ============================================================================
+
+def batched(iterable, n: int):
+    it = iter(iterable)
+    while True:
+        batch = list(islice(it, n))
+        if not batch:
+            break
+        yield batch
+
+
+def accuracy(output: torch.Tensor, target: torch.Tensor, topk: Tuple[int, ...] = (1,)) -> List[float]:
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) for k in topk]
+
+
+def get_input_dtype(precision: str) -> torch.dtype:
+    if precision in ('bf16', 'bfloat16'):
+        return torch.bfloat16
+    if precision in ('fp16', 'float16'):
+        return torch.float16
+    return torch.float32
+
+
+def get_autocast_context(precision: str, device_type: str = 'cuda'):
+    from contextlib import nullcontext
+    if precision in ('bf16', 'bfloat16'):
+        return partial(torch.amp.autocast, device_type=device_type, dtype=torch.bfloat16)
+    if precision in ('fp16', 'float16'):
+        return partial(torch.amp.autocast, device_type=device_type, dtype=torch.float16)
+    return nullcontext
+
+
+# ============================================================================
+# Zero-shot classifier (big_vision aligned: average per class, normalize)
+# ============================================================================
+
+@torch.no_grad()
+def build_zero_shot_classifier(
+    model,
+    processor,
+    classnames: Sequence[str],
+    templates: Sequence[str],
+    num_classes_per_batch: int = 10,
+    device: torch.device = None,
+    use_tqdm: bool = True,
+    canonicalize: bool = True,
+) -> torch.Tensor:
+    """Build classifier weights. Aligned with discriminative_classifier._average_embeddings."""
+    if canonicalize:
+        classnames = get_canonicalized_class_names(list(classnames))
+        templates = get_canonicalized_templates(list(templates))
+
+    num_templates = len(templates)
+    num_classes = len(classnames)
+
+    if use_tqdm:
+        num_iter = (num_classes - 1) // num_classes_per_batch + 1
+        iter_wrap = partial(tqdm, total=num_iter, desc="Building classifier", unit_scale=num_classes_per_batch)
+    else:
+        iter_wrap = iter
+
+    def _process_batch(batch_classnames: List[str]) -> torch.Tensor:
+        num_batch_classes = len(batch_classnames)
+        texts = [
+            substitute_prompt(tpl, c)
+            for c in batch_classnames
+            for tpl in templates
+        ]
+        inputs = processor(text=texts, return_tensors="pt", padding="max_length", max_length=64, truncation=True)
+        input_ids = inputs["input_ids"].to(device)
+        attention_mask = inputs.get("attention_mask")
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+
+        text_features = model.get_text_features(input_ids=input_ids, attention_mask=attention_mask)
+        text_features = F.normalize(text_features, dim=-1)
+        text_features = text_features.reshape(num_batch_classes, num_templates, -1).mean(dim=1)
+        text_features = F.normalize(text_features, dim=1)
+        return text_features.T  # (embed_dim, num_batch_classes)
+
+    batched_embeds = [
+        _process_batch(batch)
+        for batch in iter_wrap(batched(classnames, num_classes_per_batch))
+    ]
+    return torch.cat(batched_embeds, dim=1)
+
+
+# ============================================================================
+# Evaluation (big_vision aligned: raw cosine, no logit_scale/logit_bias)
+# ============================================================================
+
+@torch.no_grad()
+def evaluate(
+    model,
+    classifier: torch.Tensor,
+    dataloader: DataLoader,
+    device: torch.device,
+    precision: str = 'fp32',
+) -> Tuple[float, float]:
+    """Evaluate. Uses raw cosine similarity like discriminative_classifier."""
+    autocast_ctx = get_autocast_context(precision, device_type=device.type)
+    input_dtype = get_input_dtype(precision)
+
+    top1, top5, n = 0.0, 0.0, 0
+    with torch.inference_mode():
+        for images, targets in tqdm(dataloader, desc="Evaluating"):
+            images = images.to(device=device, dtype=input_dtype)
+            targets = targets.to(device)
+
+            with autocast_ctx():
+                image_features = model.get_image_features(pixel_values=images)
+                image_features = F.normalize(image_features, dim=-1)
+                # Raw cosine similarity (no logit_scale, no logit_bias)
+                logits = image_features @ classifier
+
+            acc1, acc5 = accuracy(logits, targets, topk=(1, 5))
+            top1 += acc1
+            top5 += acc5
+            n += images.size(0)
+
+    return top1 / n * 100, top5 / n * 100
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Zero-shot ImageNet evaluation (big_vision discriminative_classifier aligned)"
+    )
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--data_path", type=str, required=True)
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--precision", type=str, default="fp32", choices=["fp32", "fp16", "bf16"])
+    parser.add_argument("--image_size", type=int, default=224,
+                        help="Image size for resize (big_vision default 224)")
+    parser.add_argument("--no_canonicalize", action="store_true",
+                        help="Disable canonicalize of class names and templates")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+
+    print("=" * 60)
+    print("Zero-Shot ImageNet (big_vision aligned)")
+    print("=" * 60)
+    print(f"Model: {args.model_path}")
+    print(f"Data: {args.data_path}")
+    print(f"Image: resize({args.image_size})|value_range(-1,1)")
+    print(f"Logit: raw cosine (no scale/bias)")
+    print(f"Canonicalize: {not args.no_canonicalize}")
+    print()
+
+    model = SiglipModel.from_pretrained(args.model_path)
+    processor = AutoProcessor.from_pretrained(args.model_path)
+    model = model.to(device)
+    model.eval()
+
+    transform = get_bigvision_imagenet_transform(args.image_size)
+    dataset = ImageFolder(root=args.data_path, transform=transform)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=True,
+    )
+
+    print("Building zero-shot classifier...")
+    autocast_ctx = get_autocast_context(args.precision, device_type=device.type)
+    with autocast_ctx():
+        classifier = build_zero_shot_classifier(
+            model,
+            processor,
+            classnames=CLIP_IMAGENET_CLASS_NAMES,
+            templates=CLIP_PAPER_PROMPT_TEMPLATES,
+            num_classes_per_batch=10,
+            device=device,
+            use_tqdm=True,
+            canonicalize=not args.no_canonicalize,
+        )
+
+    print("Evaluating...")
+    top1, top5 = evaluate(model, classifier, dataloader, device=device, precision=args.precision)
+
+    print()
+    print("=" * 60)
+    print("Results:")
+    print(f"  Top-1: {top1:.2f}%")
+    print(f"  Top-5: {top5:.2f}%")
+    print("=" * 60)
+
+    out_dir = Path(args.model_path)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_json = out_dir / "zeroshot_imagenet_bigvision_results.json"
+    with open(out_json, "w", encoding="utf-8") as f:
+        json.dump(
+            {"top1": top1, "top5": top5, "aligned": "big_vision", "data_path": args.data_path},
+            f, indent=2, ensure_ascii=False,
+        )
+    print(f"Results written to {out_json}")
+
+
+if __name__ == "__main__":
+    main()

patch-reparameterization/src/eval/fid.py

@@ -0,0 +1,238 @@
+import argparse
+from pathlib import Path
+import numpy as np
+
+from torch_fidelity import calculate_metrics
+from .utils import ImgArrDataset
+
+from torch_fidelity.feature_extractor_inceptionv3 import FeatureExtractorInceptionV3  # original torch-fidelity :contentReference[oaicite:2]{index=2}
+import scipy.linalg
+import torch
+
+def _fid_from_moments(mu1, sigma1, mu2, sigma2) -> float:
+    # FID formula :contentReference[oaicite:3]{index=3}
+
+    mu1 = np.asarray(mu1, dtype=np.float64)
+    mu2 = np.asarray(mu2, dtype=np.float64)
+    sigma1 = np.asarray(sigma1, dtype=np.float64)
+    sigma2 = np.asarray(sigma2, dtype=np.float64)
+
+    diff = mu1 - mu2
+    covmean = scipy.linalg.sqrtm(sigma1 @ sigma2)
+
+    if np.iscomplexobj(covmean):  # numerical noise
+        covmean = covmean.real
+
+    fid = diff.dot(diff) + np.trace(sigma1 + sigma2 - 2.0 * covmean)
+    return float(max(fid, 0.0))
+
+
+@torch.no_grad()
+def _compute_inception_moments_from_arr(arr: np.ndarray, batch_size: int, device: str) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Uses torch-fidelity's InceptionV3 feature extractor to get 2048-d pool features.
+    Assumes arr is [N,H,W,C] or [N,C,H,W], uint8 (0..255) or float (0..1 or 0..255).
+    """
+    # Convert to torch in NCHW uint8 as the safest default.
+    x = arr
+    if x.ndim != 4:
+        raise ValueError(f"Expected 4D array, got shape {x.shape}")
+
+    if x.shape[-1] == 3:  # NHWC -> NCHW
+        x = np.transpose(x, (0, 3, 1, 2))
+
+    if x.dtype != np.uint8:
+        # If float in [0,1], scale up; otherwise assume already 0..255-ish
+        x_f = x.astype(np.float32)
+        if x_f.max() <= 1.5:
+            x_f = x_f * 255.0
+        x = np.clip(x_f, 0, 255).astype(np.uint8)
+
+    xt = torch.from_numpy(x).to(device=device, dtype=torch.uint8)
+
+    fe = FeatureExtractorInceptionV3(name="inception-v3-compat", features_list=['2048']).to(device).eval()  # preregistered extractor name :contentReference[oaicite:4]{index=4}
+
+    feats = []
+    for i in range(0, xt.shape[0], batch_size):
+        batch = xt[i : i + batch_size]
+        f = fe(batch)[0]              # (B, 2048)
+        feats.append(f.detach().cpu())
+
+    feats = torch.cat(feats, dim=0).double().numpy()  # (N, 2048) float64
+    mu = feats.mean(axis=0)
+    sigma = np.cov(feats, rowvar=False)
+    return mu, sigma
+
+
+@torch.no_grad()
+def compute_inception_features(arr: np.ndarray, batch_size: int, device: str) -> np.ndarray:
+    """Returns raw 2048-d InceptionV3 features for distributed FID computation."""
+    x = arr
+    if x.ndim != 4:
+        raise ValueError(f"Expected 4D array, got shape {x.shape}")
+    if x.shape[-1] == 3:  # NHWC -> NCHW
+        x = np.transpose(x, (0, 3, 1, 2))
+    if x.dtype != np.uint8:
+        x_f = x.astype(np.float32)
+        if x_f.max() <= 1.5:
+            x_f = x_f * 255.0
+        x = np.clip(x_f, 0, 255).astype(np.uint8)
+    xt = torch.from_numpy(x).to(device=device, dtype=torch.uint8)
+    fe = FeatureExtractorInceptionV3(name="inception-v3-compat", features_list=['2048']).to(device).eval()
+    feats = []
+    for i in range(0, xt.shape[0], batch_size):
+        batch = xt[i: i + batch_size]
+        f = fe(batch)[0]
+        feats.append(f.detach().cpu())
+    return torch.cat(feats, dim=0).double().numpy()  # (N, 2048)
+
+
+def calculate_rfid_from_features(features1: np.ndarray, features2: np.ndarray) -> float:
+    """Compute FID from pre-computed inception features (for distributed use)."""
+    mu1 = features1.mean(axis=0)
+    sigma1 = np.cov(features1, rowvar=False)
+    mu2 = features2.mean(axis=0)
+    sigma2 = np.cov(features2, rowvar=False)
+    return _fid_from_moments(mu1, sigma1, mu2, sigma2)
+
+
+def calculate_gfid(
+    arr1: np.ndarray,
+    ref_arr: dict,
+    batch_size: int = 64,
+    device: str = "cuda",
+) -> float:
+    mu_ref, sigma_ref = ref_arr['mu'], ref_arr['sigma']
+    mu_gen, sigma_gen = _compute_inception_moments_from_arr(arr1, batch_size=batch_size, device=device)
+    return _fid_from_moments(mu_gen, sigma_gen, mu_ref, sigma_ref)
+
+def calculate_rfid(
+    arr1,
+    arr2=None,
+    bs=64,
+    device="cuda",
+    fid_statistics_file=None,
+):
+    arr1_ds = ImgArrDataset(arr1)
+
+    if fid_statistics_file is not None:
+        metrics_kwargs = dict(
+            input1=arr1_ds,
+            input2=None,
+            fid_statistics_file=fid_statistics_file,
+            batch_size=bs,
+            fid=True,
+            cuda=(device == "cuda"),
+        )
+    else:
+        if arr2 is None:
+            raise ValueError("Either arr2 or fid_statistics_file must be provided.")
+        arr2_ds = ImgArrDataset(arr2)
+        metrics_kwargs = dict(
+            input1=arr1_ds,
+            input2=arr2_ds,
+            batch_size=bs,
+            fid=True,
+            cuda=(device == "cuda"),
+        )
+
+    metrics = calculate_metrics(**metrics_kwargs)
+    return metrics["frechet_inception_distance"]
+
+
+if __name__ == "__main__":
+    import argparse
+    from pathlib import Path
+    import numpy as np
+    import torch
+
+    parser = argparse.ArgumentParser(description="Compute FID using original torch-fidelity")
+    parser.add_argument("--arr1", type=str, required=True,
+                        help="Path to generated images array (.npy or .npz). Array should be N x H x W x 3 or N x 3 x H x W.")
+    parser.add_argument("--arr2", type=str, default=None,
+                        help="Optional path to reference images array (.npy or .npz). If set, uses torch_fidelity.calculate_metrics.")
+
+    # reference stats path (npz with mu/sigma keys)
+    parser.add_argument("--ref", "--fid-statistics-file", dest="ref", type=str, default=None,
+                        help="Path to reference stats (.npz with mu/sigma, mu_s/sigma_s, mu_clip/sigma_clip, etc.). "
+                             "If set, computes moments for arr1 using torch-fidelity InceptionV3 and FID vs these stats.")
+
+    # which keys in the ref npz to use
+    parser.add_argument("--ref-mu-key", type=str, default="mu",
+                        help="Key in --ref .npz for reference mean (default: mu).")
+    parser.add_argument("--ref-sigma-key", type=str, default="sigma",
+                        help="Key in --ref .npz for reference covariance (default: sigma).")
+
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"])
+    parser.add_argument("--no-scipy", action="store_true",
+                        help="Do not use scipy.linalg.sqrtm; use torch eig fallback (slower/less identical).")
+
+    args = parser.parse_args()
+
+    # exactly one of arr2 or ref must be provided
+    if (args.arr2 is None) == (args.ref is None):
+        parser.error("Specify exactly one of --arr2 or --ref/--fid-statistics-file")
+
+    def load_array(path: str) -> np.ndarray:
+        p = Path(path)
+        if not p.exists():
+            raise FileNotFoundError(f"File not found: {p}")
+        if p.suffix == ".npy":
+            return np.load(p)
+        elif p.suffix == ".npz":
+            z = np.load(p)
+            # common conventions: 'arr_0' or 'images'
+            if "arr_0" in z.files:
+                return z["arr_0"]
+            if "images" in z.files:
+                return z["images"]
+            raise KeyError(f"{p} is .npz but has no 'arr_0' or 'images'. Keys={z.files}")
+        else:
+            raise ValueError(f"Unsupported array file type: {p.suffix} (expected .npy or .npz)")
+
+    arr1 = load_array(args.arr1)
+    print("[INFO] arr1:", arr1.shape, arr1.dtype)
+
+    if args.arr2 is not None:
+        arr2 = load_array(args.arr2)
+        print("[INFO] arr2:", arr2.shape, arr2.dtype)
+
+        fid = calculate_rfid(
+            arr1=arr1,
+            arr2=arr2,
+            bs=args.batch_size,
+            device=args.device,
+            fid_statistics_file=None,   # unused in upstream path
+        )
+        print(f"[RESULT] FID: {fid:.6f}")
+        raise SystemExit(0)
+
+    # stats mode: arr1 vs ref npz moments
+    ref_path = Path(args.ref)
+    if not ref_path.exists():
+        raise FileNotFoundError(f"Ref stats not found: {ref_path}")
+
+    if ref_path.suffix != ".npz":
+        raise ValueError(
+            f"--ref must be a .npz containing mu/sigma. Got: {ref_path.suffix}. "
+            "Original torch-fidelity cannot use .pt stats without input2."
+        )
+
+    stats = np.load(ref_path)
+    if args.ref_mu_key not in stats.files or args.ref_sigma_key not in stats.files:
+        raise KeyError(
+            f"Missing '{args.ref_mu_key}'/'{args.ref_sigma_key}' in {ref_path}. "
+            f"Available keys: {list(stats.files)}"
+        )
+    print(f"[INFO] ref stats: {ref_path}")
+    # If you want to respect --no-scipy, thread it through:
+    # (edit your fid_arr_vs_npzstats / _fid_from_moments accordingly)
+    fid = calculate_gfid(
+        arr1=arr1,
+        ref_arr=stats,
+        batch_size=args.batch_size,
+        device=args.device,
+        # optionally add: use_scipy=(not args.no_scipy)
+    )
+    print(f"[RESULT] FID: {fid:.6f}")

patch-reparameterization/src/eval/ref_iqa.py

@@ -0,0 +1,93 @@
+from collections import namedtuple
+import numpy as np
+import torch
+import math
+from tqdm import trange
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import models
+from torchmetrics.functional.image.ssim import structural_similarity_index_measure
+import os
+import sys
+from pathlib import Path
+from typing import Optional
+from .utils import LPIPS, to_torch_tensor
+
+@torch.no_grad()
+def calculate_psnr(arr1, arr2, bs, device="cuda", disable_bar: bool = True, psnr_resize: Optional[int] = None) -> float:
+    """
+    Computes PSNR between two single images or batches of images.
+    PSNR is averaged over the batch if given as (B, C, H, W).
+
+    Args:
+        psnr_resize: If set, both images are resized to (psnr_resize x psnr_resize)
+                     before computing PSNR. Useful for evaluating at lower resolutions.
+    """
+    B = arr1.shape[0]
+    n_batches = math.ceil(B / bs)
+    psnr = torch.zeros(B, device=device)
+
+    for i in trange(n_batches, desc="Calculating PSNR", file=sys.stdout, disable=disable_bar):
+        start_idx = i * bs
+        end_idx = min((i + 1) * bs, B)
+        # PSNR expects input in [0, 1], (B, C, H, W)
+        batch_arr1 = to_torch_tensor(arr1[start_idx:end_idx]).to(device)
+        batch_arr2 = to_torch_tensor(arr2[start_idx:end_idx]).to(device)
+        if psnr_resize is not None:
+            batch_arr1 = F.interpolate(batch_arr1, size=(psnr_resize, psnr_resize), mode="bilinear", align_corners=False)
+            batch_arr2 = F.interpolate(batch_arr2, size=(psnr_resize, psnr_resize), mode="bilinear", align_corners=False)
+        batch_mse = torch.mean((batch_arr1 - batch_arr2) ** 2, dim=[1,2,3])  # shape [bs]
+        batch_mse = torch.clamp(batch_mse, min=1e-10)
+        batch_psnr = 20.0 * torch.log10(1 / torch.sqrt(batch_mse))  # shape [bs]
+        psnr[start_idx:end_idx] = batch_psnr
+    return psnr.mean().item()
+@torch.no_grad()
+def calculate_lpips(arr1, arr2, bs, device="cuda", disable_bar: bool = True) -> float:
+    """
+    Computes LPIPS between two single images or batches of images.
+    LPIPS is averaged over the batch if given as (B, C, H, W).
+    """
+    B = arr1.shape[0]
+    n_batches = math.ceil(B / bs)
+    loss_fn = LPIPS().eval().to(device)
+    lpips = torch.zeros(B, device=device)
+
+    for i in trange(n_batches, desc="Calculating LPIPS", file=sys.stdout, disable=disable_bar):
+        start_idx = i * bs
+        end_idx = min((i + 1) * bs, B)
+        # LPIPS expects input in [-1, 1], (B, C, H, W)
+        batch_arr1 = (to_torch_tensor(arr1[start_idx:end_idx]).to(device) - 0.5) * 2.
+        batch_arr2 = (to_torch_tensor(arr2[start_idx:end_idx]).to(device) - 0.5) * 2.
+        batch_lpips = loss_fn(batch_arr1, batch_arr2).squeeze()  # shape [bs]
+        lpips[start_idx:end_idx] = batch_lpips
+    return lpips.mean().item()
+
+
+######################################################
+# 4. SSIM Calculation
+######################################################
+@torch.no_grad()
+def calculate_ssim(arr1, arr2, bs, device="cuda", disable_bar: bool = True) -> float:
+    """
+    Computes SSIM between two single images or batches of images.
+    SSIM is averaged over the batch if given as (B, C, H, W).
+    """
+    B = arr1.shape[0]
+    n_batches = math.ceil(B / bs)
+    ssim_val = torch.zeros(B, device=device)
+
+    for i in trange(n_batches, desc="Calculating SSIM", file=sys.stdout, disable=disable_bar):
+        start_idx = i * bs
+        end_idx = min((i + 1) * bs, B)
+        # SSIM expects input in [0, 1], (B, C, H, W)
+        batch_arr1 = to_torch_tensor(arr1[start_idx:end_idx]).to(device)
+        batch_arr2 = to_torch_tensor(arr2[start_idx:end_idx]).to(device)
+        # SSIM expects input in [0, 1], (B, C, H, W)
+        batch_ssim = structural_similarity_index_measure(
+            target=batch_arr1,
+            preds=batch_arr2,
+            data_range=1.0,
+            reduction="none"
+        )
+        ssim_val[start_idx:end_idx] = batch_ssim
+    return ssim_val.mean().item()

patch-reparameterization/src/eval/utils.py

@@ -0,0 +1,155 @@
+from collections import namedtuple
+import numpy as np
+import torch
+import math
+import torch.nn as nn
+from torchvision import models
+from torch.utils.data import Dataset
+import os
+import sys
+from pathlib import Path
+
+# insert ../ as a module path
+current_dir = Path(__file__).resolve().parent
+parent_dir = current_dir.parent
+if str(parent_dir) not in sys.path:
+    sys.path.insert(0, str(parent_dir))
+from disc.lpips_utils import get_ckpt_path
+
+def normalize_tensor(x, eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
+    return x / (norm_factor + eps)
+
+
+def spatial_average(x, keepdim=True):
+    return x.mean([2, 3],keepdim=keepdim)
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+class VGG16(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super(VGG16, self).__init__()
+        vgg_pretrained_features = models.vgg16(weights='DEFAULT').features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout=True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = VGG16(requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self, name="vgg_lpips"):
+        ckpt = get_ckpt_path(name)
+        self.load_state_dict(
+            torch.load(
+                ckpt,
+                map_location=torch.device("cpu")),
+                strict=False
+            )
+        print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+    def forward(self, input, target):
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+class ImgArrDataset(Dataset):
+    """
+    Serve as the wrapper purpose for `torch-fidelity` to calculate the metrics (FID),
+    which expects [B, C, H, W] tensor in uint8. We are returning a [C, H, W] tensor.
+    """
+    def __init__(self, arr):
+        self.arr = arr
+
+    def __len__(self):
+        return len(self.arr)
+
+    def __getitem__(self, idx):
+        return torch.from_numpy(self.arr[idx]).permute(2, 0, 1)
+
+
+######################################################
+# 1. Convert NumPy arrays to Torch tensors, if needed
+######################################################
+def to_torch_tensor(np_array):
+    """
+    Converts a NumPy array of shape (B, H, W, C) to a (B, C, H, W) float32 Torch tensor.
+    Expects image to be in [0, 1] or [0, 255].
+    Adjust if necessary.
+    """
+    tensor = torch.from_numpy(np_array).permute(0, 3, 1, 2)
+
+    # If your data is in [0, 255], convert to [0, 1].
+    if tensor.max() > 1.0:
+        tensor = tensor.float() / 255.0
+    else:
+        tensor = tensor.float()
+    return tensor

patch-reparameterization/src/eval_only.py

@@ -0,0 +1,927 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms.
+# Licensed under the MIT license.
+"""
+Standalone evaluation script for PatchReparam models (Stage-1 / Stage-2 checkpoints).
+
+Loads a checkpoint, runs reconstruction metrics (PSNR, SSIM, rFID) and optionally
+zero-shot ImageNet classification. Supports both single-GPU and multi-GPU via
+accelerate launch.
+
+Usage:
+    # Single GPU (with checkpoint)
+    python eval_only.py --config configs/stage2/xxx.yaml --ckpt path/to/ep-last.pt
+
+    # Pretrained only (no checkpoint, model created from config)
+    python eval_only.py --config configs/stage1/pretrained/xxx.yaml --output-dir ./eval_pretrained
+
+    # Multi-GPU
+    accelerate launch --num_processes 4 eval_only.py --config configs/stage2/xxx.yaml --ckpt path/to/ep-last.pt
+
+    # Override eval data from config
+    python eval_only.py --config xxx.yaml --ckpt xxx.pt --eval-data /path/to/val/ --reference-npz /path/to/val.npz
+
+    # Save reconstructed images (PNG) to eval_results/recon_images/
+    python eval_only.py --config xxx.yaml --ckpt xxx.pt --save-recon-images
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from accelerate import Accelerator
+from omegaconf import OmegaConf
+from utils.config_utils import load_config  # noqa: F401  (registers oc.env resolver)
+from torch.cuda.amp import autocast
+from torch.utils.data import DataLoader, Subset
+from torchvision import transforms
+from torchvision.datasets import ImageFolder
+from tqdm import tqdm
+
+from eval import evaluate_reconstruction_distributed, compute_reconstruction_metrics
+from eval.clip_zero_shot import (
+    IMAGENET_CLASSNAMES,
+    CLIP_PAPER_PROMPT_TEMPLATES,
+    PatchReparamZeroShotWrapper,
+    build_zero_shot_classifier,
+    create_imagenet_dataloader_patch_reparam,
+    evaluate as evaluate_zero_shot,
+)
+from stage1 import PatchReparam
+from utils.train_utils import center_crop_arr, parse_configs
+
+
+class ExternalVAEWrapper(torch.nn.Module):
+    """Wraps diffusers AutoencoderKL to match the PatchReparam forward interface.
+
+    Input:  images in [0, 1], shape (B, 3, H, W)
+    Output: reconstructed images in [0, 1], shape (B, 3, H, W)
+    """
+
+    def __init__(self, vae):
+        super().__init__()
+        self.vae = vae
+        self.scaling_factor = vae.config.scaling_factor
+
+    @torch.no_grad()
+    def forward(self, images, **kwargs):
+        # Disable autocast: Flux VAE requires fp32 precision for accurate reconstruction
+        with torch.amp.autocast("cuda", enabled=False):
+            images = images.float()
+            # AutoencoderKL expects [-1, 1]
+            z = self.vae.encode(images * 2 - 1).latent_dist.sample() * self.scaling_factor
+            decoded = self.vae.decode(z / self.scaling_factor).sample
+            return (decoded + 1) / 2  # back to [0, 1]
+
+class Qwen3VarlenWrapper(torch.nn.Module):
+    """Wraps Qwen3Unified PatchReparam to match the standard model(images) eval interface.
+
+    Input:  images in [0, 1], shape (B, 3, 256, 256)
+    Output: reconstructed images in [0, 1], shape (B, 3, 256, 256)
+    """
+
+    def __init__(self, pr):
+        super().__init__()
+        self.pr = pr
+
+    def forward(self, images, **kwargs):
+        # images: [B, 3, H, W] in [0, 1]
+        B = images.shape[0]
+        H, W = images.shape[2], images.shape[3]
+        p  = self.pr.decoder.config.patch_size
+        tp = int(self.pr.encoder.temporal_patch_size)
+        H_grid, W_grid = H // p, W // p
+
+        # 从模型取 normalize 参数，不 hardcode
+        enc_mean = self.pr.encoder_mean.to(images.device)  # [1, 3, 1, 1]
+        enc_std  = self.pr.encoder_std.to(images.device)   # [1, 3, 1, 1]
+        frames = images.unsqueeze(1).expand(-1, tp, -1, -1, -1).contiguous()  # [B, tp, 3, H, W]
+        frames = (frames - enc_mean.unsqueeze(1)) / enc_std.unsqueeze(1)
+        frames = frames.reshape(B, tp, 3, H_grid, p, W_grid, p)
+        pixel_format = getattr(self.pr.encoder, 'pixel_format', 't_major')
+        if pixel_format == 'c_major':
+            # C-major: [B, H_g, W_g, C, tp, p, p] → [B*N, C*tp*p*p]
+            frames = frames.permute(0, 3, 5, 2, 1, 4, 6).contiguous().reshape(B * H_grid * W_grid, 3 * tp * p * p)
+        else:
+            # T-major（旧行为，兼容已有 checkpoint）
+            frames = frames.permute(0, 3, 5, 1, 2, 4, 6).contiguous().reshape(B * H_grid * W_grid * tp, 3, p, p)
+        grid_thw = torch.tensor([[1, H_grid, W_grid]], dtype=torch.long, device=images.device).expand(B, -1)
+
+        x_recs = self.pr(varlen=True, pixel_values=frames, grid_thw=grid_thw)
+        # fast path (all same res) returns tensor [B,3,H,W]; slow path returns list
+        if isinstance(x_recs, list):
+            x_recs = torch.stack(x_recs)
+        return x_recs.clamp(0, 1)
+
+
+try:
+    from peft import LoraConfig, get_peft_model, TaskType
+    PEFT_AVAILABLE = True
+except ImportError:
+    PEFT_AVAILABLE = False
+    class LoraConfig:
+        pass
+    class TaskType:
+        FEATURE_EXTRACTION = None
+
+
+def safe_load_checkpoint(path: str, map_location: str = "cpu"):
+    try:
+        return torch.load(path, map_location=map_location, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=map_location)
+
+
+def load_pr_from_checkpoint(
+    pr_config,
+    ckpt_path: str,
+    device: torch.device,
+    use_ema: bool = True,
+) -> PatchReparam:
+    """Build PatchReparam from config and load weights from checkpoint."""
+    from utils.model_utils import get_obj_from_str
+
+    cfg = OmegaConf.to_container(pr_config, resolve=True) if hasattr(pr_config, "_metadata") else dict(pr_config)
+    target = cfg.get("target")
+    params = cfg.get("params", {})
+    if not target:
+        raise KeyError("Config must have 'target' for PatchReparam instantiation.")
+
+    pr = get_obj_from_str(target)(**params)
+    pr = pr.to(device).eval()
+    pr.requires_grad_(False)
+
+    ckpt = safe_load_checkpoint(ckpt_path, map_location="cpu")
+    use_lora = ckpt.get("use_lora", False)
+    lora_config = ckpt.get("lora_config", None)
+
+    if use_lora and lora_config and PEFT_AVAILABLE:
+        if isinstance(lora_config, dict) and "default" in lora_config:
+            lora_config = lora_config["default"]
+        if isinstance(lora_config, dict):
+            lora_config = LoraConfig(
+                task_type=TaskType.FEATURE_EXTRACTION,
+                r=int(lora_config.get("r", 8)),
+                lora_alpha=int(lora_config.get("lora_alpha", 16)),
+                target_modules=lora_config.get("target_modules", ["q_proj", "v_proj", "k_proj", "out_proj"]),
+                lora_dropout=float(lora_config.get("lora_dropout", 0.1)),
+            )
+        pr.encoder = get_peft_model(pr.encoder, lora_config)
+        logging.info("Applied LoRA adapter from checkpoint.")
+
+    state_dict = ckpt.get("ema" if use_ema else "model")
+    if state_dict is None:
+        raise KeyError(f"Checkpoint must contain 'ema' or 'model'. Got keys: {list(ckpt.keys())}")
+
+    # Handle DDP/compile prefixes
+    clean_state = {}
+    for k, v in state_dict.items():
+        clean_k = k
+        for prefix in ["module._orig_mod.", "_orig_mod.", "module."]:
+            if clean_k.startswith(prefix):
+                clean_k = clean_k[len(prefix):]
+                break
+        clean_state[clean_k] = v
+
+    missing, unexpected = pr.load_state_dict(clean_state, strict=False)
+    if missing:
+        logging.warning(f"Missing keys ({len(missing)}): {missing[:5]}...")
+    if unexpected:
+        logging.warning(f"Unexpected keys ({len(unexpected)}): {unexpected[:5]}...")
+    logging.info(f"Loaded {'EMA' if use_ema else 'model'} weights from {ckpt_path}")
+    return pr
+
+
+def create_pr_from_config(pr_config, device: torch.device) -> PatchReparam:
+    """Build PatchReparam from config only (pretrained weights, no checkpoint loading)."""
+    from utils.model_utils import get_obj_from_str
+
+    cfg = OmegaConf.to_container(pr_config, resolve=True) if hasattr(pr_config, "_metadata") else dict(pr_config)
+    target = cfg.get("target")
+    params = cfg.get("params", {})
+    if not target:
+        raise KeyError("Config must have 'target' for PatchReparam instantiation.")
+
+    pr = get_obj_from_str(target)(**params)
+    pr = pr.to(device).eval()
+    pr.requires_grad_(False)
+    logging.info("Created PatchReparam from config (pretrained, no checkpoint).")
+    return pr
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Evaluate PatchReparam checkpoint (reconstruction + optional zero-shot).")
+    parser.add_argument("--config", type=str, required=True, help="YAML config with stage_1 section.")
+    parser.add_argument("--ckpt", type=str, default=None, help="Path to checkpoint (.pt). If omitted, model is created from config only (pretrained).")
+    parser.add_argument("--eval-data", type=str, default=None, help="Eval dataset path (ImageFolder). Overrides config.")
+    parser.add_argument("--reference-npz", type=str, default=None, help="Reference NPZ for metrics. Overrides config.")
+    parser.add_argument("--output-dir", type=str, default=None, help="Experiment dir (parent of checkpoints). Default: ckpt's parent's parent.")
+    parser.add_argument("--image-size", type=int, default=256, help="Image size for center crop.")
+    parser.add_argument("--batch-size", type=int, default=64, help="Batch size per GPU for reconstruction.")
+    parser.add_argument("--metrics", type=str, nargs="+", default=None, help="Metrics: psnr, ssim, rfid. Default from config.")
+    parser.add_argument("--eval-weights", type=str, choices=["both", "ema_only", "model_only"], default="ema_only",
+        help="Which weights to eval: both (default), ema_only, model_only. Overrides config eval.eval_weights.")
+    parser.add_argument("--no-zeroshot", action="store_true", help="Disable zero-shot eval even if enabled in config.")
+    parser.add_argument("--save-recon-images", action="store_true", help="Save reconstructed images as PNG files to eval_results/recon_images/.")
+    parser.add_argument("--image-features-mode", type=str, default=None,
+        help="For SigLIP2 encoder: which tokens to use for reconstruction. Overrides config eval.eval_modes when set.")
+    parser.add_argument("--precision", choices=["fp32", "fp16", "bf16"], default="bf16")
+    parser.add_argument("--encoder-cos-sim-model", type=str, default=None,
+        help="HuggingFace model path for 'cos_sim' metric. Overrides config eval.encoder_cos_sim_model.")
+    parser.add_argument("--psnr-resize", type=int, default=None,
+        help="If set, resize images to this resolution (square) before computing PSNR. Overrides config eval.psnr_resize.")
+    parser.add_argument("--l1-resize", type=int, default=None,
+        help="If set, resize images to this resolution (square) before computing L1. Overrides config eval.l1_resize.")
+    # Denoise augmentation: encode -> noise -> DiT denoise -> decode
+    parser.add_argument("--denoise-augment", action="store_true",
+        help="Enable denoise augmentation: encode GT, add flow-matching noise, DiT denoise, then decode.")
+    parser.add_argument("--dit-config", type=str, default=None,
+        help="DiT sampling config YAML (must contain stage_2, transport, misc sections).")
+    parser.add_argument("--denoise-t", type=float, default=0.1,
+        help="Noise level t for flow-matching interpolation (0~1). Default 0.1.")
+    parser.add_argument("--denoise-steps", type=int, default=1,
+        help="Number of euler steps from t to 0. Default 1.")
+    parser.add_argument("--denoise-t-list", type=str, default=None,
+        help="Comma-separated t schedule (e.g. '0.8,0.67,0.5,0.0'). Overrides --denoise-t/--denoise-steps.")
+    parser.add_argument("--denoise-shift", type=float, default=None,
+        help="Time shift value. When set, auto-compute t schedule from shifted 50-step euler, taking last --denoise-steps steps from t_start.")
+    parser.add_argument("--denoise-shift-total-steps", type=int, default=50,
+        help="Total euler steps for shift schedule computation. Default 50.")
+    parser.add_argument("--save-denoise-latents", action="store_true",
+        help="Save denoised latents (.pt) and decoded images (.png) under output-dir/denoise_latents/.")
+    parser.add_argument("--num-samples", type=int, default=None,
+        help="Limit number of eval samples. Default: use all.")
+    # Recon gFID: reconstruct train images per-class, compute gFID vs precomputed stats
+    parser.add_argument("--recon-gfid", action="store_true",
+        help="Reconstruct train images (per-class sampling) and compute gFID.")
+    parser.add_argument("--train-data", type=str,
+        default=os.environ.get("DATA_ROOT", "/path/to/ImageNet1k") + "/train/",
+        help="Train dataset path (ImageFolder) for --recon-gfid.")
+    parser.add_argument("--gfid-npz", type=str,
+        default=os.path.join(os.environ.get("FID_STATS_ROOT", "/path/to/fid_stats"), "VIRTUAL_imagenet256_labeled.npz"),
+        help="Precomputed inception stats NPZ (keys: mu, sigma) for gFID.")
+    parser.add_argument("--samples-per-class", type=int, default=50,
+        help="Number of samples per class for --recon-gfid. Default 50.")
+    return parser.parse_args()
+
+
+def sample_per_class_indices(dataset, samples_per_class):
+    """Return indices for uniform per-class sampling from an ImageFolder dataset."""
+    from collections import defaultdict
+    class_to_indices = defaultdict(list)
+    for idx, (_, label) in enumerate(dataset.samples):
+        class_to_indices[label].append(idx)
+    selected = []
+    for label in sorted(class_to_indices.keys()):
+        indices = class_to_indices[label]
+        if len(indices) >= samples_per_class:
+            selected.extend(indices[:samples_per_class])
+        else:
+            selected.extend(indices)
+    return selected
+
+
+@torch.no_grad()
+def evaluate_recon_gfid(
+    model,
+    dataset,
+    gfid_npz_path,
+    batch_size,
+    autocast_kwargs,
+    accelerator,
+    experiment_dir=None,
+    num_samples=None,
+    image_features_mode=None,
+    save_recon_dir=None,
+):
+    """Reconstruct train images and compute gFID vs precomputed inception stats."""
+    from eval.fid import calculate_gfid
+    rank = accelerator.process_index
+    world_size = accelerator.num_processes
+    device = accelerator.device
+    is_main = accelerator.is_main_process
+
+    model.eval()
+    N = num_samples if num_samples else len(dataset)
+    chunk = N // world_size
+    start = rank * chunk
+    end = (rank + 1) * chunk if rank < world_size - 1 else N
+
+    subset = Subset(dataset, list(range(start, end)))
+    loader = DataLoader(subset, batch_size=batch_size, shuffle=False,
+                        num_workers=4, pin_memory=True, drop_last=False)
+
+    reconstructions = []
+    iterator = tqdm(loader, desc=f"[Rank {rank}] Recon gFID", file=sys.stdout) if is_main else loader
+
+    with torch.inference_mode():
+        for images, _ in iterator:
+            images = images.to(device, non_blocking=True)
+            with autocast(**autocast_kwargs):
+                if image_features_mode is not None:
+                    recon = model(images, image_features_mode=image_features_mode)
+                else:
+                    recon = model(images)
+            if hasattr(recon, 'x_rec'):
+                recon = recon.x_rec
+            recon = recon.clamp(0, 1)
+            recon_np = recon.mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+            for img in recon_np:
+                reconstructions.append(img)
+
+    reconstructions = np.stack(reconstructions)
+
+    # Gather across ranks
+    temp_dir = os.path.join(experiment_dir or ".", "eval_npzs_gfid")
+    if is_main:
+        os.makedirs(temp_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    shard_path = os.path.join(temp_dir, f"recon_gfid_{rank:02d}.npz")
+    np.savez(shard_path, arr_0=reconstructions)
+    accelerator.wait_for_everyone()
+
+    metrics = None
+    if is_main:
+        all_recons = []
+        for r in range(world_size):
+            shard = np.load(os.path.join(temp_dir, f"recon_gfid_{r:02d}.npz"))["arr_0"]
+            all_recons.append(shard)
+        combined = np.concatenate(all_recons, axis=0)[:N]
+
+        ref_stats = np.load(gfid_npz_path)
+        print(f"[ReconGFID] Recon shape: {combined.shape}")
+
+        # Save reconstructed images in background thread (non-blocking, before metrics)
+        if save_recon_dir:
+            import threading
+            recons_copy = np.copy(combined)
+            def _save_images():
+                from PIL import Image as _Image
+                os.makedirs(save_recon_dir, exist_ok=True)
+                for i, img in enumerate(recons_copy):
+                    _Image.fromarray(img).save(os.path.join(save_recon_dir, f"recon_{i:06d}.png"))
+                print(f"[ReconGFID] Saved {len(recons_copy)} images to {save_recon_dir}")
+            t = threading.Thread(target=_save_images, daemon=False)
+            t.start()
+            print(f"[ReconGFID] Saving {len(combined)} images to {save_recon_dir} in background...")
+
+        device_str = "cuda" if device.type == "cuda" else "cpu"
+        gfid = calculate_gfid(combined, ref_stats, batch_size=128, device=device_str)
+        metrics = {"gfid": float(gfid)}
+        print(f"[ReconGFID] gFID: {gfid:.4f}")
+
+        # cleanup
+        for r in range(world_size):
+            p = os.path.join(temp_dir, f"recon_gfid_{r:02d}.npz")
+            if os.path.exists(p):
+                os.remove(p)
+
+    accelerator.wait_for_everyone()
+    return metrics
+
+
+def compute_denoise_t_list(t_start, denoise_steps, shift, shift_total_steps=50):
+    """Compute shifted t schedule from t_start to 0, taking last k steps from full schedule."""
+    t_uniform = 1 - torch.linspace(0, 1, shift_total_steps + 1)
+    t_shifted = shift * t_uniform / (1 + (shift - 1) * t_uniform)
+    # Find first index where t_shifted <= t_start
+    below = (t_shifted <= t_start).nonzero(as_tuple=True)[0]
+    if len(below) == 0:
+        return torch.tensor([t_start, 0.0])
+    start_idx = below[0].item()
+    # Take denoise_steps steps from start_idx, always end at 0
+    remaining = t_shifted[start_idx:]
+    if denoise_steps >= len(remaining):
+        return remaining
+    # Evenly sample denoise_steps points from remaining (including endpoints)
+    indices = torch.linspace(0, len(remaining) - 1, denoise_steps + 1).long()
+    return remaining[indices]
+
+
+def euler_denoise(x, dit_model, labels, t_start, steps, t_list=None):
+    """Euler ODE integration from t_start to 0.
+
+    Args:
+        t_list: Optional pre-computed t schedule tensor (including endpoints).
+                If provided, t_start and steps are ignored.
+    """
+    if t_list is not None:
+        ts = t_list.to(x.device)
+    else:
+        ts = torch.linspace(t_start, 0, steps + 1, device=x.device)
+    for i in range(len(ts) - 1):
+        t_cur = ts[i]
+        dt = ts[i + 1] - ts[i]  # negative (going from t to 0)
+        t_batch = torch.full((x.shape[0],), t_cur, device=x.device)
+        v = dit_model(x, t_batch, y=labels)
+        x = x + v * dt
+    return x
+
+
+@torch.no_grad()
+def evaluate_denoise_reconstruction(
+    pr,
+    dit_model,
+    dataset,
+    denoise_t,
+    denoise_steps,
+    batch_size,
+    autocast_kwargs,
+    reference_npz_path,
+    metrics_to_compute,
+    accelerator,
+    experiment_dir=None,
+    encoder_cos_sim_model=None,
+    psnr_resize=None,
+    l1_resize=None,
+    num_samples=None,
+    save_recon_dir=None,
+    t_list=None,
+    save_latents_dir=None,
+):
+    """Evaluate reconstruction with denoise augmentation (encode -> noise -> DiT denoise -> decode)."""
+    rank = accelerator.process_index
+    world_size = accelerator.num_processes
+    device = accelerator.device
+    is_main = accelerator.is_main_process
+
+    pr.eval()
+    dit_model.eval()
+
+    N = min(len(dataset), num_samples) if num_samples else len(dataset)
+    chunk = N // world_size
+    start = rank * chunk
+    end = (rank + 1) * chunk if rank < world_size - 1 else N
+
+    subset = Subset(dataset, list(range(start, end)))
+    loader = DataLoader(subset, batch_size=batch_size, shuffle=False,
+                        num_workers=4, pin_memory=True, drop_last=False)
+
+    reconstructions = []
+    iterator = tqdm(loader, desc=f"[Rank {rank}] Denoise Recon", file=sys.stdout) if is_main else loader
+
+    with torch.inference_mode():
+        for images, labels in iterator:
+            images = images.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
+            with autocast(**autocast_kwargs):
+                # 1. encode -> normalized latent
+                enc = pr.encode(images, return_aux=True)
+                z = enc.merged_tokens_norm  # [B, C, H, W]
+
+                # 2. flow-matching noise: x_t = (1-t)*data + t*noise
+                noise = torch.randn_like(z)
+                x_t = (1 - denoise_t) * z + denoise_t * noise
+
+                # 3. DiT euler denoise with GT label
+                z_hat = euler_denoise(x_t, dit_model, labels, denoise_t, denoise_steps, t_list=t_list)
+
+                # 4. decode (z_hat is [B,C,H,W] normalized, decode will denormalize)
+                recon = pr.decode(z_hat)
+
+            if hasattr(recon, 'x_rec'):
+                recon = recon.x_rec
+            recon = recon.clamp(0, 1)
+            recon_np = recon.mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+
+            # Save denoised latents, decoded images, and GT images
+            if save_latents_dir is not None:
+                from PIL import Image as _Image
+                latents_sub = os.path.join(save_latents_dir, "latents")
+                gt_sub = os.path.join(save_latents_dir, "gt")
+                os.makedirs(latents_sub, exist_ok=True)
+                os.makedirs(gt_sub, exist_ok=True)
+                gt_np = images.clamp(0, 1).mul(255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+                for j in range(z_hat.shape[0]):
+                    global_idx = start + len(reconstructions) + j
+                    torch.save(z_hat[j].cpu(), os.path.join(latents_sub, f"{global_idx:06d}.pt"))
+                    _Image.fromarray(recon_np[j]).save(os.path.join(save_latents_dir, f"{global_idx:06d}.png"))
+                    _Image.fromarray(gt_np[j]).save(os.path.join(gt_sub, f"{global_idx:06d}.png"))
+
+            for img in recon_np:
+                reconstructions.append(img)
+
+    reconstructions = np.stack(reconstructions)
+
+    # Gather across ranks via temp npz files
+    temp_dir = os.path.join(experiment_dir or ".", "eval_npzs_denoise")
+    if is_main:
+        os.makedirs(temp_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    shard_path = os.path.join(temp_dir, f"denoise_recon_{rank:02d}.npz")
+    np.savez(shard_path, arr_0=reconstructions)
+    accelerator.wait_for_everyone()
+
+    metrics = None
+    if is_main:
+        all_recons = []
+        for r in range(world_size):
+            shard = np.load(os.path.join(temp_dir, f"denoise_recon_{r:02d}.npz"))["arr_0"]
+            all_recons.append(shard)
+        combined = np.concatenate(all_recons, axis=0)[:N]
+
+        ref_images = np.load(reference_npz_path)["arr_0"][:N]
+        print(f"[DenoiseEval] Combined shape: {combined.shape}, ref shape: {ref_images.shape}")
+
+        # Save reconstructed images in background thread (non-blocking, before metrics)
+        if save_recon_dir:
+            import threading
+            recons_copy = np.copy(combined)
+            def _save_images():
+                from PIL import Image as _Image
+                os.makedirs(save_recon_dir, exist_ok=True)
+                for i, img in enumerate(recons_copy):
+                    _Image.fromarray(img).save(os.path.join(save_recon_dir, f"recon_{i:06d}.png"))
+                print(f"[DenoiseEval] Saved {len(recons_copy)} images to {save_recon_dir}")
+            t = threading.Thread(target=_save_images, daemon=False)
+            t.start()
+            print(f"[DenoiseEval] Saving {len(combined)} images to {save_recon_dir} in background...")
+
+        metrics = compute_reconstruction_metrics(
+            ref_images, combined, accelerator.device, 128,
+            metrics_to_compute=metrics_to_compute,
+            encoder_model=encoder_cos_sim_model,
+            psnr_resize=psnr_resize,
+            l1_resize=l1_resize,
+        )
+        print(f"[DenoiseEval] denoise_t={denoise_t}, steps={denoise_steps}")
+        for k, v in metrics.items():
+            print(f"  {k}: {v:.6f}")
+
+        # cleanup
+        for r in range(world_size):
+            p = os.path.join(temp_dir, f"denoise_recon_{r:02d}.npz")
+            if os.path.exists(p):
+                os.remove(p)
+
+    accelerator.wait_for_everyone()
+    return metrics
+
+
+def main():
+    args = parse_args()
+
+    mixed_precision = "fp16" if args.precision == "fp16" else ("bf16" if args.precision == "bf16" else "no")
+    accelerator = Accelerator(mixed_precision=mixed_precision)
+
+    # Only rank 0 logs; suppress other ranks to avoid duplicated output
+    root_logger = logging.getLogger()
+    if accelerator.is_main_process:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="[%(asctime)s] [%(name)s] %(levelname)s: %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+            force=True,
+        )
+    else:
+        root_logger.setLevel(logging.CRITICAL + 1)
+        for h in list(root_logger.handlers):
+            root_logger.removeHandler(h)
+    logger = logging.getLogger(__name__)
+
+    full_cfg = OmegaConf.load(args.config)
+
+    # External VAE mode (e.g. Flux/SDXL VAE): config has vae_decoder instead of stage_1
+    vae_decoder_cfg = full_cfg.get("vae_decoder", None)
+    use_external_vae = vae_decoder_cfg is not None
+
+    pr_config = None
+    if not use_external_vae:
+        pr_config, *_ = parse_configs(full_cfg)
+        if pr_config is None:
+            raise ValueError("Config must contain stage_1 or vae_decoder section.")
+
+    eval_section = full_cfg.get("eval", None)
+    eval_data = args.eval_data or (eval_section.get("data_path") if eval_section else None)
+    reference_npz = args.reference_npz or (eval_section.get("reference_npz_path") if eval_section else None)
+    eval_metrics = args.metrics or (list(eval_section.get("metrics", ["psnr", "ssim", "rfid"])) if eval_section else ["psnr", "ssim", "rfid"])
+    encoder_cos_sim_model = args.encoder_cos_sim_model or (eval_section.get("encoder_cos_sim_model") if eval_section else None)
+    psnr_resize = args.psnr_resize or (eval_section.get("psnr_resize") if eval_section else None)
+    l1_resize = args.l1_resize or (eval_section.get("l1_resize") if eval_section else None)
+    eval_weights = args.eval_weights or (eval_section.get("eval_weights", "both") if eval_section else "both")
+    if args.image_features_mode is not None:
+        eval_modes = [args.image_features_mode]
+    else:
+        eval_modes = list(eval_section.get("eval_modes", [None]) if eval_section else [None])
+    # print(eval_modes)
+    if use_external_vae:
+        models_to_eval = ["external_vae"]
+    elif args.ckpt is None:
+        models_to_eval = ["pretrained"]
+    else:
+        models_to_eval = ["ema", "model"] if eval_weights == "both" else (["ema"] if eval_weights == "ema_only" else ["model"])
+
+    if not eval_data or not reference_npz:
+        raise ValueError("Must specify --eval-data and --reference-npz, or ensure config eval.data_path and eval.reference_npz_path are set.")
+    if 'cos_sim' in eval_metrics:
+        assert encoder_cos_sim_model, (
+            "'cos_sim' metric requires --encoder-cos-sim-model or eval.encoder_cos_sim_model in config."
+        )
+
+    # exp_dir: same level as checkpoints (e.g. ckpts/stage1/exp_name/); when no ckpt, require output_dir
+    if args.ckpt is not None:
+        exp_dir = Path(args.output_dir) if args.output_dir else Path(args.ckpt).parent.parent
+    else:
+        exp_dir = Path(args.output_dir) if args.output_dir else Path("./eval_pretrained")
+    output_dir = str(exp_dir)  # for eval_npzs temp files
+    eval_results_dir = exp_dir / "eval_results"
+    ckpt_version = "external_vae" if use_external_vae else ("pretrained" if args.ckpt is None else Path(args.ckpt).stem)
+    ckpt_output_dir = eval_results_dir / ckpt_version
+    os.makedirs(ckpt_output_dir, exist_ok=True)
+
+    # Save config and startup args for reproducibility (main process only)
+    if accelerator.is_main_process:
+        config_save_path = ckpt_output_dir / "eval_config.yaml"
+        OmegaConf.save(full_cfg, config_save_path)
+        logger.info("Config saved to %s", config_save_path)
+        args_dict = vars(args)
+        # Ensure JSON-serializable (e.g. None, bool, int, str, list)
+        args_save_path = ckpt_output_dir / "eval_args.json"
+        with open(args_save_path, "w", encoding="utf-8") as f:
+            json.dump(args_dict, f, indent=2, ensure_ascii=False)
+        logger.info("Startup args saved to %s", args_save_path)
+
+    device = accelerator.device
+    eval_dataset = ImageFolder(
+        str(eval_data),
+        transform=transforms.Compose([
+            transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
+            transforms.ToTensor(),
+        ]),
+    )
+    logger.info(f"Eval dataset: {eval_data}, {len(eval_dataset)} images. Models to eval: {models_to_eval} eval_mode:{eval_modes}")
+
+    autocast_kwargs = {
+        "enabled": accelerator.mixed_precision != "no",
+        "dtype": torch.bfloat16 if accelerator.mixed_precision == "bf16" else torch.float16 if accelerator.mixed_precision == "fp16" else torch.float32,
+    }
+
+    do_zeroshot = (
+        not args.no_zeroshot
+        and eval_section
+        and eval_section.get("zero_shot", {}).get("enabled", False)
+    )
+    zs_cfg = eval_section.get("zero_shot", {}) if eval_section else {}
+    zs_data_path = str(zs_cfg.get("data_path") or zs_cfg.get("imagenet_val_path", ""))
+    zs_batch_size = int(zs_cfg.get("batch_size", 128))
+    zs_num_workers = int(zs_cfg.get("num_workers", 8))
+
+    # ── Load DiT model if denoise augmentation is enabled ──
+    dit_model = None
+    if args.denoise_augment:
+        if not args.dit_config:
+            raise ValueError("--dit-config is required when --denoise-augment is set.")
+        from utils.model_utils import instantiate_from_config as _instantiate
+        dit_cfg = OmegaConf.load(args.dit_config)
+        _, dit_model_config, _, _, _, _, _, _ = parse_configs(dit_cfg)
+        logger.info("Loading DiT model for denoise augmentation...")
+        dit_model = _instantiate(dit_model_config).to(device).eval()
+        dit_model.requires_grad_(False)
+        logger.info("DiT loaded. denoise_t=%.3f, denoise_steps=%d", args.denoise_t, args.denoise_steps)
+
+    # ── Build per-class sampled train dataset if recon-gfid is enabled ──
+    gfid_dataset = None
+    if args.recon_gfid:
+        full_train = ImageFolder(
+            str(args.train_data),
+            transform=transforms.Compose([
+                transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
+                transforms.ToTensor(),
+            ]),
+        )
+        selected_indices = sample_per_class_indices(full_train, args.samples_per_class)
+        gfid_dataset = Subset(full_train, selected_indices)
+        logger.info(f"Recon gFID: {len(gfid_dataset)} images ({args.samples_per_class}/class) from {args.train_data}")
+
+    evaluated_samples = min(len(eval_dataset), args.num_samples) if args.num_samples else len(eval_dataset)
+    results = {
+        "ckpt": args.ckpt,
+        "config": args.config,
+        "evaluation": {
+            "dataset_size": len(eval_dataset),
+            "requested_samples": args.num_samples,
+            "evaluated_samples": evaluated_samples,
+        },
+    }
+    zs_initialized = False
+
+    for model_name in models_to_eval:
+        use_ema = model_name == "ema"
+        if use_external_vae:
+            from diffusers import AutoencoderKL
+            vae_path = vae_decoder_cfg.get("path") if hasattr(vae_decoder_cfg, "get") else str(vae_decoder_cfg)
+            logger.info("Loading external VAE from %s", vae_path)
+            external_vae = AutoencoderKL.from_pretrained(vae_path).to(device).eval()
+            external_vae.requires_grad_(False)
+            logger.info("External VAE loaded: scaling_factor=%s", external_vae.config.scaling_factor)
+            pr = ExternalVAEWrapper(external_vae)
+        elif args.ckpt is None:
+            logger.info("Creating pretrained model from config...")
+            pr = create_pr_from_config(pr_config, device)
+        else:
+            logger.info("Loading %s weights...", model_name)
+            pr = load_pr_from_checkpoint(pr_config, args.ckpt, device, use_ema=use_ema)
+
+        # Qwen3Unified 用 varlen 接口，需要 wrapper 转成标准 model(images) 接口
+        try:
+            encoder_cls = OmegaConf.to_container(pr_config, resolve=True).get("params", {}).get("encoder_cls", "")
+        except Exception:
+            encoder_cls = ""
+        if encoder_cls == "Qwen3Unified":
+            pr = Qwen3VarlenWrapper(pr)
+            logger.info("Wrapped PatchReparam with Qwen3VarlenWrapper for varlen eval.")
+
+        pr = accelerator.prepare(pr)
+        eval_mod = accelerator.unwrap_model(pr)
+
+        model_results = {}
+
+        # ── Denoise augmentation mode ──
+        if args.denoise_augment and dit_model is not None:
+            # Build t_list: priority: --denoise-t-list > --denoise-shift > linspace
+            denoise_t_list = None
+            if args.denoise_t_list is not None:
+                denoise_t_list = torch.tensor([float(x) for x in args.denoise_t_list.split(",")])
+                logger.info("Using manual denoise t_list: %s", denoise_t_list.tolist())
+            elif args.denoise_shift is not None:
+                denoise_t_list = compute_denoise_t_list(
+                    args.denoise_t, args.denoise_steps, args.denoise_shift, args.denoise_shift_total_steps
+                )
+                logger.info("Using shift-computed denoise t_list (shift=%.1f, steps=%d): %s",
+                            args.denoise_shift, args.denoise_steps, denoise_t_list.tolist())
+            else:
+                logger.info("Running denoise augmentation evaluation for %s (t=%.3f, steps=%d)...",
+                            model_name, args.denoise_t, args.denoise_steps)
+            save_recon_dir = None
+            if args.save_recon_images and accelerator.is_main_process:
+                save_recon_dir = str(ckpt_output_dir / "recon_images" / f"{model_name}_denoise_t{args.denoise_t}_s{args.denoise_steps}")
+            # Use gfid_dataset (per-class train) if --recon-gfid is also set, else eval_dataset
+            denoise_dataset = gfid_dataset if (args.recon_gfid and gfid_dataset is not None) else eval_dataset
+            denoise_stats = evaluate_denoise_reconstruction(
+                pr=eval_mod,
+                dit_model=dit_model,
+                dataset=denoise_dataset,
+                denoise_t=args.denoise_t,
+                denoise_steps=args.denoise_steps,
+                batch_size=args.batch_size,
+                autocast_kwargs=autocast_kwargs,
+                reference_npz_path=args.gfid_npz if (args.recon_gfid and gfid_dataset is not None) else reference_npz,
+                metrics_to_compute=eval_metrics,
+                accelerator=accelerator,
+                experiment_dir=output_dir,
+                encoder_cos_sim_model=encoder_cos_sim_model,
+                psnr_resize=psnr_resize,
+                l1_resize=l1_resize,
+                num_samples=args.num_samples,
+                save_recon_dir=save_recon_dir,
+                t_list=denoise_t_list,
+                save_latents_dir=os.path.join(output_dir, "denoise_latents") if args.save_denoise_latents else None,
+            )
+            if denoise_stats:
+                model_results[f"denoise_t{args.denoise_t}_s{args.denoise_steps}"] = {
+                    k: float(v) for k, v in denoise_stats.items()
+                }
+            if accelerator.is_main_process and denoise_stats:
+                logger.info("%s denoise augment metrics: %s", model_name, denoise_stats)
+        elif args.recon_gfid and gfid_dataset is not None:
+            # ── Recon gFID mode: reconstruct train images, compute gFID vs precomputed stats ──
+            logger.info("Running recon gFID evaluation for %s (%d images)...", model_name, len(gfid_dataset))
+            save_recon_dir = None
+            if args.save_recon_images and accelerator.is_main_process:
+                save_recon_dir = str(ckpt_output_dir / "recon_images" / f"{model_name}_recon_gfid")
+            gfid_stats = evaluate_recon_gfid(
+                model=eval_mod,
+                dataset=gfid_dataset,
+                gfid_npz_path=args.gfid_npz,
+                batch_size=args.batch_size,
+                autocast_kwargs=autocast_kwargs,
+                accelerator=accelerator,
+                experiment_dir=output_dir,
+                num_samples=len(gfid_dataset),
+                image_features_mode=eval_modes[0] if eval_modes[0] else None,
+                save_recon_dir=save_recon_dir,
+            )
+            if gfid_stats:
+                model_results["recon_gfid"] = {k: float(v) for k, v in gfid_stats.items()}
+            if accelerator.is_main_process and gfid_stats:
+                logger.info("%s recon gFID metrics: %s", model_name, gfid_stats)
+        else:
+            # ── Normal reconstruction evaluation ──
+            logger.info("Running reconstruction evaluation for %s...", model_name)
+            for eval_mode in eval_modes:
+                print(eval_mode)
+                save_recon_dir = None
+                if args.save_recon_images and accelerator.is_main_process:
+                    subdir = model_name if eval_mode is None else f"{model_name}_{eval_mode}"
+                    save_recon_dir = str(ckpt_output_dir / "recon_images" / subdir)
+                eval_stats = evaluate_reconstruction_distributed(
+                    eval_mod,
+                    eval_dataset,
+                    args.num_samples if args.num_samples else len(eval_dataset),
+                    batch_size=args.batch_size,
+                    experiment_dir=output_dir,
+                    global_step=0,
+                    autocast_kwargs=autocast_kwargs,
+                    reference_npz_path=reference_npz,
+                    metrics_to_compute=eval_metrics,
+                    accelerator=accelerator,
+                    save_recon_dir=save_recon_dir,
+                    image_features_mode=eval_mode,
+                    encoder_cos_sim_model=encoder_cos_sim_model,
+                    psnr_resize=psnr_resize,
+                    l1_resize=l1_resize,
+                )
+                if eval_stats:
+                    recon_key = "reconstruction" if eval_mode is None else f"reconstruction_{eval_mode}"
+                    model_results[recon_key] = {k: float(v) for k, v in eval_stats.items()}
+                if accelerator.is_main_process and eval_stats:
+                    mode_tag = "" if eval_mode is None else f" [{eval_mode}]"
+                    logger.info("%s reconstruction metrics%s: %s", model_name, mode_tag, eval_stats)
+
+        if do_zeroshot and zs_data_path and accelerator.is_main_process and not use_external_vae:
+            enc_cls = eval_mod.encoder.__class__.__name__
+            if "siglip" not in enc_cls.lower():
+                logger.warning("Zero-shot skipped for %s: encoder %s is not SigLIP-based.", model_name, enc_cls)
+            else:
+                try:
+                    from transformers import AutoProcessor, SiglipModel
+                    sp = full_cfg.get("stage_1")
+                    params = getattr(sp, "params", None) if sp is not None else None
+                    enc_params = getattr(params, "encoder_params", None) if params is not None else None
+                    siglip_path = None
+                    if enc_params is not None:
+                        siglip_path = getattr(enc_params, "model_name", None) or getattr(enc_params, "sementic_model_name", None)
+                    if not siglip_path and params is not None:
+                        siglip_path = getattr(params, "encoder_config_path", None)
+                    siglip_path = str(siglip_path) if siglip_path else ""
+                    if siglip_path:
+                        if not zs_initialized:
+                            logger.info("Loading full SigLIP for zero-shot from %s", siglip_path)
+                            zs_full_siglip = SiglipModel.from_pretrained(siglip_path).to(device).eval()
+                            zs_processor = AutoProcessor.from_pretrained(siglip_path)
+                            zs_classifier = build_zero_shot_classifier(
+                                zs_full_siglip,
+                                zs_processor,
+                                classnames=IMAGENET_CLASSNAMES,
+                                templates=CLIP_PAPER_PROMPT_TEMPLATES,
+                                num_classes_per_batch=10,
+                                device=device,
+                                use_tqdm=True,
+                            )
+                            zs_dataloader = create_imagenet_dataloader_patch_reparam(
+                                zs_data_path,
+                                eval_mod.encoder_input_size,
+                                eval_mod.encoder_mean,
+                                eval_mod.encoder_std,
+                                batch_size=zs_batch_size,
+                                num_workers=zs_num_workers,
+                            )
+                            zs_initialized = True
+                        zs_wrapper = PatchReparamZeroShotWrapper(eval_mod.encoder, zs_full_siglip)
+                        zs_wrapper.eval()
+                        zs_top1, zs_top5 = evaluate_zero_shot(
+                            zs_wrapper,
+                            zs_processor,
+                            zs_classifier,
+                            zs_dataloader,
+                            device=device,
+                            precision=args.precision,
+                        )
+                        model_results["zeroshot"] = {"top1": float(zs_top1), "top5": float(zs_top5)}
+                        logger.info("%s zero-shot top1: %.2f, top5: %.2f", model_name, zs_top1, zs_top5)
+                    else:
+                        logger.warning("Zero-shot disabled: no SigLIP path in config.")
+                except Exception as e:
+                    logger.warning("Zero-shot failed for %s: %s", model_name, e)
+
+        results[model_name] = model_results
+
+    if do_zeroshot and not zs_data_path and accelerator.is_main_process:
+        logger.warning("Zero-shot enabled but data_path not set; skipping.")
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process and results:
+        result_stem = "external_vae" if use_external_vae else ("pretrained" if args.ckpt is None else Path(args.ckpt).stem)
+        result_path = ckpt_output_dir / f"{result_stem}.json"
+        suffix = 1
+        while result_path.exists():
+            result_path = ckpt_output_dir / f"{result_stem}_{suffix}.json"
+            suffix += 1
+        with open(result_path, "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2, ensure_ascii=False)
+        logger.info("Evaluation results saved to %s", result_path)
+    logger.info("Evaluation done.")
+
+
+if __name__ == "__main__":
+    main()

patch-reparameterization/src/sample_ddp.py

@@ -0,0 +1,753 @@
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Samples a large number of images from a pre-trained stage-2 model using DDP and
+stores results for downstream metrics. For single-device sampling, use sample.py.
+
+When --reference-npz is provided, FID is computed after sampling completes.
+Reference NPZ can be either:
+  - Precomputed stats (keys: mu, sigma) from Inception features
+  - Raw images (key: arr_0) - FID will be computed via torch_fidelity
+"""
+import sys
+import os
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import argparse
+import json
+import logging
+import math
+from typing import Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from PIL import Image
+from torch.cuda.amp import autocast
+from tqdm import tqdm
+from pathlib import Path
+from omegaconf import OmegaConf
+from utils.config_utils import load_config  # noqa: F401  (registers oc.env resolver)
+from utils.model_utils import instantiate_from_config
+from stage1 import PatchReparam
+from stage2.models import Stage2ModelProtocol
+from stage2.transport import create_transport, Sampler
+from utils.train_utils import parse_configs, center_crop_arr
+from eval.fid import calculate_gfid, calculate_rfid
+from torchvision.datasets import ImageFolder
+from torchvision import transforms
+
+
+def _safe_load_checkpoint(path: str, map_location: str = "cpu"):
+    try:
+        return torch.load(path, map_location=map_location, weights_only=False)
+    except TypeError:
+        # For older PyTorch versions without weights_only arg
+        return torch.load(path, map_location=map_location)
+
+
+def _load_patch_reparam_checkpoint(pr: PatchReparam, ckpt_path: str, use_ema: bool = True) -> None:
+    """
+    Load checkpoint weights into PatchReparam, choosing between EMA and model weights.
+    Mirrors the logic in eval_inversion_error._load_patch_reparam_checkpoint.
+    """
+    logger.info("Loading PatchReparam checkpoint from %s (use_ema=%s)", ckpt_path, use_ema)
+    ckpt = _safe_load_checkpoint(ckpt_path, map_location="cpu")
+    state_dict = ckpt.get("ema" if use_ema else "model")
+    if state_dict is None:
+        raise KeyError(f"Checkpoint must contain 'ema' or 'model'. Got keys: {list(ckpt.keys())}")
+
+    clean_state = {}
+    for k, v in state_dict.items():
+        clean_k = k
+        for prefix in ["module._orig_mod.", "_orig_mod.", "module."]:
+            if clean_k.startswith(prefix):
+                clean_k = clean_k[len(prefix):]
+                break
+        clean_state[clean_k] = v
+
+    missing, unexpected = pr.load_state_dict(clean_state, strict=False)
+    if missing:
+        logger.warning("PatchReparam checkpoint missing %d keys (first 5): %s...", len(missing), missing[:5])
+    if unexpected:
+        logger.warning("PatchReparam checkpoint unexpected %d keys (first 5): %s...", len(unexpected), unexpected[:5])
+    logger.info("PatchReparam checkpoint loaded successfully.")
+
+
+def create_npz_from_sample_folder(sample_dir, num=50_000):
+    """
+    Builds a single .npz file from a folder of .png samples.
+    """
+    samples = []
+    for i in tqdm(range(num), desc="Building .npz file from samples"):
+        sample_pil = Image.open(f"{sample_dir}/{i:06d}.png")
+        sample_np = np.asarray(sample_pil).astype(np.uint8)
+        samples.append(sample_np)
+    samples = np.stack(samples)
+    assert samples.shape == (num, samples.shape[1], samples.shape[2], 3)
+    npz_path = f"{sample_dir}.npz"
+    np.savez(npz_path, arr_0=samples)
+    logger.info("Saved .npz file to %s [shape=%s].", npz_path, samples.shape)
+    return npz_path
+
+
+def compute_fid_from_samples(
+    gen_npz_path: str,
+    reference_npz_path: str,
+    device: torch.device,
+    batch_size: int = 128,
+) -> float:
+    """
+    Compute FID between generated samples and reference.
+    Reference NPZ can be:
+      - Precomputed stats: keys 'mu', 'sigma' (Inception moments)
+      - Raw images: key 'arr_0'
+    """
+    gen_data = np.load(gen_npz_path)
+    gen_arr = gen_data["arr_0"]
+
+    ref_data = np.load(reference_npz_path)
+    ref_files = list(ref_data.files)
+
+    device_str = "cuda" if device.type == "cuda" else "cpu"
+
+    if "mu" in ref_files and "sigma" in ref_files:
+        fid = calculate_gfid(gen_arr, ref_data, batch_size=batch_size, device=device_str)
+    elif "arr_0" in ref_files:
+        ref_arr = ref_data["arr_0"]
+        fid = calculate_rfid(gen_arr, ref_arr, bs=batch_size, device=device_str)
+    else:
+        raise KeyError(
+            f"Reference NPZ must have 'mu'/'sigma' (precomputed stats) or 'arr_0' (raw images). "
+            f"Got keys: {ref_files}"
+        )
+    return float(fid)
+
+def build_label_sampler(
+    sampling_mode: str,
+    num_classes: int,
+    num_fid_samples: int,
+    total_samples: int,
+    samples_needed_this_device: int,
+    batch_size: int,
+    device: torch.device,
+    rank: int,
+    iterations: int,
+    seed: int,
+    eval_data: Optional[str] = None,
+    image_size: int = 256,
+) -> Callable[[int], torch.Tensor]:
+    """Create a callable that returns a batch of labels for the given step index.
+    When mode='dataset', labels follow val dataset order (same as training eval)."""
+
+    if sampling_mode == "dataset":
+        if not eval_data or not os.path.isdir(eval_data):
+            raise ValueError(
+                "dataset mode requires --eval-data (or config eval.data_path) pointing to ImageFolder val dir."
+            )
+        dataset = ImageFolder(
+            eval_data,
+            transform=transforms.Compose([
+                transforms.Lambda(lambda x: center_crop_arr(x, image_size)),
+                transforms.ToTensor(),
+            ]),
+        )
+        N = min(len(dataset), num_fid_samples)
+        world_size = total_samples // samples_needed_this_device
+        chunk = total_samples // world_size  # == samples_needed_this_device
+        start = rank * chunk
+        end = min(start + chunk, total_samples)
+        rank_indices = list(range(start, end))
+        device_pool = torch.tensor([dataset.targets[i] for i in rank_indices], dtype=torch.long)
+        needed = iterations * batch_size
+        if len(device_pool) < needed:
+            gen = torch.Generator().manual_seed(seed + rank)
+            device_pool = torch.cat([
+                device_pool,
+                torch.randint(0, num_classes, (needed - len(device_pool),), generator=gen),
+            ])
+        device_pool = device_pool[:needed].view(iterations, batch_size)
+
+        def dataset_sampler(step_idx: int) -> torch.Tensor:
+            return device_pool[step_idx].to(device)
+
+        return dataset_sampler
+
+    if sampling_mode == "random":
+        def random_sampler(_step_idx: int) -> torch.Tensor:
+            return torch.randint(0, num_classes, (batch_size,), device=device)
+
+        return random_sampler
+
+    if sampling_mode == "equal":
+        if num_fid_samples % num_classes != 0:
+            raise ValueError(
+                f"Equal label sampling requires num_fid_samples ({num_fid_samples}) to be divisible by num_classes ({num_classes})."
+            )
+
+        labels_per_class = num_fid_samples // num_classes
+        base_pool = torch.arange(num_classes, dtype=torch.long).repeat_interleave(labels_per_class)
+
+        generator = torch.Generator()
+        generator.manual_seed(seed)
+        permutation = torch.randperm(base_pool.numel(), generator=generator)
+        base_pool = base_pool[permutation]
+
+        if total_samples > num_fid_samples:
+            tail = torch.randint(0, num_classes, (total_samples - num_fid_samples,), generator=generator)
+            global_pool = torch.cat([base_pool, tail], dim=0)
+        else:
+            global_pool = base_pool
+
+        start = rank * samples_needed_this_device
+        end = start + samples_needed_this_device
+        device_pool = global_pool[start:end]
+        device_pool = device_pool.view(iterations, batch_size)
+
+        def equal_sampler(step_idx: int) -> torch.Tensor:
+            labels = device_pool[step_idx]
+            return labels.to(device)
+
+        return equal_sampler
+    raise ValueError(f"Unknown label sampling mode: {sampling_mode}")
+
+def main(args):
+    """Run sampling with distributed execution."""
+    if not torch.cuda.is_available():
+        raise RuntimeError("Sampling with DDP requires at least one GPU. Use sample.py for single-device usage.")
+
+    torch.backends.cuda.matmul.allow_tf32 = args.tf32
+    torch.backends.cudnn.allow_tf32 = args.tf32
+    torch.set_grad_enabled(False)
+
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    device_idx = rank % torch.cuda.device_count()
+    torch.cuda.set_device(device_idx)
+    device = torch.device("cuda", device_idx)
+
+    if rank == 0:
+        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    else:
+        logger.addHandler(logging.NullHandler())
+        logger.propagate = False
+
+    seed = args.global_seed * world_size + rank
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    if rank == 0:
+        logger.info("Starting rank=%s, seed=%s, world_size=%s.", rank, seed, world_size)
+
+    use_bf16 = args.precision == "bf16"
+    if use_bf16 and not torch.cuda.is_bf16_supported():
+        raise ValueError("Requested bf16 precision, but the current CUDA device does not support bfloat16.")
+    autocast_kwargs = dict(dtype=torch.bfloat16, enabled=use_bf16)
+    cfg = OmegaConf.load(args.config)
+    if args.ckpt is not None:
+        OmegaConf.update(cfg, "stage_2.ckpt", args.ckpt, force_add=True)
+    if args.cfg_scale is not None:
+        OmegaConf.update(cfg, "guidance.scale", args.cfg_scale, force_add=True)
+    if args.precomputed_latents_dir is not None:
+        val = None if args.precomputed_latents_dir.lower() == "none" else args.precomputed_latents_dir
+        OmegaConf.update(cfg, "misc.precomputed_latents_dir", val, force_add=True)
+    pr_config, model_config, transport_config, sampler_config, guidance_config, misc, _, _ = parse_configs(cfg)
+    # ── VAE mode: if vae_decoder is specified, use external VAE instead of PatchReparam ──
+    vae_decoder_cfg = cfg.get("vae_decoder", None)
+    use_external_vae = vae_decoder_cfg is not None
+
+    if not use_external_vae and (pr_config is None or model_config is None):
+        raise ValueError("Config must provide both stage_1 and stage_2 entries (or use vae_decoder).")
+    if use_external_vae and model_config is None:
+        raise ValueError("Config must provide stage_2 section.")
+
+    misc = {} if misc is None else dict(misc)
+
+    latent_size = tuple(int(dim) for dim in misc.get("latent_size", (768, 16, 16)))
+    shift_dim = misc.get("time_dist_shift_dim", math.prod(latent_size))
+    shift_base = misc.get("time_dist_shift_base", 4096)
+    time_dist_shift = math.sqrt(shift_dim / shift_base)
+    if rank == 0:
+        logger.info("Using time_dist_shift=%.4f.", time_dist_shift)
+
+    # ------------------------------------------------------------------
+    # Optional DiT latent caching / reuse
+    # ------------------------------------------------------------------
+    precomputed_latents_dir = None
+    save_latents = False
+    latents_dir = None
+    if misc is not None:
+        precomputed_latents_dir = misc.get("precomputed_latents_dir", None)
+        save_latents = bool(misc.get("save_latents", False))
+        latents_dir = misc.get("latents_dir", None)
+
+    external_vae = None
+    vae_scaling_factor = None
+    pr = None
+    proj_down_weight = None
+    proj_norm_mode = None
+    proj_norm_mean = None
+    proj_norm_std = None
+
+    if use_external_vae:
+        from diffusers import AutoencoderKL
+        vae_path = vae_decoder_cfg.get("path") if hasattr(vae_decoder_cfg, "get") else str(vae_decoder_cfg)
+        logger.info(f"Loading external VAE from {vae_path}")
+        external_vae = AutoencoderKL.from_pretrained(vae_path).to(device).eval()
+        external_vae.requires_grad_(False)
+        vae_scaling_factor = external_vae.config.scaling_factor
+        logger.info(f"External VAE loaded: scaling_factor={vae_scaling_factor}")
+
+        # Optional projection (for high-dim experiment decoding)
+        proj_cfg = vae_decoder_cfg.get("projection", None) if hasattr(vae_decoder_cfg, "get") else None
+        if proj_cfg is not None:
+            proj_cfg = OmegaConf.to_container(proj_cfg, resolve=True) if hasattr(proj_cfg, '_dict') or type(proj_cfg).__name__ == 'DictConfig' else dict(proj_cfg)
+            weight_path = proj_cfg.get("weight_path", None)
+            if not weight_path or not os.path.exists(weight_path):
+                raise FileNotFoundError(f"Projection weight_path required for sampling: {weight_path}")
+            proj_up_weight = torch.load(weight_path, map_location="cpu").to(device)
+            proj_down_weight = proj_up_weight.T.contiguous()
+            logger.info(f"Loaded projection matrix from {weight_path}: {proj_up_weight.shape}")
+
+            proj_norm_mode = proj_cfg.get("norm_mode", "global")
+            if proj_norm_mode == "per_dim":
+                norm_stat_path = proj_cfg.get("norm_stat_path", None)
+                if norm_stat_path and os.path.exists(norm_stat_path):
+                    stats = torch.load(norm_stat_path, map_location="cpu")
+                    proj_norm_mean = stats["mean"].to(device)
+                    proj_norm_std = torch.sqrt(stats["var"].to(device) + 1e-5)
+                    logger.info(f"Loaded per-dim norm stats from {norm_stat_path}")
+                else:
+                    raise FileNotFoundError(f"per_dim norm requires norm_stat_path: {norm_stat_path}")
+    else:
+        pr: PatchReparam = instantiate_from_config(pr_config).to(device)
+        pr_ckpt_cfg = pr_config.get("checkpoint") if hasattr(pr_config, "get") else None
+        if pr_ckpt_cfg:
+            if isinstance(pr_ckpt_cfg, (str, Path)):
+                ckpt_path = str(pr_ckpt_cfg)
+                use_ema = True
+            else:
+                ckpt_path = pr_ckpt_cfg.get("path") or pr_ckpt_cfg.get("ckpt_path")
+                use_ema = pr_ckpt_cfg.get("use_ema", True)
+            if ckpt_path and str(ckpt_path).strip():
+                _load_patch_reparam_checkpoint(pr, str(ckpt_path), use_ema=use_ema)
+            else:
+                logger.warning("PatchReparam checkpoint config present but path is empty.")
+        else:
+            logger.info("No PatchReparam checkpoint specified.")
+        pr.eval()
+
+    model: Stage2ModelProtocol = instantiate_from_config(model_config).to(device)
+    model.eval()
+
+    transport_params = {}
+    if transport_config is not None:
+        transport_params = dict(transport_config.get("params", {}))
+    transport = create_transport(
+        **transport_params,
+        time_dist_shift=time_dist_shift,
+    )
+    sampler = Sampler(transport)
+
+    sampler_config = {} if sampler_config is None else dict(sampler_config)
+    sampler_mode = sampler_config.get("mode", "ODE")
+    sampler_params = dict(sampler_config.get("params", {}))
+    mode = sampler_mode.upper()
+    if mode == "ODE":
+        sample_fn = sampler.sample_ode(**sampler_params)
+    elif mode == "SDE":
+        sample_fn = sampler.sample_sde(**sampler_params)
+    else:
+        raise NotImplementedError(f"Invalid sampling mode {sampler_mode}.")
+
+    guidance_config = {} if guidance_config is None else dict(guidance_config)
+
+    def guidance_value(key: str, default: float):
+        if key in guidance_config:
+            return guidance_config[key]
+        dashed_key = key.replace("_", "-")
+        return guidance_config.get(dashed_key, default)
+
+    guidance_scale = guidance_config.get("scale", 1.0)
+    guidance_method = guidance_config.get("method", "cfg")
+    t_min = guidance_value("t_min", 0.0)
+    t_max = guidance_value("t_max", 1.0)
+
+    guid_model_forward = None
+    if guidance_scale > 1.0 and guidance_method == "autoguidance":
+        guid_model_config = guidance_config.get("guidance_model")
+        if guid_model_config is None:
+            raise ValueError("Please provide a guidance model config when using autoguidance.")
+        guid_model: Stage2ModelProtocol = instantiate_from_config(guid_model_config).to(device)
+        guid_model.eval()
+        guid_model_forward = guid_model.forward
+
+    num_classes = int(misc.get("num_classes", 1000))
+    null_label = int(misc.get("null_label", num_classes))
+
+    model_target = model_config.get("target", "stage2")
+    model_string_name = str(model_target).split(".")[-1]
+    ckpt_path = model_config.get("ckpt")
+    ckpt_string_name = "pretrained" if not ckpt_path else os.path.splitext(os.path.basename(str(ckpt_path)))[0]
+    sampling_method = sampler_params.get("sampling_method", "na")
+    num_steps = sampler_params.get("num_steps", sampler_params.get("steps", "na"))
+    guidance_tag = f"cfg-{guidance_scale:.2f}"
+    base_components = [model_string_name, ckpt_string_name, guidance_tag, f"bs{args.per_proc_batch_size}"]
+    logger.info(f"args.per_proc_batch_size{args.per_proc_batch_size}")
+    if mode == "ODE":
+        detail_components = [mode, str(num_steps), str(sampling_method), args.precision]
+    else:
+        diffusion_form = sampler_params.get("diffusion_form", "na")
+        last_step = sampler_params.get("last_step", "na")
+        last_step_size = sampler_params.get("last_step_size", "na")
+        detail_components = [mode, str(num_steps), str(sampling_method), str(diffusion_form), str(last_step), str(last_step_size), args.precision]
+    folder_name = "-".join(component.replace(os.sep, "-") for component in base_components + detail_components)
+    possible_folder_name = os.environ.get('SAVE_FOLDER', None)
+    if possible_folder_name:
+        sample_folder_dir = os.path.join(args.sample_dir, possible_folder_name)
+    else:
+        sample_folder_dir = os.path.join(args.sample_dir, folder_name)
+    # Parent dir for result files (one level above image folder)
+    result_parent_dir = Path(sample_folder_dir).parent
+    folder_base = Path(sample_folder_dir).name
+
+    if rank == 0:
+        os.makedirs(sample_folder_dir, exist_ok=True)
+        logger.info("Saving .png samples at %s", sample_folder_dir)
+
+        # Determine directory to read/write DiT latents.
+        if latents_dir is None:
+            latents_dir = os.path.join(sample_folder_dir, "latents")
+
+        # Save config and startup args for reproducibility (in parent dir, not inside image folder)
+        config_save_path = result_parent_dir / f"{folder_base}_sample_config.yaml"
+        OmegaConf.save(cfg, config_save_path)
+        logger.info("Config saved to %s", config_save_path)
+        args_dict = vars(args)
+        args_save_path = result_parent_dir / f"{folder_base}_sample_args.json"
+        with open(args_save_path, "w", encoding="utf-8") as f:
+            json.dump(args_dict, f, indent=2, ensure_ascii=False)
+        logger.info("Args saved to %s", args_save_path)
+    # Broadcast latents_dir from rank 0 so all ranks agree.
+    latents_dir_obj = {"latents_dir": latents_dir}
+    obj_list = [latents_dir_obj]
+    dist.broadcast_object_list(obj_list, src=0)
+    latents_dir = obj_list[0]["latents_dir"]
+    if save_latents and latents_dir is not None and rank == 0:
+        os.makedirs(latents_dir, exist_ok=True)
+    dist.barrier()
+
+    # If npz already exists with enough samples, skip sampling to avoid redundant computation
+    npz_path = f"{sample_folder_dir}.npz"
+    skip_sampling = False
+    if rank == 0:
+        logger.info("NPZ path: %s, isfile: %s", npz_path, os.path.isfile(npz_path))
+        if os.path.isfile(npz_path):
+            try:
+                data = np.load(npz_path)
+                arr = data["arr_0"]
+                logger.info("NPZ already exists at %s with %d samples.", npz_path, arr.shape[0])
+                if arr.shape[0] >= args.num_fid_samples:
+                    skip_sampling = True
+                    logger.info("NPZ already exists at %s with %d samples, skipping sampling.", npz_path, arr.shape[0])
+            except Exception as e:
+                logger.exception("Error loading NPZ: %s", e)
+                # Fall back to existing PNG samples if they are complete enough.
+                try:
+                    existing_files = [
+                        name
+                        for name in os.listdir(sample_folder_dir)
+                        if os.path.isfile(os.path.join(sample_folder_dir, name)) and name.endswith(".png")
+                    ]
+                    existing_set = set(existing_files)
+                    expected_files = {f"{i:06d}.png" for i in range(args.num_fid_samples)}
+                    if expected_files.issubset(existing_set):
+                        logger.info(
+                            "Found %d PNG samples and a broken NPZ. Rebuilding NPZ from existing images and skipping sampling.",
+                            len(existing_files),
+                        )
+                        create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)
+                        skip_sampling = True
+                    else:
+                        logger.info(
+                            "Existing PNG samples are incomplete (%d found, need at least %d). Proceeding to resample.",
+                            len(existing_files),
+                            args.num_fid_samples,
+                        )
+                except Exception as e2:
+                    logger.exception("Error while checking/rebuilding NPZ from PNGs: %s", e2)
+    skip_tensor = torch.tensor(1 if skip_sampling else 0, device=device)
+    dist.broadcast(skip_tensor, 0)
+    skip_sampling = skip_tensor.item() == 1
+    logger.info("Skip sampling (existing images/NPZ): %s", skip_sampling)
+    n = args.per_proc_batch_size
+    global_batch_size = n * world_size
+    existing = [name for name in os.listdir(sample_folder_dir) if (os.path.isfile(os.path.join(sample_folder_dir, name)) and name.endswith(".png"))]
+    num_samples = len(existing)
+    total_samples = int(math.ceil(args.num_fid_samples / global_batch_size) * global_batch_size)
+    if rank == 0:
+        logger.info("Total number of images that will be sampled: %d", total_samples)
+    if total_samples % world_size != 0:
+        raise ValueError("Total samples must be divisible by world size.")
+    samples_needed_this_gpu = total_samples // world_size
+    if samples_needed_this_gpu % n != 0:
+        raise ValueError("Per-rank sample count must be divisible by the per-GPU batch size.")
+    iterations = samples_needed_this_gpu // n
+    # Resume: align to global_batch_size boundary to preserve label consistency
+    total = (num_samples // global_batch_size) * global_batch_size
+    start_step = total // global_batch_size
+    remaining_iterations = iterations - start_step
+    if rank == 0 and start_step > 0:
+        logger.info("Resuming from step %d (%d images exist, aligned to %d), %d steps remaining.",
+                     start_step, num_samples, total, remaining_iterations)
+    pbar = tqdm(range(remaining_iterations)) if rank == 0 else range(remaining_iterations)
+
+    eval_data = args.eval_data
+    if eval_data is None and cfg is not None:
+        eval_section = cfg.get("eval", None)
+        if eval_section is not None:
+            eval_data = eval_section.get("data_path", None)
+    label_sampler = build_label_sampler(
+        args.label_sampling,
+        num_classes,
+        args.num_fid_samples,
+        total_samples,
+        samples_needed_this_gpu,
+        n,
+        device,
+        rank,
+        iterations,
+        args.global_seed,
+        eval_data=eval_data if args.label_sampling == "dataset" else None,
+        image_size=args.image_size,
+    )
+
+    # Decide whether to reuse precomputed DiT latents (skip DiT forward) or run DiT.
+    use_precomputed_latents = False
+    # if precomputed_latents_dir is None:
+    #     # If a dedicated precomputed directory is not provided, reuse latents_dir if present.
+    #     precomputed_latents_dir = latents_dir
+    if (
+        not skip_sampling
+        and precomputed_latents_dir is not None
+        and os.path.isdir(precomputed_latents_dir)
+    ):
+        if rank == 0:
+            logger.info("Using precomputed DiT latents from %s.", precomputed_latents_dir)
+        use_precomputed_latents = True
+    use_precomputed_tensor = torch.tensor(1 if use_precomputed_latents else 0, device=device)
+    dist.broadcast(use_precomputed_tensor, 0)
+    use_precomputed_latents = use_precomputed_tensor.item() == 1
+    if rank == 0:
+        logger.info("Use precomputed DiT latents: %s", use_precomputed_latents)
+
+    using_cfg = guidance_scale > 1.0
+    if not skip_sampling:
+        for step_idx in pbar:
+            actual_step = start_step + step_idx
+            with autocast(**autocast_kwargs):
+                if use_precomputed_latents:
+                    # Load precomputed DiT latents and only run the decoder.
+                    latents = []
+                    for local_idx in range(n):
+                        index = local_idx * world_size + rank + total
+                        latent_path = os.path.join(precomputed_latents_dir, f"{index:06d}.pt")
+                        if not os.path.isfile(latent_path):
+                            raise FileNotFoundError(
+                                f"Expected latent file {latent_path} but it does not exist. "
+                                "Please ensure precomputed latents cover all required indices."
+                            )
+                        latent = torch.load(latent_path, map_location=device)
+                        latents.append(latent)
+                    latents = torch.stack(latents, dim=0).to(device=device)
+                else:
+                    # Run DiT to generate latents.
+                    z = torch.randn(n, *latent_size, device=device)
+                    y = label_sampler(actual_step)
+
+                    model_kwargs = dict(y=y)
+                    model_fn = model.forward
+
+                    if using_cfg:
+                        z = torch.cat([z, z], dim=0)
+                        y_null = torch.full((n,), null_label, device=device)
+                        y = torch.cat([y, y_null], dim=0)
+                        model_kwargs = dict(
+                            y=y,
+                            cfg_scale=guidance_scale,
+                            cfg_interval=(t_min, t_max),
+                        )
+                        if guidance_method == "autoguidance":
+                            if guid_model_forward is None:
+                                raise RuntimeError("Guidance model forward is not initialized.")
+                            model_kwargs["additional_model_forward"] = guid_model_forward
+                            model_fn = model.forward_with_autoguidance
+                        else:
+                            model_fn = model.forward_with_cfg
+                    latents = sample_fn(z, model_fn, **model_kwargs)[-1]
+                    if using_cfg:
+                        latents, _ = latents.chunk(2, dim=0)
+
+                    # Optionally cache DiT latents for future runs.
+                    if save_latents and latents_dir is not None:
+                        for local_idx, latent in enumerate(latents):
+                            index = local_idx * world_size + rank + total
+                            latent_path = os.path.join(latents_dir, f"{index:06d}.pt")
+                            torch.save(latent.to("cpu"), latent_path)
+
+                # Decode latents into images.
+                if use_external_vae:
+                    # Reverse projection if used
+                    if proj_down_weight is not None:
+                        if proj_norm_mode == "per_dim" and proj_norm_mean is not None:
+                            latents = latents * proj_norm_std + proj_norm_mean
+                        latents = torch.nn.functional.conv2d(latents, proj_down_weight[:, :, None, None])
+                    decoded = external_vae.decode(latents / vae_scaling_factor).sample
+                    decoded = ((decoded + 1) / 2).clamp(0, 1)
+                else:
+                    decoded = pr.decode(latents)
+                    mid_x_rec = getattr(decoded, 'mid_x_rec', None)
+                    if hasattr(decoded, 'x_rec'):
+                        decoded = decoded.x_rec
+                    decoded = decoded.clamp(0, 1)
+                if mid_x_rec is not None:
+                    mid_decoded = mid_x_rec.clamp(0, 1)
+                    mid_decoded = (
+                        mid_decoded.mul(255)
+                        .permute(0, 2, 3, 1)
+                        .to("cpu", dtype=torch.uint8)
+                        .numpy()
+                    )
+                else:
+                    mid_decoded = None
+                decoded = (
+                    decoded.mul(255)
+                    .permute(0, 2, 3, 1)
+                    .to("cpu", dtype=torch.uint8)
+                    .numpy()
+                )
+
+            for local_idx, sample in enumerate(decoded):
+                index = local_idx * world_size + rank + total
+                Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png")
+                if mid_decoded is not None:
+                    mid_dir = sample_folder_dir + "_mid"
+                    os.makedirs(mid_dir, exist_ok=True)
+                    Image.fromarray(mid_decoded[local_idx]).save(f"{mid_dir}/{index:06d}.png")
+
+            total += global_batch_size
+            dist.barrier()
+
+    dist.barrier()
+    if rank == 0:
+        if not skip_sampling:
+            npz_path = create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)
+        # else: npz_path already set and file exists
+
+        reference_npz = args.reference_npz
+        if reference_npz is None and cfg is not None:
+            eval_section = cfg.get("eval", None)
+            if eval_section is not None:
+                reference_npz = eval_section.get("reference_npz_path", None)
+
+        if reference_npz is not None and not args.skip_fid:
+            logger.info("Computing FID against reference...")
+            fid_value = compute_fid_from_samples(
+                npz_path,
+                reference_npz,
+                device,
+                batch_size=args.fid_batch_size,
+            )
+            logger.info("FID: %.6f", fid_value)
+
+            result_path = result_parent_dir / f"{folder_base}_fid_result.json"
+            with open(result_path, "w", encoding="utf-8") as f:
+                json.dump(
+                    {
+                        "fid": fid_value,
+                        "reference_npz": reference_npz,
+                        "num_samples": args.num_fid_samples,
+                        "sample_folder": sample_folder_dir,
+                    },
+                    f,
+                    indent=2,
+                )
+            logger.info("FID result saved to %s", result_path)
+
+        logger.info("Done.")
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, required=True, help="Path to the config file.")
+    parser.add_argument("--sample-dir", type=str, default="samples")
+    parser.add_argument("--per-proc-batch-size", type=int, default=125)
+    parser.add_argument("--num-fid-samples", type=int, default=50_000)
+    parser.add_argument("--global-seed", type=int, default=0)
+    parser.add_argument("--precision", type=str, choices=["fp32", "bf16"], default="fp32")
+    parser.add_argument("--tf32", action=argparse.BooleanOptionalAction, default=True,
+                        help="Enable TF32 matmuls (Ampere+). Disable if deterministic results are required.")
+    parser.add_argument(
+        "--label-sampling",
+        type=str,
+        choices=["random", "equal", "dataset"],
+        default="equal",
+        help="Choose how to sample class labels. Use 'dataset' to match training eval (val order).",
+    )
+    parser.add_argument(
+        "--eval-data",
+        type=str,
+        default=None,
+        help="Path to eval ImageFolder (val dir). Required for --label-sampling dataset. Overrides config eval.data_path.",
+    )
+    parser.add_argument(
+        "--image-size",
+        type=int,
+        default=256,
+        help="Image size for center crop when using dataset mode.",
+    )
+    parser.add_argument(
+        "--reference-npz",
+        type=str,
+        default=None,
+        help="Path to reference NPZ for FID computation. Can be precomputed stats (mu/sigma) or raw images (arr_0). "
+        "Overrides config eval.reference_npz_path if set.",
+    )
+    parser.add_argument(
+        "--fid-batch-size",
+        type=int,
+        default=128,
+        help="Batch size for FID computation.",
+    )
+    parser.add_argument(
+        "--skip-fid",
+        action="store_true",
+        help="Generate PNG/NPZ artifacts without computing FID in this process.",
+    )
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default=None,
+        help="Override stage_2.ckpt in config with this checkpoint path.",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=None,
+        help="Override guidance.scale in config.",
+    )
+    parser.add_argument(
+        "--precomputed-latents-dir",
+        type=str,
+        default=None,
+        help="Override misc.precomputed_latents_dir in config. Pass 'none' to disable precomputed latents.",
+    )
+
+    args = parser.parse_args()
+    main(args)

patch-reparameterization/src/stage1/encoders/__init__.py

@@ -34,7 +34,12 @@ def register_encoder(cls: Optional[Type] = None, *, name: Optional[str] = None)
 # Only encoders used by released configs are included and imported.
 from . import siglip2
 from . import siglip2_unified
+try:
+    from . import dinov2_unified
+except ImportError:
+    pass
 try:
     from . import qwen3_unified
+    from . import qwen3_frozen
 except ImportError:
     pass

patch-reparameterization/src/stage1/encoders/dinov2_unified.py

@@ -0,0 +1,124 @@
+import copy
+import logging
+
+import torch
+from torch import nn
+from transformers import Dinov2WithRegistersModel
+
+from . import register_encoder
+
+logger = logging.getLogger(__name__)
+
+
+@register_encoder()
+class DINOv2Unified(nn.Module):
+    """Single-backbone DINOv2-with-registers encoder with two independent embeddings.
+
+    Architecture:
+        image -> sem_embeddings   -> [B, 1+R+N, C] --|
+                                                       |--> cat [2B, ...] -> encoder -> split
+        image -> recon_embeddings -> [B, 1+R+N, C] --|
+
+        sem_out   (patch tokens only) -> sem_ln                 -> [B, N, C]
+        recon_out (patch tokens only) -> recon_proj -> recon_ln -> [B, N, recon_proj_dim]
+
+        output: cat([sem, recon], dim=-1)  [B, N, C + recon_proj_dim]
+
+    R = num_register_tokens (default 4), N = num_patches (default 256).
+    CLS and register tokens are discarded after the encoder.
+    """
+
+    def __init__(self, model_name: str, recon_proj_dim: int = 256, **kwargs):
+        super().__init__()
+
+        backbone = Dinov2WithRegistersModel.from_pretrained(model_name)
+        self.hidden_size = backbone.config.hidden_size
+        self.patch_size = backbone.config.patch_size
+        self.num_register_tokens = backbone.config.num_register_tokens
+        self.recon_proj_dim = recon_proj_dim
+
+        # Two independent embeddings (patch_embeddings + cls_token + register_tokens + position_embeddings)
+        self.sem_embeddings   = copy.deepcopy(backbone.embeddings)
+        self.recon_embeddings = copy.deepcopy(backbone.embeddings)
+        # mask_token is only used in MAE pre-training; remove to avoid DDP unused-parameter error
+        del self.recon_embeddings.mask_token
+        self.recon_embeddings.mask_token = None
+
+        # Shared encoder + layernorm
+        self.encoder   = backbone.encoder
+        self.layernorm = backbone.layernorm
+
+        # recon projection + LN
+        self.recon_proj = nn.Linear(self.hidden_size, recon_proj_dim, bias=True)
+        nn.init.xavier_uniform_(self.recon_proj.weight)
+        nn.init.zeros_(self.recon_proj.bias)
+
+        self.sem_ln   = nn.LayerNorm(self.hidden_size, elementwise_affine=False)
+        self.recon_ln = nn.LayerNorm(recon_proj_dim,   elementwise_affine=False)
+
+        self.output_dim = self.hidden_size + recon_proj_dim
+        self.recon_hidden_size = self.output_dim
+
+        # skip tokens: 1 CLS + num_register_tokens
+        self._skip = 1 + self.num_register_tokens
+
+        logger.info(
+            f"DINOv2Unified: hidden_size={self.hidden_size}, "
+            f"recon_proj_dim={recon_proj_dim}, output_dim={self.output_dim}, "
+            f"skip_tokens={self._skip}"
+        )
+
+    def forward(self, images: torch.Tensor, return_aux: bool = False, **kwargs):
+        B = images.shape[0]
+
+        sem_embeds   = self.sem_embeddings(images)    # [B, 1+R+N, C]
+        recon_embeds = self.recon_embeddings(images)  # [B, 1+R+N, C]
+
+        combined = torch.cat([sem_embeds, recon_embeds], dim=0)       # [2B, 1+R+N, C]
+        hidden   = self.encoder(combined).last_hidden_state            # [2B, 1+R+N, C]
+        hidden   = self.layernorm(hidden)
+
+        # discard CLS + register tokens
+        sem_out   = hidden[:B,  self._skip:]   # [B, N, C]
+        recon_out = hidden[B:,  self._skip:]   # [B, N, C]
+
+        sem_tokens   = self.sem_ln(sem_out)
+        recon_tokens = self.recon_ln(self.recon_proj(recon_out))
+
+        merged = torch.cat([sem_tokens, recon_tokens], dim=-1)
+
+        if not return_aux:
+            return merged
+
+        return merged, {
+            "sementic_tokens": sem_tokens,
+            "recon_tokens":    recon_tokens,
+            "merged_tokens":   merged,
+        }
+
+    def set_training_mode(self, mode: str = "full", **kwargs):
+        self.requires_grad_(False)
+        if mode == "full":
+            self.requires_grad_(True)
+        elif mode == "merger_recon_patchemb":
+            for p in self.recon_embeddings.parameters():
+                p.requires_grad = True
+            for p in self.recon_proj.parameters():
+                p.requires_grad = True
+        elif mode == "frozen":
+            pass
+        else:
+            raise ValueError(f"DINOv2Unified: unknown training mode '{mode}'")
+
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        total     = sum(p.numel() for p in self.parameters())
+        logger.info(
+            f"DINOv2Unified set_training_mode='{mode}': "
+            f"{trainable/1e6:.2f}M / {total/1e6:.2f}M trainable"
+        )
+
+    def get_last_layer(self) -> nn.Parameter:
+        return self.recon_proj.weight
+
+    def get_recon_last_layer(self) -> nn.Parameter:
+        return self.recon_proj.weight

patch-reparameterization/src/stage1/encoders/qwen3_frozen.py

@@ -0,0 +1,130 @@
+"""
+Qwen3Frozen: 完全冻结的 Qwen3-VL vision backbone，只做 semantic encoding。
+继承 Qwen3Unified 的所有辅助方法（RoPE、pos_embed、block_order），
+去掉 recon 分支，输出 [total_tokens, 1152] 直接送给 decoder。
+只训 decoder，encoder 完全不更新。
+"""
+
+import copy
+import logging
+
+import torch
+from torch import nn
+
+from . import register_encoder
+from .qwen3_unified import Qwen3Unified
+
+logger = logging.getLogger(__name__)
+
+
+@register_encoder()
+class Qwen3Frozen(Qwen3Unified):
+    """
+    完全冻结的单路 Qwen3-VIT encoder（无 recon 分支）。
+    继承 Qwen3Unified 的所有辅助方法，去掉双路设计。
+    输出：[total_tokens, hidden_size=1152]
+    """
+
+    def __init__(self, model_name: str, **kwargs):
+        # 不调用 Qwen3Unified.__init__，自己初始化
+        nn.Module.__init__(self)
+
+        from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+            Qwen3VLForConditionalGeneration,
+        )
+        from transformers import AutoConfig
+
+        full_cfg = AutoConfig.from_pretrained(model_name)
+        vis_cfg  = full_cfg.vision_config
+
+        full_model = Qwen3VLForConditionalGeneration.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            device_map="cpu",
+        )
+        vis_model = full_model.visual
+        del full_model.model
+
+        self.hidden_size         = vis_cfg.hidden_size
+        self.patch_size          = vis_cfg.patch_size
+        self.temporal_patch_size = vis_cfg.temporal_patch_size
+        self.spatial_merge_size  = vis_cfg.spatial_merge_size
+        self.recon_proj_dim      = 0  # no recon branch
+
+        # single patch embed + pos embed (no deepcopy needed, frozen anyway)
+        self.sem_patch_embed   = vis_model.patch_embed
+        self.sem_pos_embed     = vis_model.pos_embed
+        self.num_grid_per_side = vis_model.num_grid_per_side
+        self.rotary_pos_emb    = vis_model.rotary_pos_emb
+        self.blocks            = vis_model.blocks
+
+        self.output_dim        = self.hidden_size
+        self.recon_hidden_size = self.hidden_size  # decoder latent_dim
+
+        # PatchReparam reads these to skip internal resize/normalize
+        self.handles_preprocessing = True
+        self.encoder_mean = torch.tensor([0.5, 0.5, 0.5]).view(1, 3, 1, 1)
+        self.encoder_std  = torch.tensor([0.5, 0.5, 0.5]).view(1, 3, 1, 1)
+        self.encoder_input_size = 256
+
+        # caches (inherited methods use these)
+        self._pos_embed_cache:   dict = {}
+        self._rot_pos_emb_cache: dict = {}
+        self._block_idx_cache:   dict = {}
+
+        # 完全冻结
+        self.requires_grad_(False)
+
+        del vis_model
+
+        logger.info(f"Qwen3Frozen: hidden_size={self.hidden_size}, fully frozen")
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thw: torch.Tensor,
+        return_aux: bool = False,
+        **kwargs,
+    ):
+        # patch embed + pos embed (single path)
+        x = self.sem_patch_embed(pixel_values)
+        pos = self._abs_pos_embed(self.sem_pos_embed, grid_thw)
+
+        # reorder to 2x2 block order
+        perm = self._block_order_indices(grid_thw, device=x.device)
+        inv_perm = torch.empty_like(perm)
+        inv_perm[perm] = torch.arange(perm.numel(), device=perm.device)
+        x = (x + pos)[perm]
+
+        # RoPE
+        rotary = self._rot_pos_emb(grid_thw)
+        emb    = torch.cat((rotary, rotary), dim=-1)
+        position_embeddings = (emb.cos(), emb.sin())
+
+        # cu_seqlens
+        cu_seqlens = torch.repeat_interleave(
+            grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
+        ).cumsum(0, dtype=torch.int32)
+        cu_seqlens = torch.nn.functional.pad(cu_seqlens, (1, 0), value=0)
+
+        # blocks forward
+        for blk in self.blocks:
+            x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings)
+
+        # restore row-major
+        x = x[inv_perm]
+
+        if not return_aux:
+            return x
+
+        return x, {"sementic_tokens": x, "merged_tokens": x}
+
+    def set_training_mode(self, mode: str = "frozen", **kwargs):
+        self.requires_grad_(False)
+        logger.info(f"Qwen3Frozen: always frozen, ignoring mode='{mode}'")
+
+    def get_last_layer(self):
+        return None
+
+    def get_recon_last_layer(self):
+        return None

patch-reparameterization/src/stage1/encoders/qwen3_unified.py

@@ -67,16 +67,22 @@ class Qwen3Unified(nn.Module):
     ):
         super().__init__()
 
-        # Load vision backbone config only (no LLM weights download needed —
-        # the pr-qwen-vit-tokenizer checkpoint provides all trained weights)
+        # Load full Qwen3VL model to extract the vision backbone, then discard the LLM
         from transformers import AutoConfig
-        from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel
+        from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
 
         full_cfg = AutoConfig.from_pretrained(model_name)
         vis_cfg: Qwen3VLVisionConfig = full_cfg.vision_config
 
-        # Instantiate vision model from config only (random init — weights come from checkpoint)
-        vis_model = Qwen3VLVisionModel(vis_cfg).to(torch.float32)
+        full_model = Qwen3VLForConditionalGeneration.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            device_map=f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu",
+        )
+        vis_model = full_model.visual
+        full_model.model = None  # free LLM weights
+        del full_model
+        torch.cuda.empty_cache()
 
         self.hidden_size = vis_cfg.hidden_size          # 1152
         self.patch_size = vis_cfg.patch_size            # 16

patch-reparameterization/src/stage2/__init__.py

@@ -0,0 +1,88 @@
+from .models.lightningDiT import LightningDiT
+from .models.DDT import DiTwDDTHead
+from typing import Callable, Dict, Optional, Type, Union
+from typing import Protocol, Any, runtime_checkable
+
+@runtime_checkable
+class Stage2ModelProtocol(Protocol):
+    def forward(self, *args: Any, **kwargs: Any) -> Any: ...
+    def forward_with_cfg(self, *args: Any, **kwargs: Any) -> Any: ...
+    def forward_with_autoguidance(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+
+# STAGE2_ARCHS: Dict[str, Callable] = {}
+# __all__ = ["STAGE2_ARCHS", "register_stage2"]
+
+
+# def _add_to_registry(name: str, func: Callable) -> Callable:
+#     if name in STAGE2_ARCHS and STAGE2_ARCHS[name] is not func:
+#         raise ValueError(f"Stage2 func: '{name}' is already registered.")
+#     STAGE2_ARCHS[name] = func
+#     return func
+
+
+# def register_stage2(func: Optional[Callable] = None, *, name: Optional[str] = None) -> Union[Callable[[Callable], Callable], Callable]:
+#     """Register a function in ``STAGE2_ARCHS``.
+
+#     Can be used either as ``@register_stage2()`` (optionally passing ``name``) or
+#     via ``register_stage2(my_function)`` after the function definition.
+#     """
+
+#     def decorator(inner_func: Callable) -> Callable:
+#         func_name = name or inner_func.__name__
+#         return _add_to_registry(func_name, inner_func)
+
+#     if func is None:
+#         return decorator
+
+#     return decorator(func)
+
+# @register_stage2()
+# def DiTXL(token_dim: int, input_size: int) -> LightningDiT:
+#     model = LightningDiT(
+#         input_size=input_size,
+#         patch_size=1,
+#         in_channels=token_dim,
+#         hidden_size=1152,
+#         depth=28,
+#         num_heads=16,
+#         mlp_ratio=4.0,
+#         class_dropout_prob=0.1,
+#         num_classes=1000,
+#     )
+#     return model
+
+# @register_stage2()
+# def DDTXL(token_dim: int, input_size: int) -> DiTwDDTHead:
+#     model = DiTwDDTHead(
+#         input_size=input_size,
+#         patch_size=1,
+#         in_channels=token_dim,
+#         hidden_size=[1152, 2048],
+#         depth=[28, 2],
+#         num_heads=[16, 16],
+#         mlp_ratio=4.0,
+#         class_dropout_prob=0.1,
+#         num_classes=1000,
+#     )
+#     return model
+
+# @register_stage2()
+# def DDTS(token_dim: int, input_size: int) -> DiTwDDTHead:
+#     model = DiTwDDTHead(
+#         input_size=input_size,
+#         patch_size=1,
+#         in_channels=token_dim,
+#         hidden_size=[384, 2048],
+#         depth=[12, 2],
+#         num_heads=[6, 16],
+#         mlp_ratio=4.0,
+#         class_dropout_prob=0.1,
+#         num_classes=1000,
+#     )
+#     return model
+
+
+
+

patch-reparameterization/src/stage2/models/DDT.py

@@ -0,0 +1,919 @@
+import logging
+from math import sqrt
+from re import L
+from regex import B
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+from transformers import SwinModel
+import torch
+from torch import nn
+from .lightningDiT import PatchEmbed, Mlp, NormAttention
+from timm.models.vision_transformer import PatchEmbed, Mlp
+from .model_utils import VisionRotaryEmbeddingFast, RMSNorm, SwiGLUFFN, GaussianFourierEmbedding, LabelEmbedder, NormAttention, get_2d_sincos_pos_embed
+import torch.nn.functional as F
+from typing import *
+
+
+def DDTModulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """
+    Applies per-segment modulation to x.
+
+    Args:
+        x: Tensor of shape (B, L_x, D)
+        shift: Tensor of shape (B, L, D)
+        scale: Tensor of shape (B, L, D)
+    Returns:
+        Tensor of shape (B, L_x, D): x * (1 + scale) + shift,
+        with shift and scale repeated to match L_x if necessary.
+    """
+    B, Lx, D = x.shape
+    _, L, _ = shift.shape
+    if Lx % L != 0:
+        raise ValueError(f"L_x ({Lx}) must be divisible by L ({L})")
+    repeat = Lx // L
+    if repeat != 1:
+        # repeat each of the L segments 'repeat' times along the length dim
+        shift = shift.repeat_interleave(repeat, dim=1)
+        scale = scale.repeat_interleave(repeat, dim=1)
+    # apply modulation
+    return x * (1 + scale) + shift
+
+
+def DDTGate(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    """
+    Applies per-segment modulation to x.
+
+    Args:
+        x: Tensor of shape (B, L_x, D)
+        gate: Tensor of shape (B, L, D)
+    Returns:
+        Tensor of shape (B, L_x, D): x * gate,
+        with gate repeated to match L_x if necessary.
+    """
+    B, Lx, D = x.shape
+    _, L, _ = gate.shape
+    if Lx % L != 0:
+        raise ValueError(f"L_x ({Lx}) must be divisible by L ({L})")
+    repeat = Lx // L
+    if repeat != 1:
+        # repeat each of the L segments 'repeat' times along the length dim
+        # print(f"gate shape: {gate.shape}, x shape: {x.shape}")
+        gate = gate.repeat_interleave(repeat, dim=1)
+    # apply modulation
+    return x * gate
+
+
+class LightningDDTBlock(nn.Module):
+    """
+    Lightning DiT Block. We add features including:
+    - ROPE
+    - QKNorm
+    - RMSNorm
+    - SwiGLU
+    - No shift AdaLN.
+    Not all of them are used in the final model, please refer to the paper for more details.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        mlp_ratio=4.0,
+        use_qknorm=False,
+        use_swiglu=True,
+        use_rmsnorm=True,
+        wo_shift=False,
+        **block_kwargs
+    ):
+        super().__init__()
+
+        # Initialize normalization layers
+        if not use_rmsnorm:
+            self.norm1 = nn.LayerNorm(
+                hidden_size, elementwise_affine=False, eps=1e-6)
+            self.norm2 = nn.LayerNorm(
+                hidden_size, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm1 = RMSNorm(hidden_size)
+            self.norm2 = RMSNorm(hidden_size)
+
+        # Initialize attention layer
+        self.attn = NormAttention(
+            hidden_size,
+            num_heads=num_heads,
+            qkv_bias=True,
+            qk_norm=use_qknorm,
+            use_rmsnorm=use_rmsnorm,
+            **block_kwargs
+        )
+
+        # Initialize MLP layer
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        def approx_gelu(): return nn.GELU(approximate="tanh")
+        if use_swiglu:
+            # here we did not use SwiGLU from xformers because it is not compatible with torch.compile for now.
+            self.mlp = SwiGLUFFN(hidden_size, int(2/3 * mlp_hidden_dim))
+        else:
+            self.mlp = Mlp(
+                in_features=hidden_size,
+                hidden_features=mlp_hidden_dim,
+                act_layer=approx_gelu,
+                drop=0
+            )
+
+        # Initialize AdaLN modulation
+        if wo_shift:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, 4 * hidden_size, bias=True)
+            )
+        else:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+            )
+        self.wo_shift = wo_shift
+
+    def forward(self, x, c, feat_rope=None):
+        if len(c.shape) < len(x.shape):
+            c = c.unsqueeze(1)  # (B, 1, C)
+        if self.wo_shift:
+            scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(
+                c).chunk(4, dim=-1)
+            shift_msa = None
+            shift_mlp = None
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(
+                c).chunk(6, dim=-1)
+        x = x + DDTGate(self.attn(DDTModulate(self.norm1(x),
+                        shift_msa, scale_msa), rope=feat_rope), gate_msa)
+        x = x + DDTGate(self.mlp(DDTModulate(self.norm2(x),
+                        shift_mlp, scale_mlp)), gate_mlp)
+        return x
+
+
+class DDTFinalLayer(nn.Module):
+    """
+    The final layer of DDT.
+    """
+
+    def __init__(self, hidden_size, patch_size, out_channels, use_rmsnorm=False):
+        super().__init__()
+        if not use_rmsnorm:
+            self.norm_final = nn.LayerNorm(
+                hidden_size, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm_final = RMSNorm(hidden_size)
+        self.linear = nn.Linear(
+            hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        if len(c.shape) < len(x.shape):
+            c = c.unsqueeze(1)
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = DDTModulate(self.norm_final(x), shift, scale)  # no gate
+        x = self.linear(x)
+        return x
+
+
+class DiTwDDTHead(nn.Module):
+    def __init__(
+            self,
+            input_size: int = 1,
+            patch_size: Union[list, int] = 1,
+            in_channels: int = 768,
+            hidden_size=[1152, 2048],
+            depth=[28, 2],
+            num_heads: Union[list[int], int] = [16, 16],
+            mlp_ratio=4.0,
+            class_dropout_prob=0.1,
+            num_classes=1000,
+            use_qknorm=False,
+            use_swiglu=True,
+            use_rope=True,
+            use_rmsnorm=True,
+            wo_shift=False,
+            use_pos_embed: bool = True,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+
+        self.encoder_hidden_size = hidden_size[0]
+        self.decoder_hidden_size = hidden_size[1]
+        self.num_heads = [num_heads, num_heads] if isinstance(
+            num_heads, int) else list(num_heads)
+        self.num_decoder_blocks = depth[1]
+        self.num_encoder_blocks = depth[0]
+        self.num_blocks = depth[0] + depth[1]
+        self.use_rope = use_rope
+        # analyze patch size
+        if isinstance(patch_size, int) or isinstance(patch_size, float):
+            patch_size = [patch_size, patch_size]  # patch size for s , x embed
+        assert len(
+            patch_size) == 2, f"patch size should be a list of two numbers, but got {patch_size}"
+        self.patch_size = patch_size
+        self.s_patch_size = patch_size[0]
+        self.x_patch_size = patch_size[1]
+        s_channel_per_token = in_channels * self.s_patch_size * self.s_patch_size
+        s_input_size = input_size
+        s_patch_size = self.s_patch_size
+        x_input_size = input_size
+        x_patch_size = self.x_patch_size
+        x_channel_per_token = in_channels * self.x_patch_size * self.x_patch_size
+        self.x_embedder = PatchEmbed(
+            x_input_size, x_patch_size, x_channel_per_token, self.decoder_hidden_size, bias=True)
+        self.s_embedder = PatchEmbed(
+            s_input_size, s_patch_size, s_channel_per_token, self.encoder_hidden_size, bias=True)
+        self.s_channel_per_token = s_channel_per_token
+        self.x_channel_per_token = x_channel_per_token
+        self.s_projector = nn.Linear(
+            self.encoder_hidden_size, self.decoder_hidden_size) if self.encoder_hidden_size != self.decoder_hidden_size else nn.Identity()
+        self.t_embedder = GaussianFourierEmbedding(self.encoder_hidden_size)
+        self.y_embedder = LabelEmbedder(
+            num_classes, self.encoder_hidden_size, class_dropout_prob)
+        # print(f"x_channel_per_token: {x_channel_per_token}, s_channel_per_token: {s_channel_per_token}")
+        self.final_layer = DDTFinalLayer(
+            self.decoder_hidden_size, 1, x_channel_per_token, use_rmsnorm=use_rmsnorm)
+        # Will use fixed sin-cos embedding:
+        if use_pos_embed:
+            num_patches = self.s_embedder.num_patches
+            self.pos_embed = nn.Parameter(torch.zeros(
+                1, num_patches, self.encoder_hidden_size), requires_grad=False)
+            self.x_pos_embed = None
+        self.use_pos_embed = use_pos_embed
+        enc_num_heads = self.num_heads[0]
+        dec_num_heads = self.num_heads[1]
+        # use rotary position encoding, borrow from EVA
+        if self.use_rope:
+            enc_half_head_dim = self.encoder_hidden_size // enc_num_heads // 2
+            hw_seq_len = int(sqrt(self.s_embedder.num_patches))
+            # print(f"enc_half_head_dim: {enc_half_head_dim}, hw_seq_len: {hw_seq_len}")
+            self.enc_feat_rope = VisionRotaryEmbeddingFast(
+                dim=enc_half_head_dim,
+                pt_seq_len=hw_seq_len,
+            )
+            dec_half_head_dim = self.decoder_hidden_size // dec_num_heads // 2
+            hw_seq_len = int(sqrt(self.x_embedder.num_patches))
+            # print(f"dec_half_head_dim: {dec_half_head_dim}, hw_seq_len: {hw_seq_len}")
+            self.dec_feat_rope = VisionRotaryEmbeddingFast(
+                dim=dec_half_head_dim,
+                pt_seq_len=hw_seq_len,
+            )
+        else:
+            self.feat_rope = None
+        self.blocks = nn.ModuleList([
+            LightningDDTBlock(self.encoder_hidden_size if i < self.num_encoder_blocks else self.decoder_hidden_size,
+                              enc_num_heads if i < self.num_encoder_blocks else dec_num_heads,
+                              mlp_ratio=mlp_ratio,
+                              use_qknorm=use_qknorm,
+                              use_rmsnorm=use_rmsnorm,
+                              use_swiglu=use_swiglu,
+                              wo_shift=wo_shift,
+                              ) for i in range(self.num_blocks)
+        ])
+        self.initialize_weights()
+
+    def initialize_weights(self, xavier_uniform_init: bool = False):
+        if xavier_uniform_init:
+            def _basic_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.xavier_uniform_(module.weight)
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            self.apply(_basic_init)
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.s_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.s_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        if self.use_pos_embed:
+            # Initialize (and freeze) pos_embed by sin-cos embedding:
+            pos_embed = get_2d_sincos_pos_embed(
+                self.pos_embed.shape[-1], int(self.s_embedder.num_patches ** 0.5))
+            self.pos_embed.data.copy_(
+                torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Zero-out adaLN modulation layers in LightningDiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        # c = self.out_channels
+        c = self.x_channel_per_token
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        return imgs
+
+    def encoder_forward(self, x, t, y):
+        """Encoder part — safe under bf16 autocast, can be torch.compiled."""
+        t = self.t_embedder(t)
+        y = self.y_embedder(y, self.training)
+        c = nn.functional.silu(t + y)
+        s = self.s_embedder(x)
+        if self.use_pos_embed:
+            s = s + self.pos_embed
+        for i in range(self.num_encoder_blocks):
+            s = self.blocks[i](s, c, feat_rope=self.enc_feat_rope)
+        t = t.unsqueeze(1).repeat(1, s.shape[1], 1)
+        s = nn.functional.silu(t + s)
+        return s
+
+    def decoder_forward(self, x, s):
+        """Decoder part — runs in fp32 to prevent bf16 accumulation errors in DDT modulation."""
+        with torch.cuda.amp.autocast(enabled=False):
+            s = self.s_projector(s.float())
+            x = self.x_embedder(x.float())
+            if self.use_pos_embed and self.x_pos_embed is not None:
+                x = x + self.x_pos_embed
+            for i in range(self.num_encoder_blocks, self.num_blocks):
+                x = self.blocks[i](x, s, feat_rope=self.dec_feat_rope)
+            x = self.final_layer(x, s)
+            x = self.unpatchify(x)
+        return x
+
+    def forward(self, x, t, y, s=None, mask=None):
+        if s is None:
+            s = self.encoder_forward(x, t, y)
+        return self.decoder_forward(x, s)
+
+    def forward_with_cfg(self, x, t, y, cfg_scale, cfg_interval=(0, 1)):
+        """
+        Forward pass of LightningDiT, but also batches the unconditional forward pass for classifier-free guidance.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
+        half = x[: len(x) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self.forward(combined, t, y)
+        # For exact reproducibility reasons, we apply classifier-free guidance on only
+        # three channels by default. The standard approach to cfg applies it to all channels.
+        # This can be done by uncommenting the following line and commenting-out the line following that.
+        # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
+        eps, rest = model_out[:,
+                              :self.in_channels], model_out[:, self.in_channels:]
+        # eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        guid_t_min, guid_t_max = cfg_interval
+        assert guid_t_min < guid_t_max, "cfg_interval should be (min, max) with min < max"
+        t = t[: len(t) // 2] # get t for the conditional half
+        half_eps = torch.where(
+            ((t >= guid_t_min) & (t <= guid_t_max)
+             ).view(-1, *[1] * (len(cond_eps.shape) - 1)),
+            uncond_eps + cfg_scale * (cond_eps - uncond_eps), cond_eps
+        )
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+
+    def forward_with_autoguidance(self, x, t, y, cfg_scale, additional_model_forward, cfg_interval=(0, 1)):
+        """
+        Forward pass of LightningDiT, but also contain the forward pass for the additional model
+        """
+        model_out = self.forward(x, t, y)
+        ag_model_out = additional_model_forward(x, t, y)
+        eps = model_out[:, :self.in_channels]
+        ag_eps = ag_model_out[:, :self.in_channels]
+
+        guid_t_min, guid_t_max = cfg_interval
+        assert guid_t_min < guid_t_max, "cfg_interval should be (min, max) with min < max"
+        eps = torch.where(
+            ((t >= guid_t_min) & (t <= guid_t_max)
+             ).view(-1, *[1] * (len(eps.shape) - 1)),
+            ag_eps + cfg_scale * (eps - ag_eps), eps
+        )
+
+        return eps
+
+
+class DiTwDDTHeadDualHead(DiTwDDTHead):
+    """DDT with dual decoder heads: one for semantic, one for recon.
+
+    Shared encoder (28 blocks) → two independent decoder branches:
+      - sem branch: s_projector_sem → x_embedder_sem → decoder_sem_blocks → final_layer_sem → sem_channels
+      - recon branch: s_projector_recon → x_embedder_recon → decoder_recon_blocks → final_layer_recon → recon_channels
+    Outputs are concatenated back to in_channels.
+    """
+
+    def __init__(
+            self,
+            sem_channels: int = 768,
+            input_size: int = 1,
+            patch_size: Union[list, int] = 1,
+            in_channels: int = 768,
+            hidden_size=[1152, 2048],
+            depth=[28, 2],
+            num_heads: Union[list[int], int] = [16, 16],
+            mlp_ratio=4.0,
+            class_dropout_prob=0.1,
+            num_classes=1000,
+            use_qknorm=False,
+            use_swiglu=True,
+            use_rope=True,
+            use_rmsnorm=True,
+            wo_shift=False,
+            use_pos_embed: bool = True,
+    ):
+        # Initialize parent (builds shared encoder + original decoder)
+        super().__init__(
+            input_size=input_size, patch_size=patch_size, in_channels=in_channels,
+            hidden_size=hidden_size, depth=depth, num_heads=num_heads,
+            mlp_ratio=mlp_ratio, class_dropout_prob=class_dropout_prob,
+            num_classes=num_classes, use_qknorm=use_qknorm, use_swiglu=use_swiglu,
+            use_rope=use_rope, use_rmsnorm=use_rmsnorm, wo_shift=wo_shift,
+            use_pos_embed=use_pos_embed,
+        )
+        recon_channels = in_channels - sem_channels
+        assert recon_channels > 0, f"sem_channels={sem_channels} >= in_channels={in_channels}"
+        self.sem_channels = sem_channels
+        self.recon_channels = recon_channels
+
+        p = self.x_patch_size
+        dec_num_heads = self.num_heads[1]
+        _mlp_ratio = mlp_ratio
+
+        # Input size for x_embedder (spatial dim before patchify)
+        _x_input_size = int(sqrt(self.x_embedder.num_patches)) * self.x_patch_size
+
+        # --- Semantic decoder branch ---
+        sem_ch_per_token = sem_channels * p * p
+        self.x_embedder_sem = PatchEmbed(
+            _x_input_size, self.x_patch_size, sem_ch_per_token, self.decoder_hidden_size, bias=True)
+        self.s_projector_sem = nn.Linear(
+            self.encoder_hidden_size, self.decoder_hidden_size) if self.encoder_hidden_size != self.decoder_hidden_size else nn.Identity()
+        self.decoder_blocks_sem = nn.ModuleList([
+            LightningDDTBlock(self.decoder_hidden_size, dec_num_heads,
+                              mlp_ratio=_mlp_ratio,
+                              use_qknorm=use_qknorm, use_rmsnorm=use_rmsnorm,
+                              use_swiglu=use_swiglu, wo_shift=wo_shift)
+            for _ in range(self.num_decoder_blocks)
+        ])
+        self.final_layer_sem = DDTFinalLayer(
+            self.decoder_hidden_size, 1, sem_ch_per_token, use_rmsnorm=use_rmsnorm)
+
+        # --- Recon decoder branch ---
+        recon_ch_per_token = recon_channels * p * p
+        self.x_embedder_recon = PatchEmbed(
+            _x_input_size, self.x_patch_size, recon_ch_per_token, self.decoder_hidden_size, bias=True)
+        self.s_projector_recon = nn.Linear(
+            self.encoder_hidden_size, self.decoder_hidden_size) if self.encoder_hidden_size != self.decoder_hidden_size else nn.Identity()
+        self.decoder_blocks_recon = nn.ModuleList([
+            LightningDDTBlock(self.decoder_hidden_size, dec_num_heads,
+                              mlp_ratio=_mlp_ratio,
+                              use_qknorm=use_qknorm, use_rmsnorm=use_rmsnorm,
+                              use_swiglu=use_swiglu, wo_shift=wo_shift)
+            for _ in range(self.num_decoder_blocks)
+        ])
+        self.final_layer_recon = DDTFinalLayer(
+            self.decoder_hidden_size, 1, recon_ch_per_token, use_rmsnorm=use_rmsnorm)
+
+        # RoPE for new decoder branches (shared same config as original dec_feat_rope)
+        if self.use_rope:
+            self.dec_feat_rope_sem = self.dec_feat_rope
+            self.dec_feat_rope_recon = self.dec_feat_rope
+
+        # Remove parent's original decoder modules (unused, would cause DDP errors)
+        del self.x_embedder
+        del self.s_projector
+        del self.final_layer
+        # Remove parent's original decoder blocks (blocks[num_encoder_blocks:])
+        original_blocks = list(self.blocks)
+        self.blocks = nn.ModuleList(original_blocks[:self.num_encoder_blocks])
+
+        # Initialize new modules
+        self._init_dual_head_weights()
+
+        logger.info(f"DiTwDDTHeadDualHead: sem_channels={sem_channels}, recon_channels={recon_channels}, "
+                    f"decoder_blocks={self.num_decoder_blocks} x2, decoder_hidden={self.decoder_hidden_size}")
+
+    def _init_dual_head_weights(self):
+        """Initialize weights for the dual decoder branches."""
+        for embedder in [self.x_embedder_sem, self.x_embedder_recon]:
+            w = embedder.proj.weight.data
+            nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+            nn.init.constant_(embedder.proj.bias, 0)
+
+        for blocks in [self.decoder_blocks_sem, self.decoder_blocks_recon]:
+            for block in blocks:
+                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        for fl in [self.final_layer_sem, self.final_layer_recon]:
+            nn.init.constant_(fl.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(fl.adaLN_modulation[-1].bias, 0)
+            nn.init.constant_(fl.linear.weight, 0)
+            nn.init.constant_(fl.linear.bias, 0)
+
+    def _unpatchify_branch(self, x, channels):
+        """Unpatchify for a branch with specific number of channels."""
+        c = channels * self.x_patch_size * self.x_patch_size
+        p = self.x_patch_size
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        return imgs
+
+    def decoder_forward(self, x, s):
+        """Dual decoder: sem and recon branches run independently, outputs concatenated."""
+        with torch.cuda.amp.autocast(enabled=False):
+            # Split input x along channel dim: [B, C, H, W] → sem [B, sem_ch, H, W] + recon [B, recon_ch, H, W]
+            x_float = x.float()
+            x_sem_input = x_float[:, :self.sem_channels, :, :]
+            x_recon_input = x_float[:, self.sem_channels:, :, :]
+
+            s_float = s.float()
+
+            # --- Semantic branch ---
+            s_sem = self.s_projector_sem(s_float)
+            h_sem = self.x_embedder_sem(x_sem_input)
+            rope_sem = self.dec_feat_rope_sem if self.use_rope else None
+            for block in self.decoder_blocks_sem:
+                h_sem = block(h_sem, s_sem, feat_rope=rope_sem)
+            h_sem = self.final_layer_sem(h_sem, s_sem)
+            out_sem = self._unpatchify_branch(h_sem, self.sem_channels)  # [B, sem_ch, H, W]
+
+            # --- Recon branch ---
+            s_recon = self.s_projector_recon(s_float)
+            h_recon = self.x_embedder_recon(x_recon_input)
+            rope_recon = self.dec_feat_rope_recon if self.use_rope else None
+            for block in self.decoder_blocks_recon:
+                h_recon = block(h_recon, s_recon, feat_rope=rope_recon)
+            h_recon = self.final_layer_recon(h_recon, s_recon)
+            out_recon = self._unpatchify_branch(h_recon, self.recon_channels)  # [B, recon_ch, H, W]
+
+            # Concatenate: [B, sem_ch + recon_ch, H, W] = [B, in_channels, H, W]
+            out = torch.cat([out_sem, out_recon], dim=1)
+        return out
+
+
+class PartialRoPE(nn.Module):
+    """
+    A RoPE wrapper that only applies rotation to the last `num_spatial` tokens
+    in the sequence, leaving the leading `num_globaltoken` tokens unchanged.
+
+    The underlying `feat_rope` (VisionRotaryEmbeddingFast) expects input of shape
+    [B, num_heads, num_spatial, head_dim] and returns the same shape with RoPE applied.
+    This wrapper slices q/k, applies RoPE to the spatial slice, then stitches back.
+
+    Args:
+        feat_rope:        a VisionRotaryEmbeddingFast instance sized for num_spatial tokens.
+        num_globaltoken:  number of leading global tokens to skip.
+    """
+    def __init__(self, feat_rope, num_globaltoken: int):
+        super().__init__()
+        self.feat_rope = feat_rope
+        self.num_globaltoken = num_globaltoken
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, num_heads, N, head_dim]  where N = num_globaltoken + num_spatial
+        Returns:
+            x: same shape, with RoPE applied only to x[:, :, num_globaltoken:, :]
+        """
+        g = self.num_globaltoken
+        x_global  = x[:, :, :g, :]   # (B, H, num_globaltoken, D)  — untouched
+        x_spatial = x[:, :, g:, :]   # (B, H, num_spatial,     D)  — apply RoPE
+        x_spatial = self.feat_rope(x_spatial)
+        return torch.cat([x_global, x_spatial], dim=2)
+
+
+class DiTwDDTHeadTokenInput(nn.Module):
+    """
+    A variant of DiTwDDTHead that accepts pre-tokenized input x of shape [B, N, D],
+    where N = num_globaltoken + H * W (global tokens come first).
+
+    Key differences from DiTwDDTHead:
+    - Input x is already tokenized: no PatchEmbed needed; a linear projector maps D -> hidden_size.
+    - Position encoding:
+        * spatial tokens (H*W) use fixed 2D sin-cos positional embedding.
+        * global tokens use learnable positional embedding.
+    - Encoder (DDT-style): all N tokens pass through encoder blocks together with AdaLN condition c.
+      After encoding, timestep t is broadcast and fused into the token sequence as the condition s.
+    - Decoder: all N tokens are denoised conditioned on s from the encoder via AdaLN.
+    - Output: [B, N, D], same shape as input (both global and spatial tokens are denoised).
+    """
+
+    def __init__(
+            self,
+            in_channels: int = 768,         # D: input/output token dimension
+            hw_size: int = 16,              # spatial grid side length, num spatial tokens = hw_size^2
+            num_globaltoken: int = 8,       # number of global tokens (prepended before spatial tokens)
+            hidden_size: Union[list, int] = [1152, 2048],  # [encoder_hidden, decoder_hidden]
+            depth: list = [28, 2],          # [num_encoder_blocks, num_decoder_blocks]
+            num_heads: Union[list, int] = [16, 16],
+            mlp_ratio: float = 4.0,
+            class_dropout_prob: float = 0.1,
+            num_classes: int = 1000,
+            use_qknorm: bool = False,
+            use_swiglu: bool = True,
+            use_rope: bool = False,         # RoPE is disabled by default for mixed token sequences
+            use_rmsnorm: bool = True,
+            wo_shift: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.hw_size = hw_size
+        self.num_spatial = hw_size * hw_size
+        self.num_globaltoken = num_globaltoken
+        self.num_tokens = num_globaltoken + self.num_spatial  # total N
+
+        # hidden sizes
+        if isinstance(hidden_size, int):
+            hidden_size = [hidden_size, hidden_size]
+        self.encoder_hidden_size = hidden_size[0]
+        self.decoder_hidden_size = hidden_size[1]
+
+        # num heads
+        if isinstance(num_heads, int):
+            num_heads = [num_heads, num_heads]
+        self.num_heads = list(num_heads)
+        enc_num_heads = self.num_heads[0]
+        dec_num_heads = self.num_heads[1]
+
+        # block depths
+        self.num_encoder_blocks = depth[0]
+        self.num_decoder_blocks = depth[1]
+        self.num_blocks = depth[0] + depth[1]
+        self.use_rope = use_rope
+
+        # ---------- input projectors ----------
+        # project D -> encoder_hidden_size for the encoder stage
+        self.enc_input_proj = nn.Linear(in_channels, self.encoder_hidden_size, bias=True)
+        # project D -> decoder_hidden_size for the decoder stage
+        self.dec_input_proj = nn.Linear(in_channels, self.decoder_hidden_size, bias=True)
+
+        # ---------- position embeddings ----------
+        # spatial tokens: fixed 2D sin-cos, shape (1, num_spatial, encoder_hidden_size)
+        self.enc_spatial_pos_embed = nn.Parameter(
+            torch.zeros(1, self.num_spatial, self.encoder_hidden_size), requires_grad=False)
+        self.dec_spatial_pos_embed = nn.Parameter(
+            torch.zeros(1, self.num_spatial, self.decoder_hidden_size), requires_grad=False)
+
+        # global tokens: learnable, shape (1, num_globaltoken, hidden_size)
+        self.enc_global_pos_embed = nn.Parameter(
+            torch.zeros(1, num_globaltoken, self.encoder_hidden_size), requires_grad=True)
+        self.dec_global_pos_embed = nn.Parameter(
+            torch.zeros(1, num_globaltoken, self.decoder_hidden_size), requires_grad=True)
+
+        # ---------- condition embedders ----------
+        self.t_embedder = GaussianFourierEmbedding(self.encoder_hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes, self.encoder_hidden_size, class_dropout_prob)
+
+        # project encoder hidden -> decoder hidden for condition s
+        self.s_projector = (
+            nn.Linear(self.encoder_hidden_size, self.decoder_hidden_size)
+            if self.encoder_hidden_size != self.decoder_hidden_size
+            else nn.Identity()
+        )
+
+        # ---------- RoPE (optional, only applied to spatial tokens) ----------
+        # global tokens are skipped; PartialRoPE wraps VisionRotaryEmbeddingFast
+        # so that only the trailing num_spatial positions are rotated.
+        if self.use_rope:
+            enc_half_head_dim = self.encoder_hidden_size // enc_num_heads // 2
+            self.enc_feat_rope = PartialRoPE(
+                feat_rope=VisionRotaryEmbeddingFast(
+                    dim=enc_half_head_dim,
+                    pt_seq_len=hw_size,
+                ),
+                num_globaltoken=num_globaltoken,
+            )
+            dec_half_head_dim = self.decoder_hidden_size // dec_num_heads // 2
+            self.dec_feat_rope = PartialRoPE(
+                feat_rope=VisionRotaryEmbeddingFast(
+                    dim=dec_half_head_dim,
+                    pt_seq_len=hw_size,
+                ),
+                num_globaltoken=num_globaltoken,
+            )
+        else:
+            self.enc_feat_rope = None
+            self.dec_feat_rope = None
+
+        # ---------- transformer blocks ----------
+        self.blocks = nn.ModuleList([
+            LightningDDTBlock(
+                self.encoder_hidden_size if i < self.num_encoder_blocks else self.decoder_hidden_size,
+                enc_num_heads if i < self.num_encoder_blocks else dec_num_heads,
+                mlp_ratio=mlp_ratio,
+                use_qknorm=use_qknorm,
+                use_rmsnorm=use_rmsnorm,
+                use_swiglu=use_swiglu,
+                wo_shift=wo_shift,
+            ) for i in range(self.num_blocks)
+        ])
+
+        # ---------- final layer ----------
+        # outputs in_channels per token (no patch expansion)
+        self.final_layer = DDTFinalLayer(
+            self.decoder_hidden_size, 1, in_channels, use_rmsnorm=use_rmsnorm)
+
+        # ---------- output projector ----------
+        # map decoder_hidden_size back to in_channels
+        # (DDTFinalLayer already does linear(hidden -> 1*1*in_channels), so this is identity-like)
+        # We keep it explicit for clarity; DDTFinalLayer.linear outputs in_channels directly.
+
+        self.initialize_weights()
+
+    # ------------------------------------------------------------------
+    def initialize_weights(self):
+        # sin-cos 2D pos embed for spatial tokens (encoder)
+        enc_pos = get_2d_sincos_pos_embed(self.encoder_hidden_size, self.hw_size)
+        self.enc_spatial_pos_embed.data.copy_(
+            torch.from_numpy(enc_pos).float().unsqueeze(0))
+
+        # sin-cos 2D pos embed for spatial tokens (decoder)
+        dec_pos = get_2d_sincos_pos_embed(self.decoder_hidden_size, self.hw_size)
+        self.dec_spatial_pos_embed.data.copy_(
+            torch.from_numpy(dec_pos).float().unsqueeze(0))
+
+        # learnable global pos embed: small normal init
+        nn.init.normal_(self.enc_global_pos_embed, std=0.02)
+        nn.init.normal_(self.dec_global_pos_embed, std=0.02)
+
+        # input projectors
+        nn.init.xavier_uniform_(self.enc_input_proj.weight)
+        nn.init.constant_(self.enc_input_proj.bias, 0)
+        nn.init.xavier_uniform_(self.dec_input_proj.weight)
+        nn.init.constant_(self.dec_input_proj.bias, 0)
+
+        # label embedding
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        # timestep embedding MLP
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # zero-out adaLN modulation layers
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # zero-out final layer
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    # ------------------------------------------------------------------
+    def _build_pos_embed(self, hidden_size, global_pos_embed, spatial_pos_embed):
+        """
+        Concatenate [global_pos_embed, spatial_pos_embed] along dim=1.
+        global_pos_embed: (1, num_globaltoken, hidden_size)  — learnable
+        spatial_pos_embed: (1, num_spatial, hidden_size)     — fixed sin-cos
+        Returns: (1, num_tokens, hidden_size)
+        """
+        return torch.cat([global_pos_embed, spatial_pos_embed], dim=1)
+
+    # ------------------------------------------------------------------
+    def forward(self, x, t, y, s=None):
+        """
+        Args:
+            x: [B, N, D]  noisy token sequence, N = num_globaltoken + H*W
+            t: [B]        diffusion timestep
+            y: [B]        class label (int)
+            s: optional pre-computed encoder condition [B, N, encoder_hidden_size].
+               If provided, the encoder stage is skipped.
+        Returns:
+            out: [B, N, D]  denoised token sequence (same shape as input)
+        """
+        B, N, D = x.shape
+        assert N == self.num_tokens, (
+            f"Expected N={self.num_tokens} (num_globaltoken={self.num_globaltoken} + "
+            f"num_spatial={self.num_spatial}), got {N}"
+        )
+
+        # ---- condition embeddings ----
+        t_emb = self.t_embedder(t)                       # (B, encoder_hidden_size)
+        y_emb = self.y_embedder(y, self.training)        # (B, encoder_hidden_size)
+        c = nn.functional.silu(t_emb + y_emb)           # (B, encoder_hidden_size)
+
+        # ---- encoder stage ----
+        if s is None:
+            # project input to encoder hidden size
+            xe = self.enc_input_proj(x)                  # (B, N, encoder_hidden_size)
+
+            # add position embeddings:
+            #   global tokens (front): learnable enc_global_pos_embed
+            #   spatial tokens (back): fixed sin-cos enc_spatial_pos_embed
+            pos = self._build_pos_embed(
+                self.encoder_hidden_size,
+                self.enc_global_pos_embed,               # (1, num_globaltoken, enc_h)
+                self.enc_spatial_pos_embed,              # (1, num_spatial,     enc_h)
+            )  # (1, N, encoder_hidden_size)
+            xe = xe + pos
+
+            # pass through encoder blocks (self-attention conditioned on c)
+            # PartialRoPE ensures only spatial tokens (trailing hw_size^2) are rotated
+            for i in range(self.num_encoder_blocks):
+                xe = self.blocks[i](xe, c, feat_rope=self.enc_feat_rope)
+
+            # fuse timestep into encoder output to form condition s
+            t_broadcast = t_emb.unsqueeze(1).expand(-1, N, -1)  # (B, N, enc_h)
+            s = nn.functional.silu(t_broadcast + xe)             # (B, N, encoder_hidden_size)
+
+        # project condition s to decoder hidden size
+        s = self.s_projector(s)                          # (B, N, decoder_hidden_size)
+
+        # ---- decoder stage ----
+        # project input to decoder hidden size
+        xd = self.dec_input_proj(x)                      # (B, N, decoder_hidden_size)
+
+        # add decoder position embeddings
+        dec_pos = self._build_pos_embed(
+            self.decoder_hidden_size,
+            self.dec_global_pos_embed,                   # (1, num_globaltoken, dec_h)
+            self.dec_spatial_pos_embed,                  # (1, num_spatial,     dec_h)
+        )  # (1, N, decoder_hidden_size)
+        xd = xd + dec_pos
+
+        # pass through decoder blocks (conditioned on s)
+        # PartialRoPE ensures only spatial tokens (trailing hw_size^2) are rotated
+        for i in range(self.num_encoder_blocks, self.num_blocks):
+            xd = self.blocks[i](xd, s, feat_rope=self.dec_feat_rope)
+
+        # final layer: outputs (B, N, in_channels)
+        out = self.final_layer(xd, s)                    # (B, N, in_channels)
+        return out
+
+    # ------------------------------------------------------------------
+    def forward_with_cfg(self, x, t, y, cfg_scale, cfg_interval=(0, 1)):
+        """
+        Classifier-free guidance forward pass.
+        Assumes x is already doubled: first half conditional, second half unconditional.
+        """
+        half = x[: len(x) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self.forward(combined, t, y)         # (2B, N, D)
+
+        eps, rest = model_out[:, :, :self.in_channels], model_out[:, :, self.in_channels:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+
+        guid_t_min, guid_t_max = cfg_interval
+        assert guid_t_min < guid_t_max
+        t_half = t[: len(t) // 2]
+        half_eps = torch.where(
+            ((t_half >= guid_t_min) & (t_half <= guid_t_max)
+             ).view(-1, *[1] * (len(cond_eps.shape) - 1)),
+            uncond_eps + cfg_scale * (cond_eps - uncond_eps),
+            cond_eps,
+        )
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=-1)
+
+    # ------------------------------------------------------------------
+    def forward_with_autoguidance(self, x, t, y, cfg_scale, additional_model_forward, cfg_interval=(0, 1)):
+        """
+        Auto-guidance forward pass using an additional (weaker) model.
+        """
+        model_out = self.forward(x, t, y)
+        ag_model_out = additional_model_forward(x, t, y)
+
+        eps = model_out[:, :, :self.in_channels]
+        ag_eps = ag_model_out[:, :, :self.in_channels]
+
+        guid_t_min, guid_t_max = cfg_interval
+        assert guid_t_min < guid_t_max
+        eps = torch.where(
+            ((t >= guid_t_min) & (t <= guid_t_max)
+             ).view(-1, *[1] * (len(eps.shape) - 1)),
+            ag_eps + cfg_scale * (eps - ag_eps),
+            eps,
+        )
+        return eps

patch-reparameterization/src/stage2/models/__init__.py

@@ -0,0 +1,15 @@
+"""Encoder registry for stage 2 models."""
+
+
+
+# Import modules that perform registration on import.
+from typing import Callable, Dict, Optional, Type, Union
+from typing import Protocol, Any, runtime_checkable
+
+@runtime_checkable
+class Stage2ModelProtocol(Protocol):
+    def forward(self, *args: Any, **kwargs: Any) -> Any: ...
+    def forward_with_cfg(self, *args: Any, **kwargs: Any) -> Any: ...
+    def forward_with_autoguidance(self, *args: Any, **kwargs: Any) -> Any: ...
+    def train(self, mode: bool = True) -> Any: ...
+    def eval(self) -> Any: ...

patch-reparameterization/src/stage2/models/lightningDiT.py

@@ -0,0 +1,305 @@
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import PatchEmbed, Mlp
+import torch
+from torch import nn
+
+from .model_utils import VisionRotaryEmbeddingFast, SwiGLUFFN, RMSNorm, NormAttention, LabelEmbedder, get_2d_sincos_pos_embed, GaussianFourierEmbedding, modulate
+
+
+
+
+class LightningDiTBlock(nn.Module):
+    """
+    Lightning DiT Block. We add features including:
+    - ROPE
+    - QKNorm
+    - RMSNorm
+    - SwiGLU
+    - No shift AdaLN.
+    Not all of them are used in the final model, please refer to the paper for more details.
+    """
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        mlp_ratio=4.0,
+        use_qknorm=False,
+        use_swiglu=True,
+        use_rmsnorm=True,
+        wo_shift=False,
+        **block_kwargs
+    ):
+        super().__init__()
+
+        # Initialize normalization layers
+        if not use_rmsnorm:
+            self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+            self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm1 = RMSNorm(hidden_size)
+            self.norm2 = RMSNorm(hidden_size)
+
+        # Initialize attention layer
+        self.attn = NormAttention(
+            hidden_size,
+            num_heads=num_heads,
+            qkv_bias=True,
+            qk_norm=use_qknorm,
+            use_rmsnorm=use_rmsnorm,
+            **block_kwargs
+        )
+
+        # Initialize MLP layer
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        if use_swiglu:
+            # here we did not use SwiGLU from xformers because it is not compatible with torch.compile for now.
+            self.mlp = SwiGLUFFN(hidden_size, int(2/3 * mlp_hidden_dim))
+        else:
+            self.mlp = Mlp(
+                in_features=hidden_size,
+                hidden_features=mlp_hidden_dim,
+                act_layer=approx_gelu,
+                drop=0
+            )
+
+        # Initialize AdaLN modulation
+        if wo_shift:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, 4 * hidden_size, bias=True)
+            )
+        else:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+            )
+        self.wo_shift = wo_shift
+
+    def forward(self, x, c, feat_rope=None):
+        if self.wo_shift:
+            scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(4, dim=1)
+            shift_msa = None
+            shift_mlp = None
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), rope=feat_rope)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+class LightningFinalLayer(nn.Module):
+    """
+    The final layer of LightningDiT.
+    """
+    def __init__(self, hidden_size, patch_size, out_channels, use_rmsnorm=False):
+        super().__init__()
+        if not use_rmsnorm:
+            self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm_final = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+class LightningDiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        input_size=16,
+        patch_size=1,
+        in_channels=768,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        class_dropout_prob=0.1,
+        num_classes=1000,
+        learn_sigma=False,
+        use_qknorm=False,
+        use_swiglu=True,
+        use_rope=True,
+        use_rmsnorm=True,
+        wo_shift=False,
+        use_gembed: bool = True,
+    ):
+        super().__init__()
+        self.learn_sigma = learn_sigma
+        self.in_channels = in_channels
+        self.out_channels = in_channels if not learn_sigma else in_channels * 2
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.use_rope = use_rope
+        self.use_rmsnorm = use_rmsnorm
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.use_gembed = use_gembed
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
+        self.t_embedder = GaussianFourierEmbedding(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
+        self.ssl_supervise = False
+        num_patches = self.x_embedder.num_patches
+        # Will use fixed sin-cos embedding:
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+
+        # use rotary position encoding, borrow from EVA
+        if self.use_rope:
+            half_head_dim = hidden_size // num_heads // 2
+            hw_seq_len = input_size // patch_size
+            self.feat_rope = VisionRotaryEmbeddingFast(
+                dim=half_head_dim,
+                pt_seq_len=hw_seq_len,
+            )
+        else:
+            self.feat_rope = None
+
+        self.blocks = nn.ModuleList([
+            LightningDiTBlock(hidden_size,
+                     num_heads,
+                     mlp_ratio=mlp_ratio,
+                     use_qknorm=use_qknorm,
+                     use_swiglu=use_swiglu,
+                     use_rmsnorm=use_rmsnorm,
+                     wo_shift=wo_shift,
+                     ) for _ in range(depth)
+        ])
+        self.final_layer = LightningFinalLayer(hidden_size, patch_size, self.out_channels, use_rmsnorm=use_rmsnorm)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5))
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in LightningDiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        return imgs
+
+    def forward(self, x, t=None, y=None):
+        """
+        Forward pass of LightningDiT.
+        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        y: (N,) tensor of class labels
+        use_checkpoint: boolean to toggle checkpointing
+        """
+        x = self.x_embedder(x) + self.pos_embed  # (N, T, D), where T = H * W / patch_size ** 2
+        t = self.t_embedder(t)                   # (N, D)
+        y = self.y_embedder(y, self.training)    # (N, D)
+        c = t + y                                # (N, D)
+
+        for block in self.blocks:
+            x = block(x, c, feat_rope=self.feat_rope)
+
+        x = self.final_layer(x, c)                # (N, T, patch_size ** 2 * out_channels)
+        x = self.unpatchify(x)                   # (N, out_channels, H, W)
+
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        return x
+
+    def forward_with_cfg(self, x, t, y, cfg_scale, cfg_interval=(-1e4, -1e4), interval_cfg: float = 0.0):
+        """
+        Forward pass of LightningDiT, but also batches the unconditional forward pass for classifier-free guidance.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
+        half = x[: len(x) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self.forward(combined, t, y)
+        if self.ssl_supervise:
+            model_out = model_out[0] # take the output only
+        # For exact reproducibility reasons, we apply classifier-free guidance on only
+        # three channels by default. The standard approach to cfg applies it to all channels.
+        # This can be done by uncommenting the following line and commenting-out the line following that.
+        # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
+        eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
+        #eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        t = t[0] # check if t < cfg_interval
+        in_interval = False
+        for i in range(len(cfg_interval)):
+            if t >= cfg_interval[i][0] and t < cfg_interval[i][1]:
+                #print(f't:{t}, cfg_interval: {cfg_interval[i]}')
+                if interval_cfg > 1.0:
+                    half_eps = uncond_eps  + interval_cfg * (cond_eps - uncond_eps)
+                else:
+                    half_eps = cond_eps # only use conditional generation
+                in_interval = True
+                break
+        if not in_interval:
+            #print(f't:{t} not in cfg_interval')
+            half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+    def forward_with_autoguidance(self, x, t, y, cfg_scale, additional_model_forward, cfg_interval=(-1e4, -1e4), interval_cfg: float = 0.0):
+        """
+        Forward pass of LightningDiT, but also contain the forward pass for the additional model
+        """
+        half = x[: len(x) // 2] # cut the x by half, autoguidance does not need repeated input
+        t = t[: len(t) // 2]
+        y = y[: len(y) // 2]
+        model_out = self.forward(half, t, y)
+        ag_model_out = additional_model_forward(half, t, y)
+        eps = model_out[:, :self.in_channels]
+        ag_eps = ag_model_out[:, :self.in_channels]
+        t = t[0]
+        in_interval = False
+        for i in range(len(cfg_interval)):
+            if t >= cfg_interval[i][0] and t < cfg_interval[i][1]:
+                if interval_cfg > 1.0:
+                    eps = ag_eps + interval_cfg * (eps - ag_eps)
+                in_interval = True
+                break
+        if not in_interval:
+            eps = ag_eps + cfg_scale * (eps - ag_eps)
+        return torch.cat([eps, eps], dim=0)

patch-reparameterization/src/stage2/models/model_utils.py

@@ -0,0 +1,438 @@
+import numpy.random as random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange, repeat
+from typing import *
+from math import pi
+from timm.models.vision_transformer import Attention
+from functools import partial
+from einops import rearrange, repeat
+from collections.abc import Callable
+import numpy as np
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+def broadcat(tensors, dim=-1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+    assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]
+               ), 'invalid dimensions for broadcastable concatentation'
+    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x):
+    x = rearrange(x, '... (d r) -> ... d r', r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, '... d r -> ... (d r)')
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+class VisionRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        pt_seq_len,
+        ft_seq_len=None,
+        custom_freqs=None,
+        freqs_for='lang',
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f'unknown modality {freqs_for}')
+
+        if ft_seq_len is None:
+            ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs_h = torch.einsum('..., f -> ... f', t, freqs)
+        freqs_h = repeat(freqs_h, '... n -> ... (n r)', r=2)
+
+        freqs_w = torch.einsum('..., f -> ... f', t, freqs)
+        freqs_w = repeat(freqs_w, '... n -> ... (n r)', r=2)
+
+        freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
+
+        self.register_buffer("freqs_cos", freqs.cos())
+        self.register_buffer("freqs_sin", freqs.sin())
+
+        # print('======== shape of rope freq', self.freqs_cos.shape, '========')
+
+    def forward(self, t, start_index=0):
+        rot_dim = self.freqs_cos.shape[-1]
+        end_index = start_index + rot_dim
+        assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
+        t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]
+        t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)
+        return torch.cat((t_left, t, t_right), dim=-1)
+
+
+class VisionRotaryEmbeddingFast(nn.Module):
+    def __init__(
+        self,
+        dim,
+        pt_seq_len=16,
+        ft_seq_len=None,
+        custom_freqs=None,
+        freqs_for='lang',
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f'unknown modality {freqs_for}')
+
+        if ft_seq_len is None:
+            ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs = torch.einsum('..., f -> ... f', t, freqs)
+        freqs = repeat(freqs, '... n -> ... (n r)', r=2)
+        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
+
+        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
+        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
+
+        self.register_buffer("freqs_cos", freqs_cos)
+        self.register_buffer("freqs_sin", freqs_sin)
+
+        # print('======== shape of rope freq', self.freqs_cos.shape,freqs_sin.shape, '========')
+
+    def forward(self, t):
+        # print('======== shape of t', t.shape, '========')
+        _, _, Lt, _ = t.shape # B, num_heads, L, dim
+        L, _ = self.freqs_cos.shape # L, dim
+        repeat = Lt // L
+        freqs_cos, freqs_sin = self.freqs_cos, self.freqs_sin
+        if repeat != 1:
+            freqs_cos = freqs_cos.repeat_interleave(repeat, dim=0)
+            freqs_sin = freqs_sin.repeat_interleave(repeat, dim=0)
+            # print('======== shape of rope freq', freqs_cos.shape,freqs_sin.shape, '========')
+            # print(f'======== repeat {repeat} times ========')
+            # print(f'======== shape of t {t.shape} ========')
+            # # assert the repeat is twice
+            # #assert repeat == 2, f'repeat should be 2, but got {repeat}'
+            # # check the content of the repeated freqs
+            # # the content at odd index should be the same as the content at even index
+            # assert torch.allclose(freqs_cos[::2], freqs_cos[1::2]), 'repeated freqs_cos are not the same'
+            # assert torch.allclose(freqs_sin[::2], freqs_sin[1::2]), 'repeated freqs_sin are not the same'
+        # apply repeated freqs
+        return t * freqs_cos + rotate_half(t) * freqs_sin
+
+
+class RelativePositionBias2D(nn.Module):
+    """
+    2D relative positional bias for full self-attention.
+    Creates a learnable bias table of size (2*H-1) (2*W-1) per head,
+    and a fixed index map to look up bias for any pair of token positions.
+    """
+
+    def __init__(self, height: int, width: int, num_heads: int):
+        super().__init__()
+        self.height = height
+        self.width = width
+        self.num_heads = num_heads
+
+        # Create a bias table: one bias for every possible relative offset
+        # in y ∈ [-(H-1)..(H-1)] and x ∈ [-(W-1)..(W-1)]
+        self.relative_bias_table = nn.Parameter(
+            torch.zeros((2 * height - 1) * (2 * width - 1), num_heads)
+        )
+        # Precompute a (H*W)×(H*W) index matrix of which bias entry each pair (i,j) uses
+        coords_h = torch.arange(height)
+        coords_w = torch.arange(width)
+        # meshgrid of absolute coords, shape (H*W, 2)
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'), dim=-1).view(-1, 2)
+
+        # Compute all pairwise relative coords
+        relative_coords = coords[:, None, :] - coords[None, :, :]   # shape (HW, HW, 2)
+        # shift to positive
+        relative_coords[..., 0] += height - 1  # y
+        relative_coords[..., 1] += width - 1   # x
+
+        # flatten 2D index into single index: idx = y*(2W-1) + x
+        relative_index = relative_coords[..., 0] * (2 * width - 1) + relative_coords[..., 1]
+        # register as buffer so it’s on the right device / dtype
+        self.register_buffer("relative_index", relative_index.long())
+
+    def forward(self):
+        """
+        Returns:
+           bias: Tensor of shape (1, num_heads, HW, HW)
+        to be added to the raw attention logits before softmax.
+        """
+        # Lookup and reshape to (HW, HW, num_heads)
+        bias = self.relative_bias_table[self.relative_index.view(-1)]  # (HW*HW, num_heads)
+        bias = bias.view(self.height * self.width,
+                         self.height * self.width,
+                         self.num_heads)  # (HW, HW, heads)
+        # permute to (heads, HW, HW) and add batch-dim
+        bias = bias.permute(2, 0, 1).unsqueeze(0)  # (1, heads, HW, HW)
+        return bias
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class NormAttention(nn.Module):
+    """
+    Attention module of LightningDiT.
+    """
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.,
+        proj_drop: float = 0.,
+        norm_layer: nn.Module = nn.LayerNorm,
+        fused_attn: bool = True,
+        use_rmsnorm: bool = False,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = fused_attn
+
+        if use_rmsnorm:
+            norm_layer = RMSNorm
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor, rope=None) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if rope is not None:
+            q = rope(q)
+            k = rope(k)
+
+        if self.fused_attn:
+            q = q.to(v.dtype)
+            k = k.to(v.dtype) # rope may change the q,k's dtype
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            )
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class GaussianFourierEmbedding(nn.Module):
+    """
+    Gaussian Fourier Embedding for timesteps.
+    """
+    embedding_size: int = 256
+    scale: float = 1.0
+    def __init__(self, hidden_size: int, embedding_size: int = 256, scale: float = 1.0):
+        super().__init__()
+        self.embedding_size = embedding_size
+        self.scale = scale
+        self.register_buffer('W', torch.normal(0, self.scale, (embedding_size,)))
+        self.mlp = nn.Sequential(
+            nn.Linear(embedding_size * 2, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+    def forward(self, t):
+        with torch.no_grad():
+            W = self.W # stop gradient manually
+        t = t[:, None] * W[None, :] * 2 * torch.pi
+        # Concatenate sine and cosine transformations
+        t_embed =  torch.cat([torch.sin(t), torch.cos(t)], dim=-1)
+        t_embed = self.mlp(t_embed)
+        return t_embed
+
+
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size, dropout_prob):
+        super().__init__()
+        use_cfg_embedding = dropout_prob > 0
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+    def token_drop(self, labels, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        labels = torch.where(drop_ids, self.num_classes, labels)
+        return labels
+
+    def forward(self, labels, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            labels = self.token_drop(labels, force_drop_ids)
+        embeddings = self.embedding_table(labels)
+        return embeddings

patch-reparameterization/src/stage2/transport/__init__.py

@@ -0,0 +1,79 @@
+from .transport import Transport, ModelType, WeightType, PathType, Sampler
+
+def create_transport(
+    path_type='Linear',
+    prediction="velocity",
+    loss_weight=None,
+    train_eps=None,
+    sample_eps=None,
+    time_dist_type="uniform",
+    time_dist_shift=1.0,
+    sem_dim=None,
+    recon_loss_weight=1.0,
+    normalize_channel_weight=False,
+    **kwargs
+):
+    """function for creating Transport object
+    **Note**: model prediction defaults to velocity
+    Args:
+    - path_type: type of path to use; default to linear
+    - learn_score: set model prediction to score
+    - learn_noise: set model prediction to noise
+    - velocity_weighted: weight loss by velocity weight
+    - likelihood_weighted: weight loss by likelihood weight
+    - train_eps: small epsilon for avoiding instability during training
+    - sample_eps: small epsilon for avoiding instability during sampling
+    - time_dist_type: type of time distribution to use; default to uniform
+    - time_dist_shift: shift for time distribution; default to 1.0
+    """
+
+    if prediction == "noise":
+        model_type = ModelType.NOISE
+    elif prediction == "score":
+        model_type = ModelType.SCORE
+    else:
+        model_type = ModelType.VELOCITY
+
+    if loss_weight == "velocity":
+        loss_type = WeightType.VELOCITY
+    elif loss_weight == "likelihood":
+        loss_type = WeightType.LIKELIHOOD
+    else:
+        loss_type = WeightType.NONE
+
+    path_choice = {
+        "Linear": PathType.LINEAR,
+        "GVP": PathType.GVP,
+        "VP": PathType.VP,
+        "SemanticRecon": PathType.SEMANTIC_RECON,
+    }
+
+    path_type = path_choice[path_type]
+
+    if (path_type in [PathType.VP]):
+        train_eps = 1e-5 if train_eps is None else train_eps
+        sample_eps = 1e-3 if train_eps is None else sample_eps
+    elif (path_type in [PathType.GVP, PathType.LINEAR, PathType.SEMANTIC_RECON] and model_type != ModelType.VELOCITY):
+        train_eps = 1e-3 if train_eps is None else train_eps
+        sample_eps = 1e-3 if train_eps is None else sample_eps
+    else: # velocity & [GVP, LINEAR, SEMANTIC_RECON] is stable everywhere
+        train_eps = 0
+        sample_eps = 0
+
+    use_time_shift = kwargs.get("use_time_shift", True)
+    # create flow state
+    state = Transport(
+        model_type=model_type,
+        path_type=path_type,
+        loss_type=loss_type,
+        time_dist_type=time_dist_type,
+        time_dist_shift=time_dist_shift,
+        train_eps=train_eps,
+        sample_eps=sample_eps,
+        use_time_shift=use_time_shift,
+        sem_dim=sem_dim,
+        recon_loss_weight=recon_loss_weight,
+        normalize_channel_weight=normalize_channel_weight,
+    )
+
+    return state

patch-reparameterization/src/stage2/transport/integrators.py

@@ -0,0 +1,134 @@
+import numpy as np
+import torch as th
+import torch.nn as nn
+from torchdiffeq import odeint
+from functools import partial
+from tqdm import tqdm
+import logging
+logger = logging.getLogger(__name__)
+class sde:
+    """SDE solver class"""
+    def __init__(
+        self,
+        drift,
+        diffusion,
+        *,
+        t0,
+        t1,
+        num_steps,
+        sampler_type,
+        time_dist_shift,
+    ):
+        assert t0 < t1, "SDE sampler has to be in forward time"
+
+        self.num_timesteps = num_steps
+        self.t = 1 - th.linspace(t0, t1, num_steps)
+        self.t = time_dist_shift * self.t / (1 + (time_dist_shift - 1) * self.t)
+        self.drift = drift
+        self.diffusion = diffusion
+        self.sampler_type = sampler_type
+        self.time_dist_shift = time_dist_shift
+
+    def __Euler_Maruyama_step(self, x, mean_x, t_curr, t_next, model, **model_kwargs):
+        w_cur = th.randn(x.size()).to(x)
+        t = th.ones(x.size(0)).to(x) * t_curr
+        dw = w_cur * th.sqrt(t_curr - t_next)
+        drift = self.drift(x, t, model, **model_kwargs)
+        diffusion = self.diffusion(x, t)
+        mean_x = x - drift * (t_curr - t_next)
+        x = mean_x + th.sqrt(2 * diffusion) * dw
+        return x, mean_x
+
+    def __Heun_step(self, x, _, t_curr, t_next, model, **model_kwargs):
+        w_cur = th.randn(x.size()).to(x)
+        dw = w_cur * th.sqrt(t_curr - t_next)
+        diffusion = self.diffusion(x, th.ones(x.size(0)).to(x) * t_curr)
+        xhat = x + th.sqrt(2 * diffusion) * dw
+        K1 = self.drift(
+            xhat, th.ones(x.size(0)).to(x) * t_curr, model, **model_kwargs
+        )
+        xp = xhat - (t_curr - t_next) * K1
+        K2 = self.drift(
+            xp, th.ones(x.size(0)).to(x) * t_next, model, **model_kwargs
+        )
+        return xhat - 0.5 * (t_curr - t_next) * (K1 + K2), xhat # at last time point we do not perform the heun step
+
+    def __forward_fn(self):
+        """TODO: generalize here by adding all private functions ending with steps to it"""
+        sampler_dict = {
+            "euler": self.__Euler_Maruyama_step,
+            "heun": self.__Heun_step,
+        }
+
+        try:
+            sampler = sampler_dict[self.sampler_type]
+        except:
+            raise NotImplementedError("Smapler type not implemented.")
+
+        return sampler
+
+    def sample(self, init, model, **model_kwargs) -> tuple[th.Tensor]:
+        """forward loop of sde"""
+        x = init
+        mean_x = init
+        samples = []
+        sampler = self.__forward_fn()
+        for t_curr, t_next in zip(self.t[:-1], self.t[1:]):
+            with th.no_grad():
+                x, mean_x = sampler(x, mean_x, t_curr, t_next, model, **model_kwargs)
+                samples.append(x)
+
+        return samples
+
+class ode:
+    """ODE solver class"""
+    def __init__(
+        self,
+        drift,
+        *,
+        t0,
+        t1,
+        sampler_type,
+        num_steps,
+        atol,
+        rtol,
+        time_dist_shift,
+        use_time_shift=True
+    ):
+        # Allow t0 > t1 for reverse (data -> noise) integration
+        assert t0 != t1, "ODE sampler needs distinct endpoints"
+        if t0 > t1:
+            t0, t1 = t1, t0  # linspace(t0, t1) must go low -> high
+        self.drift = drift
+        # num_steps=1 means 1 euler step, need 2 time points (start + end)
+        num_t_points = max(num_steps, 2)
+        self.t = 1 - th.linspace(t0, t1, num_t_points)
+        if use_time_shift:
+            self.t = time_dist_shift * self.t / (1 + (time_dist_shift - 1) * self.t)
+        else:
+            logger.warning("Sampling Time shift is disabled.")
+
+        self.atol = atol
+        self.rtol = rtol
+        self.sampler_type = sampler_type
+
+    def sample(self, x, model, **model_kwargs) -> tuple[th.Tensor]:
+
+        device = x[0].device if isinstance(x, tuple) else x.device
+        def _fn(t, x):
+            t = th.ones(x[0].size(0)).to(device) * t if isinstance(x, tuple) else th.ones(x.size(0)).to(device) * t
+            model_output = self.drift(x, t, model, **model_kwargs)
+            return model_output
+
+        t = self.t.to(device)
+        atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol]
+        rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol]
+        samples = odeint(
+            _fn,
+            x,
+            t,
+            method=self.sampler_type,
+            atol=atol,
+            rtol=rtol
+        )
+        return samples

patch-reparameterization/src/stage2/transport/path.py

@@ -0,0 +1,242 @@
+import torch as th
+import numpy as np
+from functools import partial
+import logging
+logger = logging.getLogger(__name__)
+def expand_t_like_x(t, x):
+    """Function to reshape time t to broadcastable dimension of x
+    Args:
+      t: [batch_dim,], time vector
+      x: [batch_dim,...], data point
+    """
+    dims = [1] * (len(x.size()) - 1)
+    t = t.view(t.size(0), *dims)
+    return t
+
+
+#################### Coupling Plans ####################
+
+class ICPlan:
+    """Linear Coupling Plan"""
+    def __init__(self, sigma=0.0):
+        self.sigma = sigma
+
+    def compute_alpha_t(self, t):
+        """Compute the data coefficient along the path"""
+        return 1 - t, -1
+
+    def compute_sigma_t(self, t):
+        """Compute the noise coefficient along the path"""
+        return t, 1
+
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Compute the ratio between d_alpha and alpha"""
+        return -1 / (1 - t)
+
+    def compute_drift(self, x, t):
+        """We always output sde according to score parametrization; """
+        t = expand_t_like_x(t, x)
+        alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        drift = alpha_ratio * x
+        diffusion = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t
+
+        return -drift, diffusion
+
+    def compute_diffusion(self, x, t, form="constant", norm=1.0):
+        """Compute the diffusion term of the SDE
+        Args:
+          x: [batch_dim, ...], data point
+          t: [batch_dim,], time vector
+          form: str, form of the diffusion term
+          norm: float, norm of the diffusion term
+        """
+        t = expand_t_like_x(t, x)
+        choices = {
+            "constant": norm,
+            "SBDM": norm * self.compute_drift(x, t)[1],
+            "sigma": norm * self.compute_sigma_t(t)[0],
+            # "linear": norm * (1 - t),
+            "linear": norm * t,
+            "decreasing": 0.25 * (norm * th.cos(np.pi * (1 - t)) + 1) ** 2,
+            "inccreasing-decreasing": norm * th.sin(np.pi * (1 - t)) ** 2,
+        }
+
+        try:
+            diffusion = choices[form]
+        except KeyError:
+            raise NotImplementedError(f"Diffusion form {form} not implemented")
+
+        return diffusion
+
+    def get_score_from_velocity(self, velocity, x, t):
+        """Wrapper function: transfrom velocity prediction model to score
+        Args:
+            velocity: [batch_dim, ...] shaped tensor; velocity model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        mean = x
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+        score = (reverse_alpha_ratio * velocity - mean) / var
+        return score
+
+    def get_noise_from_velocity(self, velocity, x, t):
+        """Wrapper function: transfrom velocity prediction model to denoiser
+        Args:
+            velocity: [batch_dim, ...] shaped tensor; velocity model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        mean = x
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = reverse_alpha_ratio * d_sigma_t - sigma_t
+        noise = (reverse_alpha_ratio * velocity - mean) / var
+        return noise
+
+    def get_velocity_from_score(self, score, x, t):
+        """Wrapper function: transfrom score prediction model to velocity
+        Args:
+            score: [batch_dim, ...] shaped tensor; score model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        drift, var = self.compute_drift(x, t)
+        velocity = var * score - drift
+        return velocity
+
+    def compute_mu_t(self, t, x0, x1):
+        """Compute the mean of time-dependent density p_t"""
+        t = expand_t_like_x(t, x1)
+        alpha_t, _ = self.compute_alpha_t(t)
+        sigma_t, _ = self.compute_sigma_t(t)
+        return alpha_t * x1 + sigma_t * x0
+
+    def compute_xt(self, t, x0, x1):
+        """Sample xt from time-dependent density p_t; rng is required"""
+        xt = self.compute_mu_t(t, x0, x1)
+        return xt
+
+    def compute_ut(self, t, x0, x1, xt):
+        """Compute the vector field corresponding to p_t"""
+        t = expand_t_like_x(t, x1)
+        _, d_alpha_t = self.compute_alpha_t(t)
+        _, d_sigma_t = self.compute_sigma_t(t)
+        return d_alpha_t * x1 + d_sigma_t * x0
+
+    def plan(self, t, x0, x1, **kwargs):
+        """Compute (t, xt, ut).  Extra kwargs are accepted but ignored so that
+        callers can unconditionally pass semantic=/recon= without branching."""
+        xt = self.compute_xt(t, x0, x1)
+        ut = self.compute_ut(t, x0, x1, xt)
+        return t, xt, ut
+
+
+class VPCPlan(ICPlan):
+    """class for VP path flow matching"""
+
+    def __init__(self, sigma_min=0.1, sigma_max=20.0):
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.log_mean_coeff = lambda t: -0.25 * ((1 - t) ** 2) * (self.sigma_max - self.sigma_min) - 0.5 * (1 - t) * self.sigma_min
+        self.d_log_mean_coeff = lambda t: 0.5 * (1 - t) * (self.sigma_max - self.sigma_min) + 0.5 * self.sigma_min
+
+    def compute_alpha_t(self, t):
+        """Compute coefficient of x0"""
+        p_alpha_t = 2 * self.log_mean_coeff(t)
+        alpha_t = th.sqrt(1 - th.exp(p_alpha_t))
+        d_alpha_t = th.exp(p_alpha_t) * (2 * self.d_log_mean_coeff(t)) / (-2 * alpha_t)
+        return alpha_t, d_alpha_t
+
+    def compute_sigma_t(self, t):
+        """Compute coefficient of x1"""
+        sigma_t = self.log_mean_coeff(t)
+        sigma_t = th.exp(sigma_t)
+        d_sigma_t = sigma_t * self.d_log_mean_coeff(t)
+        return sigma_t, d_sigma_t
+
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        return d_alpha_t / alpha_t
+
+    def compute_drift(self, x, t):
+        """Compute the drift term of the SDE"""
+        t = expand_t_like_x(t, x)
+        beta_t = self.sigma_min + (1 - t) * (self.sigma_max - self.sigma_min)
+        return -0.5 * beta_t * x, beta_t / 2
+
+
+class GVPCPlan(ICPlan):
+    def __init__(self, sigma=0.0):
+        super().__init__(sigma)
+
+    def compute_alpha_t(self, t):
+        """Compute coefficient of x1"""
+        alpha_t = th.cos(t * np.pi / 2)
+        d_alpha_t = -np.pi / 2 * th.sin(t * np.pi / 2)
+        return alpha_t, d_alpha_t
+
+    def compute_sigma_t(self, t):
+        """Compute coefficient of x0"""
+        sigma_t = th.sin(t * np.pi / 2)
+        d_sigma_t = np.pi / 2 * th.cos(t * np.pi / 2)
+        return sigma_t, d_sigma_t
+
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
+        return -np.pi / 2 * th.tan(t * np.pi / 2)
+
+class SemanticReconPlan(ICPlan):
+    """Flow Matching plan with a time-dependent endpoint.
+
+    Instead of a fixed x1, the target endpoint varies with time:
+        x1(t) = semantic + (1 - t) * recon
+
+    where ``semantic`` (s) and ``recon`` (z) are the normalised semantic and
+    reconstruction components of the latent, satisfying s + z = x1_standard.
+
+    Derived interpolant and velocity field:
+        xt  = (1-t) * x1(t) + t * x0
+            = (1-t) * [s + (1-t)*z] + t * x0
+
+        ut  = d(xt)/dt
+            = x0 - s - 2*(1-t)*z
+
+    When ``semantic`` or ``recon`` is None (e.g. at inference time or when the
+    feature is unavailable), the plan falls back to the standard ICPlan
+    behaviour so that existing code paths are not broken.
+    """
+
+    def plan(self, t, x0, x1, semantic=None, recon=None):
+        """Compute (t, xt, ut) for training.
+
+        Args:
+            t:        [B] time samples in [0, 1]
+            x0:       [B, ...] noise samples ~ N(0, I)
+            x1:       [B, ...] data samples (= semantic + recon, normalised)
+            semantic: [B, ...] normalised semantic component s  (optional)
+            recon:    [B, ...] normalised recon component z     (optional)
+
+        Returns:
+            t, xt, ut  – all [B, ...]
+        """
+        # logger.debug(f"sementic: {semantic}")
+        # logger.debug(f"recon: {recon}")
+        if semantic is None or recon is None:
+            # Graceful fallback: standard linear interpolant
+            return super().plan(t, x0, x1)
+
+        t_exp = expand_t_like_x(t, semantic)           # [B, 1, 1, ...] broadcastable
+        x1_t  = semantic + (1 - t_exp) * recon         # time-dependent endpoint
+        xt    = (1 - t_exp) * x1_t + t_exp * x0        # interpolant
+        ut    = x0 - semantic - 2 * (1 - t_exp) * recon  # analytic velocity field
+        return t, xt, ut

patch-reparameterization/src/stage2/transport/transport.py

@@ -0,0 +1,614 @@
+import torch as th
+import numpy as np
+import logging
+
+import enum
+
+from . import path
+from .utils import EasyDict, log_state, mean_flat
+from .integrators import ode, sde
+import logging
+logger = logging.getLogger(__name__)
+
+class ModelType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    NOISE = enum.auto()  # the model predicts epsilon
+    SCORE = enum.auto()  # the model predicts \nabla \log p(x)
+    VELOCITY = enum.auto()  # the model predicts v(x)
+
+class PathType(enum.Enum):
+    """
+    Which type of path to use.
+    """
+
+    LINEAR = enum.auto()
+    GVP = enum.auto()
+    VP = enum.auto()
+    SEMANTIC_RECON = enum.auto()  # time-dependent endpoint: x1(t) = s + (1-t)*z
+
+class WeightType(enum.Enum):
+    """
+    Which type of weighting to use.
+    """
+
+    NONE = enum.auto()
+    VELOCITY = enum.auto()
+    LIKELIHOOD = enum.auto()
+
+
+def truncated_logitnormal_sample(
+    shape, mu, sigma, low=0.0, high=1.0
+):
+    """
+    Samples X in (0,1) with Z = logit(X) ~ Normal(mu, sigma^2), truncated so X in [low, high].
+    Works for scalars or tensors mu/sigma/low/high with broadcasting.
+
+    Args:
+        shape: output batch shape (e.g., (N,) or (N,M)). Leave () to broadcast to mu.shape.
+        mu, sigma: tensors or floats (sigma > 0).
+        low, high: truncation bounds in [0,1]. (low can be 0, high can be 1).
+        device, dtype: optional overrides.
+
+    Returns:
+        Tensor of samples with shape = broadcast(shape, mu.shape, ...)
+    """
+    mu   = th.as_tensor(mu)
+    sigma= th.as_tensor(sigma)
+    low  = th.as_tensor(low)
+    high = th.as_tensor(high)
+
+    # Map truncation bounds to logit space; handles 0/1 → ±inf automatically.
+    z_low  = th.logit(low)   # = -inf if low==0
+    z_high = th.logit(high)  # = +inf if high==1
+
+    # Standardize bounds for the base Normal(0,1)
+    base = th.distributions.Normal(th.zeros_like(mu), th.ones_like(sigma))
+    alpha = (z_low  - mu) / sigma
+    beta  = (z_high - mu) / sigma
+
+    # Truncated-normal inverse CDF sampling:
+    # U ~ Uniform(Φ(alpha), Φ(beta));  Z = mu + sigma * Φ^{-1}(U);  X = sigmoid(Z)
+    cdf_alpha = base.cdf(alpha)
+    cdf_beta  = base.cdf(beta)
+
+    # Draw uniforms on the truncated interval
+    out_shape = th.broadcast_shapes(shape, mu.shape, sigma.shape, low.shape, high.shape)
+    U = th.rand(out_shape, device=mu.device, dtype=mu.dtype)
+    U = cdf_alpha + (cdf_beta - cdf_alpha) * U.clamp_(0, 1)
+
+    Z = mu + sigma * base.icdf(U)
+    X = th.sigmoid(Z)
+
+    # Numerical safety when low/high are extremely close; clamp back into [low, high].
+    return X.clamp(low, high)
+
+
+class Transport:
+
+    def __init__(
+        self,
+        *,
+        model_type,
+        path_type,
+        loss_type,
+        time_dist_type,
+        time_dist_shift,
+        train_eps,
+        sample_eps,
+        use_time_shift,
+        sem_dim=None,
+        recon_loss_weight=1.0,
+        normalize_channel_weight=False,
+    ):
+        path_options = {
+            PathType.LINEAR: path.ICPlan,
+            PathType.GVP: path.GVPCPlan,
+            PathType.VP: path.VPCPlan,
+            PathType.SEMANTIC_RECON: path.SemanticReconPlan,
+        }
+
+        self.loss_type = loss_type
+        self.model_type = model_type
+        self.time_dist_type = time_dist_type
+        self.time_dist_shift = time_dist_shift
+        assert self.time_dist_shift >= 1.0, "time distribution shift must be >= 1.0."
+        self.path_sampler = path_options[path_type]()
+        self.train_eps = train_eps
+        self.sample_eps = sample_eps
+        self.use_time_shift = use_time_shift
+        self.sem_dim = sem_dim
+        self.recon_loss_weight = float(recon_loss_weight)
+        self.normalize_channel_weight = normalize_channel_weight
+        if self.sem_dim is not None:
+            logger.info(f"Channel-wise loss: sem_dim={self.sem_dim}, recon_loss_weight={self.recon_loss_weight}, normalize={self.normalize_channel_weight}")
+
+    def prior_logp(self, z):
+        '''
+            Standard multivariate normal prior
+            Assume z is batched
+        '''
+        shape = th.tensor(z.size())
+        N = th.prod(shape[1:])
+        _fn = lambda x: -N / 2. * np.log(2 * np.pi) - th.sum(x ** 2) / 2.
+        return th.vmap(_fn)(z)
+
+
+    def check_interval(
+        self,
+        train_eps,
+        sample_eps,
+        *,
+        diffusion_form="SBDM",
+        sde=False,
+        reverse=False,
+        eval=False,
+        last_step_size=0.0,
+    ):
+        t0 = 0
+        t1 = 1 - 1 / 1000
+        eps = train_eps if not eval else sample_eps
+        if (type(self.path_sampler) in [path.VPCPlan]):
+
+            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
+
+        elif (type(self.path_sampler) in [path.ICPlan, path.GVPCPlan, path.SemanticReconPlan]) \
+            and (self.model_type != ModelType.VELOCITY or sde): # avoid numerical issue by taking a first semi-implicit step
+
+            t0 = eps if (diffusion_form == "SBDM" and sde) or self.model_type != ModelType.VELOCITY else 0
+            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
+
+        if reverse:
+            t0, t1 = 1 - t0, 1 - t1
+
+        return t0, t1
+
+
+    def sample(self, x1):
+        """Sampling x0 & t based on shape of x1 (if needed)
+          Args:
+            x1 - data point; [batch, *dim]
+        """
+
+        x0 = th.randn_like(x1)
+        dist_options = self.time_dist_type.split("_")
+        t0, t1 = self.check_interval(self.train_eps, self.sample_eps)
+        if dist_options[0] == "uniform":
+            t = th.rand((x1.shape[0],)) * (t1 - t0) + t0
+            # print('UNIFORM IS CALLED')
+        elif dist_options[0] == "logit-normal":
+            assert len(dist_options) == 3, "Logit-normal distribution must specify the mean and variance."
+            mu, sigma = float(dist_options[1]), float(dist_options[2])
+            assert sigma > 0, "Logit-normal distribution must have positive variance."
+            t = truncated_logitnormal_sample(
+                (x1.shape[0],), mu=mu, sigma=sigma, low=t0, high=t1
+            )
+            # print('LOGITNORMAL IS CALLED')
+        elif dist_options[0] == "fixed":
+            # Fixed t for 1-step training (e.g., "fixed_1.0")
+            fixed_t = float(dist_options[1]) if len(dist_options) > 1 else 1.0
+            t = th.full((x1.shape[0],), fixed_t)
+        else:
+            raise NotImplementedError(f"Unknown time distribution type {self.time_dist_type}")
+
+        t = t.to(x1)
+        # logger.debug(f"t before shift: {t}")
+        #sqrt_size_ratio = 1 / self.time_dist_shift # already sqrted
+        if self.use_time_shift:
+            t = self.time_dist_shift * t / (1 + (self.time_dist_shift - 1) * t)
+        # logger.debug(f"t after shift: {t}")
+        return t, x0, x1
+
+
+    def training_losses(
+        self,
+        model,
+        x1,
+        model_kwargs=None
+    ):
+        """Loss for training the score model
+        Args:
+        - model: backbone model; could be score, noise, or velocity
+        - x1: datapoint
+        - model_kwargs: additional arguments for the model
+        """
+        if model_kwargs == None:
+            model_kwargs = {}
+
+        # Extract auxiliary loss parameters so they aren't passed to the model's forward pass
+        aux_losses_cfg = model_kwargs.pop('aux_losses_cfg', {})
+        latent_var = model_kwargs.pop('latent_var', None)
+
+        # Pre-emptively pop all potential projection modules to avoid passing them to the model
+        proj_modules = {}
+        keys_to_pop = [k for k in model_kwargs.keys() if k.endswith('_proj')]
+        for k in keys_to_pop:
+            proj_name = k.replace('_proj', '')
+            proj_modules[proj_name] = model_kwargs.pop(k)
+
+        # Handle legacy 'semetic_proj' key if present
+        if 'semetic_proj' in model_kwargs:
+            proj_modules['semantic'] = model_kwargs.pop('semetic_proj')
+
+        # Pop semantic/recon components for SemanticReconPlan path construction.
+        # These must NOT be forwarded to the model's forward().
+        path_semantic = model_kwargs.pop('path_semantic', None)
+        path_recon    = model_kwargs.pop('path_recon', None)
+
+        logger.debug(f"x1 shape{x1.shape}")
+        t, x0, x1 = self.sample(x1)
+        t, xt, ut = self.path_sampler.plan(t, x0, x1, semantic=path_semantic, recon=path_recon)
+        logger.debug(f"xt shape{xt.shape}")
+        model_output = model(xt, t, **model_kwargs)
+        logger.debug(f"model output shape{model_output.shape}")
+        B, *_, C = xt.shape
+        assert model_output.size() == (B, *xt.size()[1:-1], C)
+
+        terms = {}
+        terms['pred'] = model_output
+        terms['xt'] = xt
+        terms['t'] = t
+        terms['x0'] = x0
+        if self.model_type == ModelType.VELOCITY:
+            mse = (model_output - ut) ** 2
+            if self.sem_dim is not None and self.sem_dim > 0:
+                # Channel dim is dim=1 for [B,C,H,W] or dim=-1 for [B,N,C]
+                if mse.dim() == 4:  # [B, C, H, W]
+                    terms['loss_sem'] = mean_flat(mse[:, :self.sem_dim])
+                    terms['loss_recon'] = mean_flat(mse[:, self.sem_dim:])
+                    if self.recon_loss_weight != 1.0:
+                        C = mse.shape[1]
+                        recon_dim = C - self.sem_dim
+                        sem_w = 1.0
+                        recon_w = self.recon_loss_weight
+                        if self.normalize_channel_weight:
+                            norm_factor = C / (self.sem_dim * sem_w + recon_dim * recon_w)
+                            sem_w *= norm_factor
+                            recon_w *= norm_factor
+                        weight = th.ones(1, C, 1, 1, device=mse.device)
+                        weight[:, :self.sem_dim] = sem_w
+                        weight[:, self.sem_dim:] = recon_w
+                        mse = mse * weight
+                else:  # [B, N, C]
+                    terms['loss_sem'] = mean_flat(mse[..., :self.sem_dim])
+                    terms['loss_recon'] = mean_flat(mse[..., self.sem_dim:])
+                    if self.recon_loss_weight != 1.0:
+                        C = mse.shape[-1]
+                        recon_dim = C - self.sem_dim
+                        sem_w = 1.0
+                        recon_w = self.recon_loss_weight
+                        if self.normalize_channel_weight:
+                            norm_factor = C / (self.sem_dim * sem_w + recon_dim * recon_w)
+                            sem_w *= norm_factor
+                            recon_w *= norm_factor
+                        weight = th.ones(1, 1, C, device=mse.device)
+                        weight[..., :self.sem_dim] = sem_w
+                        weight[..., self.sem_dim:] = recon_w
+                        mse = mse * weight
+            terms['loss'] = mean_flat(mse)
+        else:
+            _, drift_var = self.path_sampler.compute_drift(xt, t)
+            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, xt))
+            if self.loss_type in [WeightType.VELOCITY]:
+                weight = (drift_var / sigma_t) ** 2
+            elif self.loss_type in [WeightType.LIKELIHOOD]:
+                weight = drift_var / (sigma_t ** 2)
+            elif self.loss_type in [WeightType.NONE]:
+                weight = 1
+            else:
+                raise NotImplementedError()
+
+            if self.model_type == ModelType.NOISE:
+                terms['loss'] = mean_flat(weight * ((model_output - x0) ** 2))
+            else:
+                terms['loss'] = mean_flat(weight * ((model_output * sigma_t + x0) ** 2))
+
+        # --- Compute Auxiliary Losses (e.g., Semantic Loss) ---
+        if self.model_type == ModelType.VELOCITY and aux_losses_cfg:
+            # Apply base weight to the original loss
+            base_weight = aux_losses_cfg.get('base_weight', 1.0)
+            logger.debug(f"flow loss: {terms['loss'].detach().mean().item()}")
+            terms['loss'] = terms['loss'] * base_weight
+
+            # Prepare unnormalized velocity if any projection loss is enabled
+            # We check if any proj loss is enabled to avoid unnecessary computation
+            proj_losses_cfg = aux_losses_cfg.get('proj_losses', {})
+            need_unnorm = any(
+                isinstance(cfg, dict) and cfg.get('enabled', False)
+                for cfg in proj_losses_cfg.values()
+            )
+
+            v_pred_unnorm, v_gt_unnorm = None, None
+            if need_unnorm and latent_var is not None:
+                logger.debug(f"model_output shape: {model_output.shape}")
+                logger.debug(f"latent_var shape: {latent_var.shape}")
+
+                # 1. Convert model_output and ut to [B, N, C] if they are [B, C, H, W]
+                if model_output.dim() == 4:
+                    B, C, H, W = model_output.shape
+                    N = H * W
+                    # [B, C, H, W] -> [B, C, N] -> [B, N, C]
+                    v_pred_flat = model_output.view(B, C, N).transpose(1, 2)
+                    v_gt_flat = ut.view(B, C, N).transpose(1, 2)
+                else:
+                    v_pred_flat = model_output
+                    v_gt_flat = ut
+                    B, N, C = v_pred_flat.shape
+
+                # 2. Reshape latent_var to [1, N, C] or [1, 1, C]
+                if latent_var.dim() == 1:
+                    # [C] -> [1, 1, C]
+                    scale = th.sqrt(latent_var + 1e-5).view(1, 1, -1)
+                elif latent_var.dim() == 2:
+                    # [N, C] -> [1, N, C]
+                    scale = th.sqrt(latent_var + 1e-5).unsqueeze(0)
+                elif latent_var.dim() == 3:
+                    # [C, H, W] -> [C, N] -> [N, C] -> [1, N, C]
+                    scale = th.sqrt(latent_var + 1e-5).view(latent_var.shape[0], -1).transpose(0, 1).unsqueeze(0)
+                else:
+                    scale = th.sqrt(latent_var + 1e-5)
+
+                # 3. Unnormalize
+                v_pred_unnorm = v_pred_flat * scale
+                v_gt_unnorm = v_gt_flat * scale
+
+            # --- Loop over projection losses ---
+            # This makes it easy to add more projection losses in the future
+            for proj_name, proj_cfg in proj_losses_cfg.items():
+                if isinstance(proj_cfg, dict) and proj_cfg.get('enabled', False):
+                    # Get the projection module from our pre-popped dict
+                    proj_module = proj_modules.get(proj_name)
+
+                    if proj_module is not None:
+                        if v_pred_unnorm is None:
+                            logger.warning(f"{proj_name} loss enabled but latent_var is missing.")
+                            continue
+
+                        # Map to projection space
+                        proj_v_pred = proj_module(v_pred_unnorm)
+                        proj_v_gt = proj_module(v_gt_unnorm)
+
+                        # Compute MSE
+                        proj_loss = mean_flat(((proj_v_pred - proj_v_gt) ** 2))
+
+                        logger.debug(f"{proj_name}_loss: {proj_loss.detach().mean().item()}")
+                        weight = proj_cfg.get('weight', 1.0)
+
+                        # Accumulate to total loss
+                        terms['loss'] = terms['loss'] + weight * proj_loss
+                        terms[f'{proj_name}_loss'] = proj_loss.detach()
+                    else:
+                        logger.warning(f"{proj_name} loss enabled but {proj_name}_proj not found in model_kwargs")
+        # logger.debug(terms)
+        return terms
+
+
+    def get_drift(
+        self
+    ):
+        """member function for obtaining the drift of the probability flow ODE"""
+        def score_ode(x, t, model, **model_kwargs):
+            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
+            model_output = model(x, t, **model_kwargs)
+            return (-drift_mean + drift_var * model_output) # by change of variable
+
+        def noise_ode(x, t, model, **model_kwargs):
+            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
+            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))
+            model_output = model(x, t, **model_kwargs)
+            score = model_output / -sigma_t
+            return (-drift_mean + drift_var * score)
+
+        def velocity_ode(x, t, model, **model_kwargs):
+            model_output = model(x, t, **model_kwargs)
+            return model_output
+
+        if self.model_type == ModelType.NOISE:
+            drift_fn = noise_ode
+        elif self.model_type == ModelType.SCORE:
+            drift_fn = score_ode
+        else:
+            drift_fn = velocity_ode
+
+        def body_fn(x, t, model, **model_kwargs):
+            model_output = drift_fn(x, t, model, **model_kwargs)
+            assert model_output.shape == x.shape, "Output shape from ODE solver must match input shape"
+            return model_output
+
+        return body_fn
+
+
+    def get_score(
+        self,
+    ):
+        """member function for obtaining score of
+            x_t = alpha_t * x + sigma_t * eps"""
+        if self.model_type == ModelType.NOISE:
+            score_fn = lambda x, t, model, **kwargs: model(x, t, **kwargs) / -self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))[0]
+        elif self.model_type == ModelType.SCORE:
+            score_fn = lambda x, t, model, **kwagrs: model(x, t, **kwagrs)
+        elif self.model_type == ModelType.VELOCITY:
+            score_fn = lambda x, t, model, **kwargs: self.path_sampler.get_score_from_velocity(model(x, t, **kwargs), x, t)
+        else:
+            raise NotImplementedError()
+
+        return score_fn
+
+
+class Sampler:
+    """Sampler class for the transport model"""
+    def __init__(
+        self,
+        transport,
+    ):
+        """Constructor for a general sampler; supporting different sampling methods
+        Args:
+        - transport: an tranport object specify model prediction & interpolant type
+        """
+
+        self.transport = transport
+        self.drift = self.transport.get_drift()
+        self.score = self.transport.get_score()
+
+    def __get_sde_diffusion_and_drift(
+        self,
+        *,
+        diffusion_form="SBDM",
+        diffusion_norm=1.0,
+    ):
+
+        def sde_diffusion_fn(x, t):
+            diffusion = self.transport.path_sampler.compute_diffusion(x, t, form=diffusion_form, norm=diffusion_norm)
+            return diffusion
+
+        def sde_drift_fn(x, t, model, **kwargs):
+            drift_mean = self.drift(x, t, model, **kwargs) - sde_diffusion_fn(x, t) * self.score(x, t, model, **kwargs)
+            return drift_mean
+
+
+        return sde_drift_fn, sde_diffusion_fn
+
+    def __get_last_step(
+        self,
+        sde_drift,
+        *,
+        last_step,
+        last_step_size,
+    ):
+        """Get the last step function of the SDE solver"""
+
+        if last_step is None:
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x
+        elif last_step == "Mean":
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x - sde_drift(x, t, model, **model_kwargs) * last_step_size
+        elif last_step == "Tweedie":
+            alpha = self.transport.path_sampler.compute_alpha_t # simple aliasing; the original name was too long
+            sigma = self.transport.path_sampler.compute_sigma_t
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x / alpha(t)[0][0] + (sigma(t)[0][0] ** 2) / alpha(t)[0][0] * self.score(x, t, model, **model_kwargs)
+        elif last_step == "Euler":
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x - self.drift(x, t, model, **model_kwargs) * last_step_size
+        else:
+            raise NotImplementedError()
+
+        return last_step_fn
+
+    def sample_sde(
+        self,
+        *,
+        sampling_method="Euler",
+        diffusion_form="SBDM",
+        diffusion_norm=1.0,
+        last_step="Mean",
+        last_step_size=0.04,
+        num_steps=250,
+    ):
+        """returns a sampling function with given SDE settings
+        Args:
+        - sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama
+        - diffusion_form: function form of diffusion coefficient; default to be matching SBDM
+        - diffusion_norm: function magnitude of diffusion coefficient; default to 1
+        - last_step: type of the last step; default to identity
+        - last_step_size: size of the last step; default to match the stride of 250 steps over [0,1]
+        - num_steps: total integration step of SDE
+        """
+
+        if last_step is None:
+            last_step_size = 0.0
+
+        sde_drift, sde_diffusion = self.__get_sde_diffusion_and_drift(
+            diffusion_form=diffusion_form,
+            diffusion_norm=diffusion_norm,
+        )
+
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            diffusion_form=diffusion_form,
+            sde=True,
+            eval=True,
+            reverse=False,
+            last_step_size=last_step_size,
+        )
+
+        _sde = sde(
+            sde_drift,
+            sde_diffusion,
+            t0=t0,
+            t1=t1,
+            num_steps=num_steps,
+            sampler_type=sampling_method,
+            time_dist_shift=self.transport.time_dist_shift,
+        )
+
+        last_step_fn = self.__get_last_step(sde_drift, last_step=last_step, last_step_size=last_step_size)
+
+
+        def _sample(init, model, **model_kwargs):
+            xs = _sde.sample(init, model, **model_kwargs)
+            ts = th.ones(init.size(0), device=init.device) * (1 - t1)
+            x = last_step_fn(xs[-1], ts, model, **model_kwargs)
+            xs.append(x)
+
+            assert len(xs) == num_steps, "Samples does not match the number of steps"
+
+            return xs
+
+        return _sample
+
+    def sample_ode(
+        self,
+        *,
+        sampling_method="dopri5",
+        num_steps=50,
+        atol=1e-6,
+        rtol=1e-3,
+        reverse=False,
+        **kwargs
+    ):
+        """returns a sampling function with given ODE settings
+        Args:
+        - sampling_method: type of sampler used in solving the ODE; default to be Dopri5
+        - num_steps:
+            - fixed solver (Euler, Heun): the actual number of integration steps performed
+            - adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
+        - atol: absolute error tolerance for the solver
+        - rtol: relative error tolerance for the solver
+        - reverse: whether solving the ODE in reverse (data to noise); default to False
+        """
+        if reverse:
+            drift = lambda x, t, model, **kwargs: self.drift(x, th.ones_like(t) * (1 - t), model, **kwargs)
+        else:
+            drift = self.drift
+
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            sde=False,
+            eval=True,
+            reverse=reverse,
+            last_step_size=0.0,
+        )
+        # print(t0, t1)
+        _ode = ode(
+            drift=drift,
+            t0=t0,
+            t1=t1,
+            sampler_type=sampling_method,
+            num_steps=num_steps,
+            atol=atol,
+            rtol=rtol,
+            time_dist_shift=self.transport.time_dist_shift,
+            use_time_shift=kwargs.get("use_time_shift", True),
+        )
+
+        return _ode.sample

patch-reparameterization/src/stage2/transport/utils.py

@@ -0,0 +1,29 @@
+import torch as th
+
+class EasyDict:
+
+    def __init__(self, sub_dict):
+        for k, v in sub_dict.items():
+            setattr(self, k, v)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+def mean_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return th.mean(x, dim=list(range(1, len(x.size()))))
+
+def log_state(state):
+    result = []
+
+    sorted_state = dict(sorted(state.items()))
+    for key, value in sorted_state.items():
+        # Check if the value is an instance of a class
+        if "<object" in str(value) or "object at" in str(value):
+            result.append(f"{key}: [{value.__class__.__name__}]")
+        else:
+            result.append(f"{key}: {value}")
+
+    return '\n'.join(result)

patch-reparameterization/src/stage2/transport/xpred.py

@@ -0,0 +1,131 @@
+import torch
+from torch import nn
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def make_xpred_sampler(num_steps: int, sampling_method: str = "euler", t_eps: float = 1e-3):
+    """
+    Create a JiT-style x-pred sampler operating in latent space.
+
+    The underlying model is assumed to output x_pred given (z_t, t, **kwargs),
+    and we convert it to a velocity field via:
+
+        v_pred = (x_pred - z_t) / (1 - t)
+
+    The returned object has the same callable interface as the existing
+    Transport Sampler ODE/SDE samplers:
+
+        xs = sampler(z_init, model_fn, **model_kwargs)
+        final_z = xs[-1]
+    """
+    method = sampling_method.lower()
+    if method not in ("euler", "heun"):
+        raise NotImplementedError(f"x-pred sampler only supports euler/heun, got {sampling_method}")
+
+    def _forward_velocity(z: torch.Tensor, t: torch.Tensor, model_fn, model_kwargs):
+        """
+        Map x-pred model output to velocity:
+
+            v_pred = (x_pred - z) / (1 - t)
+        """
+        bsz = z.size(0)
+        if t.ndim == 0:
+            t = t.expand(bsz)
+        t = t.to(z.device)
+        t_view = t.view(bsz, *([1] * (z.ndim - 1)))
+        x_pred = model_fn(z, t, **model_kwargs)
+        v_pred = (x_pred - z) / (1.0 - t_view).clamp_min(t_eps)
+        return v_pred
+
+    def _sample(init: torch.Tensor, model_fn, **model_kwargs):
+        """
+        Fixed-step ODE sampler (Euler / Heun) for JiT-style x-pred velocity.
+
+        Args:
+            init: initial latent z_0
+            model_fn: callable taking (z_t, t, **model_kwargs) and returning x_pred
+            model_kwargs: passed through to model_fn
+        """
+        z = init
+        xs = []
+        bsz = z.size(0)
+        device = z.device
+        ts = torch.linspace(0.0, 1.0, num_steps + 1, device=device)
+
+        for i in range(num_steps):
+            t = ts[i].expand(bsz)
+            dt = ts[i + 1] - ts[i]
+
+            if method == "euler":
+                v_pred = _forward_velocity(z, t, model_fn, model_kwargs)
+                z = z + dt * v_pred
+            else:  # heun
+                v_t = _forward_velocity(z, t, model_fn, model_kwargs)
+                z_euler = z + dt * v_t
+                t_next = ts[i + 1].expand(bsz)
+                v_next = _forward_velocity(z_euler, t_next, model_fn, model_kwargs)
+                v_pred = 0.5 * (v_t + v_next)
+                z = z + dt * v_pred
+
+            xs.append(z)
+
+        return xs
+
+    return _sample
+
+
+def compute_xpred_loss(
+    model: nn.Module,
+    x: torch.Tensor,
+    labels: torch.Tensor,
+    *,
+    transport=None,
+    t_eps: float = 1e-3,
+    sample_t_method: str = "transport",
+) -> torch.Tensor:
+    """
+    JiT-style x-pred loss in latent space.
+
+    Given latent "clean" target x, we construct:
+
+        z_t = t * x + (1 - t) * e
+        v    = (x      - z_t) / (1 - t)
+        x̂    = model(z_t, t, y=labels)
+        v̂    = (x_hat - z_t) / (1 - t)
+        L    = E[ ||v - v̂||^2 ]
+    """
+    device = x.device
+    bsz = x.size(0)
+
+    # If a Transport object is provided, reuse its time sampler so that
+    # x-pred shares the same time distribution (uniform/logit-normal,
+    # eps truncation, time_dist_shift, etc.) as the underlying FM setup.
+    # Linear plan is parameterized data->noise; JiT is noise->data,
+    # so we flip t -> 1 - t to keep the *geometric* end with higher
+    # sampling density consistent.
+    if sample_t_method == "transport" and transport is not None:
+        # transport.sample(x1) -> (t, x0, x1)
+        t, _, _ = transport.sample(x)
+        t = 1.0 - t
+    elif sample_t_method == "normal_sigmoid":
+        t = torch.randn(bsz, device=device)*0.8 - 0.8
+        t = torch.sigmoid(t)
+    else:
+        raise NotImplementedError(f"Invalid sample_t_method: {sample_t_method}")
+    logger.debug(f"t: {t}")
+    t_view = t.view(bsz, *([1] * (x.ndim - 1)))
+    e = torch.randn_like(x)
+    z_t = t_view * x + (1.0 - t_view) * e
+
+    v = (x - z_t) / (1.0 - t_view).clamp_min(t_eps)
+    x_pred = model(z_t, t, y=labels)
+    v_pred = (x_pred - z_t) / (1.0 - t_view).clamp_min(t_eps)
+
+    loss = (v - v_pred) ** 2
+    logger.debug(f"xpred loss: {loss}")
+    return loss.mean()
+

requirements.txt

@@ -1,18 +1,15 @@
-# UniSpace inference dependencies
-# Do NOT list gradio, spaces, huggingface_hub — preinstalled and platform-managed
-torch
-torchvision
-transformers
-accelerate>=1.12
-safetensors>=0.7
-omegaconf
-timm
-einops
-numpy>=1.26,<3
-scipy
-scikit-image
-pillow
+# Matches the verified UniSpace environment (unispace/requirements.txt):
+# the repo vendors custom Qwen3 / Qwen3-VL modeling code that tracks
+# transformers 4.57.x internals, so both torch and transformers are pinned.
+torch==2.8.0
+torchvision==0.23.0
+transformers==4.57.3
+accelerate==1.12.0
+omegaconf==2.3.0
+safetensors>=0.4.3
+einops>=0.8.1
+timm>=1.0.13
 sentencepiece
-PyYAML>=6.0
 tqdm
-opencv-python
+numpy>=1.26,<3
+pillow

unispace/LICENSE-APACHE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2025 Bytedance Ltd. and/or its affiliates
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+   implied. See the License for the specific language governing
+   permissions and limitations under the License.

unispace/NOTICE

@@ -0,0 +1,29 @@
+unispace (unified multimodal model)
+====================================
+
+This subdirectory contains code originally extracted and refactored from
+the ByteDance Bagel project (https://github.com/ByteDance-Seed/Bagel).
+
+Copyright 2025 Bytedance Ltd. and/or its affiliates.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+The original Apache-2.0 license and copyright notices in the source
+files (header lines `Copyright 2025 Bytedance Ltd. and/or its affiliates`
+/ `SPDX-License-Identifier: Apache-2.0`) are retained as required by the
+license.
+
+Modifications and additions made in this subdirectory, as well as all
+other components of this repository (including the `patch-reparameterization/`
+subdirectory), are licensed under the MIT License — see the repository
+root LICENSE file.

unispace/modeling/__init__.py

@@ -1 +1,4 @@
-# UniSpace modeling package — no auto-imports to avoid circular dependencies
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+from . import unimm, siglip, autoencoder

unispace/modeling/autoencoder.py

@@ -0,0 +1,382 @@
+# Copyright (c) 2024 Black Forest Labs.
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
+#
+# Original file was released under Apache-2.0, with the full license text
+# available at https://github.com/black-forest-labs/flux/blob/main/LICENSE.
+#
+# This modified file is released under the same license.
+
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+from safetensors.torch import load_file as load_sft
+
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int
+    in_channels: int
+    downsample: int
+    ch: int
+    out_ch: int
+    ch_mult: list[int]
+    num_res_blocks: int
+    z_channels: int
+    scale_factor: float
+    shift_factor: float
+
+
+def swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        # downsampling
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True, chunk_dim: int = 1):
+        super().__init__()
+        self.sample = sample
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        if self.sample:
+            std = torch.exp(0.5 * logvar)
+            return mean + std * torch.randn_like(mean)
+        else:
+            return mean
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__()
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian()
+
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = self.reg(self.encoder(x))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+    if len(missing) > 0 and len(unexpected) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+        print("\n" + "-" * 79 + "\n")
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    elif len(missing) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+    elif len(unexpected) > 0:
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+
+
+def load_ae(local_path: str, vae_type: str = "flux1-vae", rae_config_path: str = None, **kwargs):
+    """
+    Load VAE model.
+
+    Args:
+        local_path: Path to the VAE checkpoint file or directory.
+        vae_type: 'flux1-vae' (z=16), 'flux2-vae' (z=32), 'rae' (z=768),
+                  or 'qwen3_unified' (Qwen3Unified encoder, z=1280).
+        rae_config_path: Path to RAE config yaml (used when vae_type='rae' or 'qwen3_unified').
+
+    Returns:
+        (model, config) tuple. Config has .z_channels and .downsample attributes.
+    """
+    if vae_type == "flux2-vae":
+        from modeling.autoencoder_flux2 import load_flux2_ae
+        return load_flux2_ae(local_path)
+
+    if vae_type in ("rae", "qwen3_unified"):
+        from modeling.autoencoder_rae import load_rae
+        return load_rae(local_path, rae_config_path=rae_config_path,
+                        rae_accelerate_root=kwargs.get('rae_accelerate_root', None),
+                        encoder_only=kwargs.get('encoder_only', False))
+
+    ae_params = AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            downsample=8,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+    )
+
+    # Loading the autoencoder
+    ae = AutoEncoder(ae_params)
+
+    if local_path is not None:
+        sd = load_sft(local_path)
+        missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
+        print_load_warning(missing, unexpected)
+    return ae, ae_params

unispace/modeling/autoencoder_flux2.py

@@ -0,0 +1,412 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+# Adapted from diffusers 0.36 AutoencoderKLFlux2, standalone implementation.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Self-contained Flux2 VAE implementation that loads official BFL weights
+# (diffusion_pytorch_model.safetensors) without requiring diffusers >= 0.36.
+# Key names match diffusers convention exactly.
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file as load_sft
+
+
+# ---------------------------------------------------------------------------
+# Config dataclass (mirrors AutoEncoderParams from autoencoder.py for Flux1)
+# ---------------------------------------------------------------------------
+@dataclass
+class Flux2VAEConfig:
+    in_channels: int = 3
+    out_channels: int = 3
+    block_out_channels: tuple = (128, 256, 512, 512)
+    layers_per_block: int = 2
+    latent_channels: int = 32
+    norm_num_groups: int = 32
+    downsample: int = 8          # 2^(num_down_blocks - 1) = 2^3 = 8
+    z_channels: int = 32         # alias for latent_channels, compat with Flux1 config
+    batch_norm_eps: float = 1e-4
+    batch_norm_momentum: float = 0.1
+    patch_size: tuple = (2, 2)
+
+
+# ---------------------------------------------------------------------------
+# Building blocks
+# ---------------------------------------------------------------------------
+class ResnetBlock2D(nn.Module):
+    """Standard ResNet block: GroupNorm -> SiLU -> Conv -> GroupNorm -> SiLU -> Conv + shortcut."""
+
+    def __init__(self, in_channels: int, out_channels: int, groups: int = 32):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(groups, in_channels, eps=1e-6)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = nn.GroupNorm(groups, out_channels, eps=1e-6)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        self.nonlinearity = nn.SiLU()
+        if in_channels != out_channels:
+            self.conv_shortcut = nn.Conv2d(in_channels, out_channels, 1)
+        else:
+            self.conv_shortcut = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        h = self.nonlinearity(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = self.nonlinearity(h)
+        h = self.conv2(h)
+        if self.conv_shortcut is not None:
+            x = self.conv_shortcut(x)
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    """Self-attention block matching diffusers Attention with _from_deprecated_attn_block=True."""
+
+    def __init__(self, channels: int, num_groups: int = 32):
+        super().__init__()
+        self.group_norm = nn.GroupNorm(num_groups, channels, eps=1e-6)
+        self.to_q = nn.Linear(channels, channels)
+        self.to_k = nn.Linear(channels, channels)
+        self.to_v = nn.Linear(channels, channels)
+        self.to_out = nn.ModuleList([nn.Linear(channels, channels)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        b, c, h, w = x.shape
+        x = self.group_norm(x)
+        x = x.reshape(b, c, h * w).permute(0, 2, 1)  # (B, HW, C)
+
+        q = self.to_q(x)
+        k = self.to_k(x)
+        v = self.to_v(x)
+
+        # scaled_dot_product_attention
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = self.to_out[0](x)
+
+        x = x.permute(0, 2, 1).reshape(b, c, h, w)
+        return x + residual
+
+
+class Downsample2D(nn.Module):
+    """Spatial downsample with padding + stride-2 conv."""
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.pad(x, (0, 1, 0, 1), mode="constant", value=0)
+        return self.conv(x)
+
+
+class Upsample2D(nn.Module):
+    """Spatial upsample with nearest interpolation + conv."""
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        return self.conv(x)
+
+
+# ---------------------------------------------------------------------------
+# Encoder / Decoder blocks (key names: down_blocks.{i}.resnets.{j} / downsamplers.0)
+# ---------------------------------------------------------------------------
+class DownEncoderBlock2D(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, num_layers: int = 2,
+                 add_downsample: bool = True, groups: int = 32):
+        super().__init__()
+        self.resnets = nn.ModuleList()
+        for i in range(num_layers):
+            self.resnets.append(ResnetBlock2D(in_ch if i == 0 else out_ch, out_ch, groups))
+        self.downsamplers = nn.ModuleList()
+        if add_downsample:
+            self.downsamplers.append(Downsample2D(out_ch))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for resnet in self.resnets:
+            x = resnet(x)
+        for ds in self.downsamplers:
+            x = ds(x)
+        return x
+
+
+class UpDecoderBlock2D(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, num_layers: int = 3,
+                 add_upsample: bool = True, groups: int = 32):
+        super().__init__()
+        self.resnets = nn.ModuleList()
+        for i in range(num_layers):
+            self.resnets.append(ResnetBlock2D(in_ch if i == 0 else out_ch, out_ch, groups))
+        self.upsamplers = nn.ModuleList()
+        if add_upsample:
+            self.upsamplers.append(Upsample2D(out_ch))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for resnet in self.resnets:
+            x = resnet(x)
+        for us in self.upsamplers:
+            x = us(x)
+        return x
+
+
+class MidBlock(nn.Module):
+    """Mid block: resnet -> attention -> resnet."""
+
+    def __init__(self, channels: int, groups: int = 32):
+        super().__init__()
+        self.resnets = nn.ModuleList([
+            ResnetBlock2D(channels, channels, groups),
+            ResnetBlock2D(channels, channels, groups),
+        ])
+        self.attentions = nn.ModuleList([AttnBlock(channels, groups)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.resnets[0](x)
+        x = self.attentions[0](x)
+        x = self.resnets[1](x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Encoder & Decoder
+# ---------------------------------------------------------------------------
+class Flux2Encoder(nn.Module):
+    def __init__(self, cfg: Flux2VAEConfig):
+        super().__init__()
+        ch = cfg.block_out_channels
+        self.conv_in = nn.Conv2d(cfg.in_channels, ch[0], 3, padding=1)
+
+        self.down_blocks = nn.ModuleList()
+        in_c = ch[0]
+        for i, out_c in enumerate(ch):
+            is_last = (i == len(ch) - 1)
+            self.down_blocks.append(
+                DownEncoderBlock2D(in_c, out_c, cfg.layers_per_block,
+                                   add_downsample=not is_last, groups=cfg.norm_num_groups)
+            )
+            in_c = out_c
+
+        self.mid_block = MidBlock(ch[-1], cfg.norm_num_groups)
+
+        self.conv_norm_out = nn.GroupNorm(cfg.norm_num_groups, ch[-1], eps=1e-6)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(ch[-1], 2 * cfg.latent_channels, 3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv_in(x)
+        for block in self.down_blocks:
+            x = block(x)
+        x = self.mid_block(x)
+        x = self.conv_norm_out(x)
+        x = self.conv_act(x)
+        x = self.conv_out(x)
+        return x
+
+
+class Flux2Decoder(nn.Module):
+    def __init__(self, cfg: Flux2VAEConfig):
+        super().__init__()
+        ch = cfg.block_out_channels
+        rev_ch = list(reversed(ch))
+
+        self.conv_in = nn.Conv2d(cfg.latent_channels, ch[-1], 3, padding=1)
+        self.mid_block = MidBlock(ch[-1], cfg.norm_num_groups)
+
+        self.up_blocks = nn.ModuleList()
+        in_c = rev_ch[0]
+        for i, out_c in enumerate(rev_ch):
+            is_last = (i == len(rev_ch) - 1)
+            # Decoder has layers_per_block + 1 resnets per block
+            self.up_blocks.append(
+                UpDecoderBlock2D(in_c, out_c, cfg.layers_per_block + 1,
+                                 add_upsample=not is_last, groups=cfg.norm_num_groups)
+            )
+            in_c = out_c
+
+        self.conv_norm_out = nn.GroupNorm(cfg.norm_num_groups, ch[0], eps=1e-6)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(ch[0], cfg.out_channels, 3, padding=1)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        z = self.conv_in(z)
+        z = self.mid_block(z)
+        for block in self.up_blocks:
+            z = block(z)
+        z = self.conv_norm_out(z)
+        z = self.conv_act(z)
+        z = self.conv_out(z)
+        return z
+
+
+# ---------------------------------------------------------------------------
+# DiagonalGaussian (sample or mode)
+# ---------------------------------------------------------------------------
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True):
+        super().__init__()
+        self.sample = sample
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=1)
+        logvar = torch.clamp(logvar, -30.0, 20.0)
+        if self.sample:
+            std = torch.exp(0.5 * logvar)
+            return mean + std * torch.randn_like(mean)
+        return mean
+
+
+# ---------------------------------------------------------------------------
+# Main model: AutoEncoderFlux2
+# ---------------------------------------------------------------------------
+class AutoEncoderFlux2(nn.Module):
+    """
+    Standalone Flux2 VAE (AutoencoderKLFlux2) implementation.
+
+    Key differences from Flux1 VAE (autoencoder.py):
+      - latent_channels = 32 (vs 16)
+      - Has quant_conv / post_quant_conv layers
+      - Uses BatchNorm-based normalization instead of scale_factor/shift_factor
+      - Architecture follows diffusers standard (DownEncoderBlock2D / UpDecoderBlock2D)
+    """
+
+    def __init__(self, config: Flux2VAEConfig | None = None):
+        super().__init__()
+        if config is None:
+            config = Flux2VAEConfig()
+        self.config = config
+
+        self.encoder = Flux2Encoder(config)
+        self.decoder = Flux2Decoder(config)
+        self.quant_conv = nn.Conv2d(2 * config.latent_channels, 2 * config.latent_channels, 1)
+        self.post_quant_conv = nn.Conv2d(config.latent_channels, config.latent_channels, 1)
+        self.reg = DiagonalGaussian(sample=False)  # default: mode (deterministic)
+
+        # BatchNorm for latent normalization (patchified space)
+        bn_channels = math.prod(config.patch_size) * config.latent_channels  # 2*2*32=128
+        self.bn = nn.BatchNorm2d(
+            bn_channels,
+            eps=config.batch_norm_eps,
+            momentum=config.batch_norm_momentum,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode image to normalized latent (BN in patchified space, then un-patchify)."""
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        z = self.reg(h)  # (B, C=32, H, W)
+
+        # Patchify → BN normalize → un-patchify to keep API compatible with Flux1
+        B, C, H, W = z.shape
+        pH, pW = self.config.patch_size  # (2, 2)
+        # (B, C, H, W) → (B, C*pH*pW, H//pH, W//pW)
+        z_p = z.reshape(B, C, H // pH, pH, W // pW, pW)
+        z_p = z_p.permute(0, 1, 3, 5, 2, 4).reshape(B, C * pH * pW, H // pH, W // pW)
+        # BN normalize
+        bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(z_p.device, z_p.dtype)
+        bn_std = torch.sqrt(
+            self.bn.running_var.view(1, -1, 1, 1).to(z_p.device, z_p.dtype)
+            + self.config.batch_norm_eps
+        )
+        z_p = (z_p - bn_mean) / bn_std
+        # Un-patchify: (B, C*pH*pW, H', W') → (B, C, H, W)
+        z = z_p.reshape(B, C, pH, pW, H // pH, W // pW)
+        z = z.permute(0, 1, 4, 2, 5, 3).reshape(B, C, H, W)
+        return z
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Decode normalized latent to image (symmetric with encode).
+
+        Pipeline: patchify → BN denormalize → unpatchify → post_quant_conv → decoder.
+        This mirrors the encode path which does: encoder → patchify → BN normalize → unpatchify.
+        """
+        B, C, H, W = z.shape
+        pH, pW = self.config.patch_size  # (2, 2)
+
+        # Patchify: (B, C, H, W) → (B, C*pH*pW, H//pH, W//pW)
+        z_p = z.reshape(B, C, H // pH, pH, W // pW, pW)
+        z_p = z_p.permute(0, 1, 3, 5, 2, 4).reshape(B, C * pH * pW, H // pH, W // pW)
+
+        # BN denormalize
+        bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(z_p.device, z_p.dtype)
+        bn_std = torch.sqrt(
+            self.bn.running_var.view(1, -1, 1, 1).to(z_p.device, z_p.dtype)
+            + self.config.batch_norm_eps
+        )
+        z_p = z_p * bn_std + bn_mean
+
+        # Un-patchify: (B, C*pH*pW, H', W') → (B, C, H, W)
+        z = z_p.reshape(B, C, pH, pW, H // pH, W // pW)
+        z = z.permute(0, 1, 4, 2, 5, 3).reshape(B, C, H, W)
+
+        z = self.post_quant_conv(z)
+        return self.decoder(z)
+
+    def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """Apply BN normalization to patchified latents (as in Flux2Pipeline)."""
+        bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+        bn_std = torch.sqrt(
+            self.bn.running_var.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+            + self.config.batch_norm_eps
+        )
+        return (latents - bn_mean) / bn_std
+
+    def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """Reverse BN normalization for decoding."""
+        bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+        bn_std = torch.sqrt(
+            self.bn.running_var.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+            + self.config.batch_norm_eps
+        )
+        return latents * bn_std + bn_mean
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.decode(self.encode(x))
+
+
+# ---------------------------------------------------------------------------
+# Loading helper
+# ---------------------------------------------------------------------------
+def _print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+    if missing:
+        print(f"[Flux2-VAE] {len(missing)} missing keys:\n\t" + "\n\t".join(missing[:20]))
+    if unexpected:
+        print(f"[Flux2-VAE] {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected[:20]))
+
+
+def load_flux2_ae(local_path: str) -> tuple[AutoEncoderFlux2, Flux2VAEConfig]:
+    """
+    Load Flux2 VAE from a safetensors checkpoint directory or file.
+
+    Args:
+        local_path: Path to the safetensors file, or directory containing
+                    diffusion_pytorch_model.safetensors.
+
+    Returns:
+        (model, config) tuple.
+    """
+    import os
+    if os.path.isdir(local_path):
+        sf_path = os.path.join(local_path, "diffusion_pytorch_model.safetensors")
+    else:
+        sf_path = local_path
+
+    config = Flux2VAEConfig()
+    model = AutoEncoderFlux2(config)
+
+    sd = load_sft(sf_path)
+    missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+    _print_load_warning(missing, unexpected)
+
+    return model, config

unispace/modeling/siglip/__init__.py

@@ -0,0 +1,98 @@
+# Copyright 2024 The HuggingFace Inc. team.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import TYPE_CHECKING
+
+from transformers.utils import (
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    is_sentencepiece_available,
+    is_torch_available,
+    is_vision_available,
+)
+
+
+_import_structure = {
+    "configuration_siglip": [
+        "SiglipConfig",
+        "SiglipTextConfig",
+        "SiglipVisionConfig",
+    ],
+    "processing_siglip": ["SiglipProcessor"],
+}
+
+try:
+    if not is_sentencepiece_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["tokenization_siglip"] = ["SiglipTokenizer"]
+
+
+try:
+    if not is_vision_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["image_processing_siglip"] = ["SiglipImageProcessor"]
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_siglip"] = [
+        "SiglipModel",
+        "SiglipPreTrainedModel",
+        "SiglipTextModel",
+        "SiglipVisionModel",
+        "SiglipForImageClassification",
+    ]
+
+
+if TYPE_CHECKING:
+    from .configuration_siglip import (
+        SiglipConfig,
+        SiglipTextConfig,
+        SiglipVisionConfig,
+    )
+    from .processing_siglip import SiglipProcessor
+
+    try:
+        if not is_sentencepiece_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .tokenization_siglip import SiglipTokenizer
+
+    try:
+        if not is_vision_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .image_processing_siglip import SiglipImageProcessor
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_siglip import (
+            SiglipForImageClassification,
+            SiglipModel,
+            SiglipPreTrainedModel,
+            SiglipTextModel,
+            SiglipVisionModel,
+        )
+
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

unispace/modeling/siglip/configuration_siglip.py

@@ -0,0 +1,287 @@
+# Copyright 2024 The HuggingFace Inc. team.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Siglip model configuration"""
+
+import os
+from typing import Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class SiglipTextConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`SiglipTextModel`]. It is used to instantiate a
+    Siglip text encoder according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the text encoder of the Siglip
+    [google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the Siglip text model. Defines the number of different tokens that can be represented by
+            the `inputs_ids` passed when calling [`SiglipModel`].
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        max_position_embeddings (`int`, *optional*, defaults to 64):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        pad_token_id (`int`, *optional*, defaults to 1):
+            The id of the padding token in the vocabulary.
+        bos_token_id (`int`, *optional*, defaults to 49406):
+            The id of the beginning-of-sequence token in the vocabulary.
+        eos_token_id (`int`, *optional*, defaults to 49407):
+            The id of the end-of-sequence token in the vocabulary.
+
+    Example:
+
+    ```python
+    >>> from transformers import SiglipTextConfig, SiglipTextModel
+
+    >>> # Initializing a SiglipTextConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipTextConfig()
+
+    >>> # Initializing a SiglipTextModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipTextModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "siglip_text_model"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        max_position_embeddings=64,
+        hidden_act="gelu_pytorch_tanh",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        # This differs from `CLIPTokenizer`'s default and from openai/siglip
+        # See https://github.com/huggingface/transformers/pull/24773#issuecomment-1632287538
+        pad_token_id=1,
+        bos_token_id=49406,
+        eos_token_id=49407,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.attention_dropout = attention_dropout
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the text config dict if we are loading from SiglipConfig
+        if config_dict.get("model_type") == "siglip":
+            config_dict = config_dict["text_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class SiglipVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`SiglipVisionModel`]. It is used to instantiate a
+    Siglip vision encoder according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the vision encoder of the Siglip
+    [google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_channels (`int`, *optional*, defaults to 3):
+            Number of channels in the input images.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        patch_size (`int`, *optional*, defaults to 16):
+            The size (resolution) of each patch.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    Example:
+
+    ```python
+    >>> from transformers import SiglipVisionConfig, SiglipVisionModel
+
+    >>> # Initializing a SiglipVisionConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipVisionConfig()
+
+    >>> # Initializing a SiglipVisionModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipVisionModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "siglip_vision_model"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        image_size=224,
+        patch_size=16,
+        hidden_act="gelu_pytorch_tanh",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the vision config dict if we are loading from SiglipConfig
+        if config_dict.get("model_type") == "siglip":
+            config_dict = config_dict["vision_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class SiglipConfig(PretrainedConfig):
+    r"""
+    [`SiglipConfig`] is the configuration class to store the configuration of a [`SiglipModel`]. It is used to
+    instantiate a Siglip model according to the specified arguments, defining the text model and vision model configs.
+    Instantiating a configuration with the defaults will yield a similar configuration to that of the Siglip
+    [google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        text_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`SiglipTextConfig`].
+        vision_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`SiglipVisionConfig`].
+        kwargs (*optional*):
+            Dictionary of keyword arguments.
+
+    Example:
+
+    ```python
+    >>> from transformers import SiglipConfig, SiglipModel
+
+    >>> # Initializing a SiglipConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipConfig()
+
+    >>> # Initializing a SiglipModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+
+    >>> # We can also initialize a SiglipConfig from a SiglipTextConfig and a SiglipVisionConfig
+    >>> from transformers import SiglipTextConfig, SiglipVisionConfig
+
+    >>> # Initializing a SiglipText and SiglipVision configuration
+    >>> config_text = SiglipTextConfig()
+    >>> config_vision = SiglipVisionConfig()
+
+    >>> config = SiglipConfig.from_text_vision_configs(config_text, config_vision)
+    ```"""
+
+    model_type = "siglip"
+
+    def __init__(self, text_config=None, vision_config=None, **kwargs):
+        super().__init__(**kwargs)
+
+        if text_config is None:
+            text_config = {}
+            logger.info("`text_config` is `None`. Initializing the `SiglipTextConfig` with default values.")
+
+        if vision_config is None:
+            vision_config = {}
+            logger.info("`vision_config` is `None`. initializing the `SiglipVisionConfig` with default values.")
+
+        self.text_config = SiglipTextConfig(**text_config)
+        self.vision_config = SiglipVisionConfig(**vision_config)
+
+        self.initializer_factor = 1.0
+
+    @classmethod
+    def from_text_vision_configs(cls, text_config: SiglipTextConfig, vision_config: SiglipVisionConfig, **kwargs):
+        r"""
+        Instantiate a [`SiglipConfig`] (or a derived class) from siglip text model configuration and siglip vision
+        model configuration.
+
+        Returns:
+            [`SiglipConfig`]: An instance of a configuration object
+        """
+
+        return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)

unispace/modeling/siglip/convert_siglip_to_hf.py

@@ -0,0 +1,401 @@
+# Copyright 2024 The HuggingFace Inc. team.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Convert SigLIP checkpoints from the original repository.
+
+URL: https://github.com/google-research/big_vision/tree/main
+"""
+
+import argparse
+import collections
+from pathlib import Path
+
+import numpy as np
+import requests
+import torch
+from huggingface_hub import hf_hub_download
+from numpy import load
+from PIL import Image
+
+from transformers import SiglipConfig, SiglipImageProcessor, SiglipModel, SiglipProcessor, SiglipTokenizer
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+model_name_to_checkpoint = {
+    # base checkpoints
+    "siglip-base-patch16-224": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_224_63724782.npz",
+    "siglip-base-patch16-256": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_256_60500360.npz",
+    "siglip-base-patch16-384": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_384_68578854.npz",
+    "siglip-base-patch16-512": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_512_68580893.npz",
+    # large checkpoints
+    "siglip-large-patch16-256": "/Users/nielsrogge/Documents/SigLIP/webli_en_l16_256_60552751.npz",
+    "siglip-large-patch16-384": "/Users/nielsrogge/Documents/SigLIP/webli_en_l16_384_63634585.npz",
+    # multilingual checkpoint
+    "siglip-base-patch16-256-i18n": "/Users/nielsrogge/Documents/SigLIP/webli_i18n_b16_256_66117334.npz",
+    # so400m checkpoints
+    "siglip-so400m-patch14-384": "/Users/nielsrogge/Documents/SigLIP/webli_en_so400m_384_58765454.npz",
+}
+
+model_name_to_image_size = {
+    "siglip-base-patch16-224": 224,
+    "siglip-base-patch16-256": 256,
+    "siglip-base-patch16-384": 384,
+    "siglip-base-patch16-512": 512,
+    "siglip-large-patch16-256": 256,
+    "siglip-large-patch16-384": 384,
+    "siglip-base-patch16-256-i18n": 256,
+    "siglip-so400m-patch14-384": 384,
+}
+
+
+def get_siglip_config(model_name):
+    config = SiglipConfig()
+
+    vocab_size = 250000 if "i18n" in model_name else 32000
+    image_size = model_name_to_image_size[model_name]
+    patch_size = 16 if "patch16" in model_name else 14
+
+    # size of the architecture
+    config.vision_config.image_size = image_size
+    config.vision_config.patch_size = patch_size
+    config.text_config.vocab_size = vocab_size
+
+    if "base" in model_name:
+        pass
+    elif "large" in model_name:
+        config.text_config.hidden_size = 1024
+        config.text_config.intermediate_size = 4096
+        config.text_config.num_hidden_layers = 24
+        config.text_config.num_attention_heads = 16
+        config.vision_config.hidden_size = 1024
+        config.vision_config.intermediate_size = 4096
+        config.vision_config.num_hidden_layers = 24
+        config.vision_config.num_attention_heads = 16
+    elif "so400m" in model_name:
+        config.text_config.hidden_size = 1152
+        config.text_config.intermediate_size = 4304
+        config.text_config.num_hidden_layers = 27
+        config.text_config.num_attention_heads = 16
+        config.vision_config.hidden_size = 1152
+        config.vision_config.intermediate_size = 4304
+        config.vision_config.num_hidden_layers = 27
+        config.vision_config.num_attention_heads = 16
+    else:
+        raise ValueError("Model not supported")
+
+    return config
+
+
+def create_rename_keys(config):
+    rename_keys = []
+    # fmt: off
+
+    # vision encoder
+
+    rename_keys.append(("params/img/embedding/kernel", "vision_model.embeddings.patch_embedding.weight"))
+    rename_keys.append(("params/img/embedding/bias", "vision_model.embeddings.patch_embedding.bias"))
+    rename_keys.append(("params/img/pos_embedding", "vision_model.embeddings.position_embedding.weight"))
+
+    for i in range(config.vision_config.num_hidden_layers):
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_0/scale", f"vision_model.encoder.layers.{i}.layer_norm1.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_0/bias", f"vision_model.encoder.layers.{i}.layer_norm1.bias"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_1/scale", f"vision_model.encoder.layers.{i}.layer_norm2.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_1/bias", f"vision_model.encoder.layers.{i}.layer_norm2.bias"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_0/kernel", f"vision_model.encoder.layers.{i}.mlp.fc1.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_0/bias", f"vision_model.encoder.layers.{i}.mlp.fc1.bias"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_1/kernel", f"vision_model.encoder.layers.{i}.mlp.fc2.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_1/bias", f"vision_model.encoder.layers.{i}.mlp.fc2.bias"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/key/kernel", f"vision_model.encoder.layers.{i}.self_attn.k_proj.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/key/bias", f"vision_model.encoder.layers.{i}.self_attn.k_proj.bias"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/value/kernel", f"vision_model.encoder.layers.{i}.self_attn.v_proj.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/value/bias", f"vision_model.encoder.layers.{i}.self_attn.v_proj.bias"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/query/kernel", f"vision_model.encoder.layers.{i}.self_attn.q_proj.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/query/bias", f"vision_model.encoder.layers.{i}.self_attn.q_proj.bias"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/out/kernel", f"vision_model.encoder.layers.{i}.self_attn.out_proj.weight"))
+        rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/out/bias", f"vision_model.encoder.layers.{i}.self_attn.out_proj.bias"))
+
+    rename_keys.append(("params/img/Transformer/encoder_norm/scale", "vision_model.post_layernorm.weight"))
+    rename_keys.append(("params/img/Transformer/encoder_norm/bias", "vision_model.post_layernorm.bias"))
+
+    rename_keys.append(("params/img/MAPHead_0/probe", "vision_model.head.probe"))
+    rename_keys.append(("params/img/MAPHead_0/LayerNorm_0/scale", "vision_model.head.layernorm.weight"))
+    rename_keys.append(("params/img/MAPHead_0/LayerNorm_0/bias", "vision_model.head.layernorm.bias"))
+    rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_0/kernel", "vision_model.head.mlp.fc1.weight"))
+    rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_0/bias", "vision_model.head.mlp.fc1.bias"))
+    rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_1/kernel", "vision_model.head.mlp.fc2.weight"))
+    rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_1/bias", "vision_model.head.mlp.fc2.bias"))
+    rename_keys.append(("params/img/MAPHead_0/MultiHeadDotProductAttention_0/out/kernel", "vision_model.head.attention.out_proj.weight"))
+    rename_keys.append(("params/img/MAPHead_0/MultiHeadDotProductAttention_0/out/bias", "vision_model.head.attention.out_proj.bias"))
+
+    # text encoder
+
+    rename_keys.append(("params/txt/Embed_0/embedding", "text_model.embeddings.token_embedding.weight"))
+    rename_keys.append(("params/txt/pos_embedding", "text_model.embeddings.position_embedding.weight"))
+
+    for i in range(config.text_config.num_hidden_layers):
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_0/scale", f"text_model.encoder.layers.{i}.layer_norm1.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_0/bias", f"text_model.encoder.layers.{i}.layer_norm1.bias"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_1/scale", f"text_model.encoder.layers.{i}.layer_norm2.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_1/bias", f"text_model.encoder.layers.{i}.layer_norm2.bias"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_0/kernel", f"text_model.encoder.layers.{i}.mlp.fc1.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_0/bias", f"text_model.encoder.layers.{i}.mlp.fc1.bias"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_1/kernel", f"text_model.encoder.layers.{i}.mlp.fc2.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_1/bias", f"text_model.encoder.layers.{i}.mlp.fc2.bias"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/key/kernel", f"text_model.encoder.layers.{i}.self_attn.k_proj.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/key/bias", f"text_model.encoder.layers.{i}.self_attn.k_proj.bias"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/value/kernel", f"text_model.encoder.layers.{i}.self_attn.v_proj.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/value/bias", f"text_model.encoder.layers.{i}.self_attn.v_proj.bias"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/query/kernel", f"text_model.encoder.layers.{i}.self_attn.q_proj.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/query/bias", f"text_model.encoder.layers.{i}.self_attn.q_proj.bias"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/out/kernel", f"text_model.encoder.layers.{i}.self_attn.out_proj.weight"))
+        rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/out/bias", f"text_model.encoder.layers.{i}.self_attn.out_proj.bias"))
+
+    rename_keys.append(("params/txt/Encoder_0/encoder_norm/scale", "text_model.final_layer_norm.weight"))
+    rename_keys.append(("params/txt/Encoder_0/encoder_norm/bias", "text_model.final_layer_norm.bias"))
+    rename_keys.append(("params/txt/head/kernel", "text_model.head.weight"))
+    rename_keys.append(("params/txt/head/bias", "text_model.head.bias"))
+
+    # learned temperature and bias
+    rename_keys.append(("params/t", "logit_scale"))
+    rename_keys.append(("params/b", "logit_bias"))
+
+    # fmt: on
+    return rename_keys
+
+
+def rename_key(dct, old, new, config):
+    val = dct.pop(old)
+
+    if ("out_proj" in new or "v_proj" in new or "k_proj" in new or "q_proj" in new) and "vision" in new:
+        val = val.reshape(-1, config.vision_config.hidden_size)
+    if ("out_proj" in new or "v_proj" in new or "k_proj" in new or "q_proj" in new) and "text" in new:
+        val = val.reshape(-1, config.text_config.hidden_size)
+
+    if "patch_embedding.weight" in new:
+        val = val.transpose(3, 2, 0, 1)
+    elif new.endswith("weight") and "position_embedding" not in new and "token_embedding" not in new:
+        val = val.T
+
+    if "position_embedding" in new and "vision" in new:
+        val = val.reshape(-1, config.vision_config.hidden_size)
+    if "position_embedding" in new and "text" in new:
+        val = val.reshape(-1, config.text_config.hidden_size)
+
+    if new.endswith("bias"):
+        val = val.reshape(-1)
+
+    dct[new] = torch.from_numpy(val)
+
+
+def read_in_q_k_v_head(state_dict, config):
+    # read in individual input projection layers
+    key_proj_weight = (
+        state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/key/kernel")
+        .reshape(-1, config.vision_config.hidden_size)
+        .T
+    )
+    key_proj_bias = state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/key/bias").reshape(-1)
+    value_proj_weight = (
+        state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/value/kernel")
+        .reshape(-1, config.vision_config.hidden_size)
+        .T
+    )
+    value_proj_bias = state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/value/bias").reshape(-1)
+    query_proj_weight = (
+        state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/query/kernel")
+        .reshape(-1, config.vision_config.hidden_size)
+        .T
+    )
+    query_proj_bias = state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/query/bias").reshape(-1)
+
+    # next, add them to the state dict as a single matrix + vector
+    state_dict["vision_model.head.attention.in_proj_weight"] = torch.from_numpy(
+        np.concatenate([query_proj_weight, key_proj_weight, value_proj_weight], axis=0)
+    )
+    state_dict["vision_model.head.attention.in_proj_bias"] = torch.from_numpy(
+        np.concatenate([query_proj_bias, key_proj_bias, value_proj_bias], axis=0)
+    )
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    image = Image.open(requests.get(url, stream=True).raw)
+    return image
+
+
+def flatten_nested_dict(params, parent_key="", sep="/"):
+    items = []
+
+    for k, v in params.items():
+        new_key = parent_key + sep + k if parent_key else k
+
+        if isinstance(v, collections.abc.MutableMapping):
+            items.extend(flatten_nested_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
+@torch.no_grad()
+def convert_siglip_checkpoint(model_name, pytorch_dump_folder_path, verify_logits=True, push_to_hub=False):
+    """
+    Copy/paste/tweak model's weights to our SigLIP structure.
+    """
+
+    # define default SigLIP configuration
+    config = get_siglip_config(model_name)
+
+    # get checkpoint
+    checkpoint = model_name_to_checkpoint[model_name]
+
+    # get vocab file
+    if "i18n" in model_name:
+        vocab_file = "/Users/nielsrogge/Documents/SigLIP/multilingual_vocab/sentencepiece.model"
+    else:
+        vocab_file = "/Users/nielsrogge/Documents/SigLIP/english_vocab/sentencepiece.model"
+
+    # load original state dict
+    data = load(checkpoint)
+    state_dict = flatten_nested_dict(data)
+
+    # remove and rename some keys
+    rename_keys = create_rename_keys(config)
+    for src, dest in rename_keys:
+        rename_key(state_dict, src, dest, config)
+
+    # qkv matrices of attention pooling head need special treatment
+    read_in_q_k_v_head(state_dict, config)
+
+    # load HuggingFace model
+    model = SiglipModel(config).eval()
+    model.load_state_dict(state_dict)
+
+    # create processor
+    # important: make tokenizer not return attention_mask since original one doesn't require it
+    image_size = config.vision_config.image_size
+    size = {"height": image_size, "width": image_size}
+    image_processor = SiglipImageProcessor(size=size)
+    tokenizer = SiglipTokenizer(vocab_file=vocab_file, model_input_names=["input_ids"])
+    processor = SiglipProcessor(image_processor=image_processor, tokenizer=tokenizer)
+
+    # verify on dummy images and texts
+    url_1 = "https://cdn.openai.com/multimodal-neurons/assets/apple/apple-ipod.jpg"
+    image_1 = Image.open(requests.get(url_1, stream=True).raw).convert("RGB")
+    url_2 = "https://cdn.openai.com/multimodal-neurons/assets/apple/apple-blank.jpg"
+    image_2 = Image.open(requests.get(url_2, stream=True).raw).convert("RGB")
+    texts = ["an apple", "a picture of an apple"]
+
+    inputs = processor(images=[image_1, image_2], text=texts, return_tensors="pt", padding="max_length")
+
+    # verify input_ids against original ones
+    if image_size == 224:
+        filename = "siglip_pixel_values.pt"
+    elif image_size == 256:
+        filename = "siglip_pixel_values_256.pt"
+    elif image_size == 384:
+        filename = "siglip_pixel_values_384.pt"
+    elif image_size == 512:
+        filename = "siglip_pixel_values_512.pt"
+    else:
+        raise ValueError("Image size not supported")
+
+    filepath = hf_hub_download(repo_id="nielsr/test-image", filename=filename, repo_type="dataset")
+    original_pixel_values = torch.load(filepath)
+    filepath = hf_hub_download(repo_id="nielsr/test-image", filename="siglip_input_ids.pt", repo_type="dataset")
+    original_input_ids = torch.load(filepath)
+
+    if "i18n" not in model_name:
+        assert inputs.input_ids.tolist() == original_input_ids.tolist()
+
+    print("Mean of original pixel values:", original_pixel_values.mean())
+    print("Mean of new pixel values:", inputs.pixel_values.mean())
+
+    # note: we're testing with original pixel values here since we don't have exact pixel values
+    with torch.no_grad():
+        outputs = model(input_ids=inputs.input_ids, pixel_values=original_pixel_values)
+
+    # with torch.no_grad():
+    #     outputs = model(input_ids=inputs.input_ids, pixel_values=inputs.pixel_values)
+
+    print(outputs.logits_per_image[:3, :3])
+
+    probs = torch.sigmoid(outputs.logits_per_image)  # these are the probabilities
+    print(f"{probs[0][0]:.1%} that image 0 is '{texts[0]}'")
+    print(f"{probs[0][1]:.1%} that image 0 is '{texts[1]}'")
+
+    if verify_logits:
+        if model_name == "siglip-base-patch16-224":
+            expected_slice = torch.tensor(
+                [[-2.9621, -2.1672], [-0.2713, 0.2910]],
+            )
+        elif model_name == "siglip-base-patch16-256":
+            expected_slice = torch.tensor(
+                [[-3.1146, -1.9894], [-0.7312, 0.6387]],
+            )
+        elif model_name == "siglip-base-patch16-384":
+            expected_slice = torch.tensor(
+                [[-2.8098, -2.1891], [-0.4242, 0.4102]],
+            )
+        elif model_name == "siglip-base-patch16-512":
+            expected_slice = torch.tensor(
+                [[-2.7899, -2.2668], [-0.4295, -0.0735]],
+            )
+        elif model_name == "siglip-large-patch16-256":
+            expected_slice = torch.tensor(
+                [[-1.5827, -0.5801], [-0.9153, 0.1363]],
+            )
+        elif model_name == "siglip-large-patch16-384":
+            expected_slice = torch.tensor(
+                [[-2.1523, -0.2899], [-0.2959, 0.7884]],
+            )
+        elif model_name == "siglip-so400m-patch14-384":
+            expected_slice = torch.tensor([[-1.2441, -0.6649], [-0.7060, 0.7374]])
+        elif model_name == "siglip-base-patch16-256-i18n":
+            expected_slice = torch.tensor(
+                [[-0.9064, 0.1073], [-0.0299, 0.5304]],
+            )
+
+        assert torch.allclose(outputs.logits_per_image[:3, :3], expected_slice, atol=1e-4)
+        print("Looks ok!")
+
+    if pytorch_dump_folder_path is not None:
+        Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
+        print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
+        model.save_pretrained(pytorch_dump_folder_path)
+        print(f"Saving processor to {pytorch_dump_folder_path}")
+        processor.save_pretrained(pytorch_dump_folder_path)
+
+    if push_to_hub:
+        model.push_to_hub(f"nielsr/{model_name}")
+        processor.push_to_hub(f"nielsr/{model_name}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Required parameters
+    parser.add_argument(
+        "--model_name",
+        default="siglip-base-patch16-224",
+        type=str,
+        choices=model_name_to_checkpoint.keys(),
+        help="Name of the model you'd like to convert.",
+    )
+    parser.add_argument(
+        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
+    )
+    parser.add_argument(
+        "--verify_logits",
+        action="store_false",
+        help="Whether to verify logits against the original implementation.",
+    )
+    parser.add_argument(
+        "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
+    )
+
+    args = parser.parse_args()
+    convert_siglip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.verify_logits, args.push_to_hub)

unispace/modeling/siglip/image_processing_siglip.py

@@ -0,0 +1,230 @@
+# Copyright 2024 The HuggingFace Inc. team.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Image processor class for SigLIP."""
+
+from typing import Dict, List, Optional, Union
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import (
+    convert_to_rgb,
+    resize,
+    to_channel_dimension_format,
+)
+from transformers.image_utils import (
+    IMAGENET_STANDARD_MEAN,
+    IMAGENET_STANDARD_STD,
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    is_scaled_image,
+    make_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+if is_vision_available():
+    import PIL
+
+
+class SiglipImageProcessor(BaseImageProcessor):
+    r"""
+    Constructs a SigLIP image processor.
+
+    Args:
+        do_resize (`bool`, *optional*, defaults to `True`):
+            Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by
+            `do_resize` in the `preprocess` method.
+        size (`Dict[str, int]` *optional*, defaults to `{"height": 224, "width": 224}`):
+            Size of the image after resizing. Can be overridden by `size` in the `preprocess` method.
+        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
+            Resampling filter to use if resizing the image. Can be overridden by `resample` in the `preprocess` method.
+        do_rescale (`bool`, *optional*, defaults to `True`):
+            Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by `do_rescale` in
+            the `preprocess` method.
+        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+            Scale factor to use if rescaling the image. Can be overridden by `rescale_factor` in the `preprocess`
+            method.
+        do_normalize (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the image by the specified mean and standard deviation. Can be overridden by
+            `do_normalize` in the `preprocess` method.
+        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+            Mean to use if normalizing the image. This is a float or list of floats the length of the number of
+            channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
+        image_std (`float` or `List[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+            Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
+            number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
+            Can be overridden by the `image_std` parameter in the `preprocess` method.
+        do_convert_rgb (`bool`, *optional*, defaults to `True`):
+            Whether to convert the image to RGB.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Dict[str, int] = None,
+        resample: PILImageResampling = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: Union[int, float] = 1 / 255,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 224, "width": 224}
+        image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
+        image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
+
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: bool = None,
+        size: Dict[str, int] = None,
+        resample: PILImageResampling = None,
+        do_rescale: bool = None,
+        rescale_factor: float = None,
+        do_normalize: bool = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: bool = None,
+    ) -> PIL.Image.Image:
+        """
+        Preprocess an image or batch of images.
+
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            size (`Dict[str, int]`, *optional*, defaults to `self.size`):
+                Size of the image after resizing.
+            resample (`int`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
+                has an effect if `do_resize` is set to `True`.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
+                `True`.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - Unset: Use the channel dimension format of the input image.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. If unset, the channel dimension format is inferred
+                from the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+        """
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, param_name="size", default_to_square=False)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+        images = make_list_of_images(images)
+
+        if not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+        # All transformations expect numpy arrays.
+        images = [to_numpy_array(image) for image in images]
+
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
+        if is_scaled_image(images[0]) and do_rescale:
+            logger.warning_once(
+                "It looks like you are trying to rescale already rescaled images. If the input"
+                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
+            )
+
+        if input_data_format is None:
+            # We assume that all images have the same channel dimension format.
+            input_data_format = infer_channel_dimension_format(images[0])
+
+        if do_resize:
+            height, width = size["height"], size["width"]
+            images = [
+                resize(image=image, size=(height, width), resample=resample, input_data_format=input_data_format)
+                for image in images
+            ]
+
+        if do_rescale:
+            images = [
+                self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
+                for image in images
+            ]
+
+        if do_normalize:
+            images = [
+                self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+                for image in images
+            ]
+
+        images = [
+            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
+        ]
+
+        data = {"pixel_values": images}
+        return BatchFeature(data=data, tensor_type=return_tensors)