Add files using upload-large-folder tool

code/umm/.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.py[cod]
+outputs/

code/umm/README.md

@@ -0,0 +1,95 @@
+# UMM 分阶段代码目录
+
+这里按训练阶段独立组织当前已经闭环的 FLUX.1 unified model 主线。每个阶段都有自己的：
+
+```text
+train.py / infer.py   Python 入口
+run_train.sh          分布式启动脚本
+config.yaml           本阶段配置
+README.md             数据、模型、loss、冻结参数、checkpoint 说明
+```
+
+没有统一总启动脚本，也不会自动串行运行所有阶段。上一个阶段验收完成后，手工把 checkpoint 路径写入下一个阶段配置，避免程序误选“最新目录”。
+
+## 目录
+
+```text
+umm/
+├── runtime/                 本地共享 models/blip3o 源码与 DeepSpeed 配置
+├── stage1_meta_query/       256 个 meta query 对齐 FLUX.1
+├── stage2_covt/             SAM/DINO/VGGT/Edge/SigLIP CoVT 课程与 LoRA merge
+├── stage3_generation/       CoVT/Query 与 DiT generation path
+├── stage4_tag_moe/          多图编辑 + TAG-MoE
+└── inference/               独立推理入口
+```
+
+## 阶段关系
+
+```text
+Stage 1
+Qwen3-VL + FLUX.1
+  -> embeddings_connector/checkpoint-N.bin
+
+Stage 2
+Qwen3-VL + Stage-1 connector + SAM/DINOv2/VGGT/PiDiNet/SigLIP teachers
+  -> CoVT LoRA/non-LoRA -> merge dense CoVT model
+
+Stage 3
+dense CoVT model + FLUX.1
+  -> dense generation model + generation projector
+
+Stage 4
+dense generation model + projector + quality500k multi-image data
+  -> TAG-MoE trainable_weights.bin
+
+Inference
+dense generation model + projector + FLUX.1 + TAG sidecar
+```
+
+## 启动方式
+
+每个 `train.py` 都是当前真实训练脚本的本地副本，不再调用仓库原位置的训练脚本。共享的 model/data/trainer 代码位于 `runtime/`，shell 会把它放在 `PYTHONPATH` 最前面，因此修改 `umm/runtime/models/blip3o` 会直接影响这些阶段。
+
+每个阶段独立启动：
+
+```bash
+bash hc_test/umm/stage1_meta_query/run_train.sh
+bash hc_test/umm/stage2_covt/run_train.sh
+bash hc_test/umm/stage3_generation/run_train.sh
+bash hc_test/umm/stage4_tag_moe/run_train.sh
+```
+
+先做命令检查、不启动训练：
+
+```bash
+DRY_RUN=1 bash hc_test/umm/stage1_meta_query/run_train.sh
+DRY_RUN=1 bash hc_test/umm/stage2_covt/run_train.sh
+DRY_RUN=1 bash hc_test/umm/stage3_generation/run_train.sh
+DRY_RUN=1 bash hc_test/umm/stage4_tag_moe/run_train.sh
+```
+
+配置项仍可通过命令行覆盖：
+
+```bash
+DRY_RUN=1 bash hc_test/umm/stage4_tag_moe/run_train.sh \
+  --learning_rate 2e-5 \
+  --output_dir /path/to/new_output
+```
+
+具体数据字段、loss 和 checkpoint 见各子目录 README。`runtime/` 只复制源码和小型配置，没有复制数据、预训练权重或 `.pth/.bin/.safetensors` checkpoint。
+
+## 运行环境
+
+当前主环境从原 `unvideo` 环境派生并持久化在：
+
+```text
+/rczhang/rczhang/conda_packs/unvideo
+```
+
+各阶段的 train/inference shell 都显式使用这里的 Python，避免机器重启后系统路径环境丢失；分布式入口通过 `python -m deepspeed.launcher.runner` 或 `python -m torch.distributed.run` 启动，不依赖可能带旧 shebang 的复制脚本。Stage 2 的 `setup_umm_env.sh` 负责克隆/校准依赖。所有入口在启动前都会清除 proxy 环境变量，只使用直连。环境与 Stage-B 代码验证方式见 `stage2_covt/README.md`。
+
+## 工作边界
+
+这个目录整理的是与你最初描述一致的 **FLUX.1 主线**：meta query -> CoVT -> generation -> TAG-MoE。仓库里的 Target-CoVT 与 FLUX.2 Klein 是独立试验分支，不是这条主线后面可以直接续训的一两个阶段，因此没有混入这里。
+
+当前 Stage 3 的 `train.py` 是较新的 `train_stage3_flux.py` 组装入口；Stage 4 的默认 dense checkpoint 则来自已经跑过的 UND/GEN 历史主线。两者的模型类和 artifact 约定还没有被项目统一。为了避免制造一个表面能串联、实际错误加载的假 pipeline，这里保留真实默认 checkpoint，并在 Stage 3、Stage 4 README 中明确 hand-off。后续要打通新 Stage 3 -> Stage 4，应先统一 Stage 4 的 `OmniGenUndForConditionalGeneration` 加载类与 condition layout。

code/umm/__init__.py

@@ -0,0 +1 @@
+"""Stage-by-stage launch layout for the latest complete FLUX.1 UMM pipeline."""

code/umm/inference/README.md

@@ -0,0 +1,108 @@
+# Inference：Stage-4 TAG-MoE 图像编辑
+
+## 需要的组件
+
+推理不是只加载一个 checkpoint，而是组合：
+
+```text
+1. merged Stage-3 dense Qwen/CoVT model
+2. FLUX.1 base（VAE + transformer + scheduler）
+3. Stage-3 generation projector
+4. Stage-4 trainable_weights.bin
+5. 与训练一致的 4 experts / top-1 / last 10 layers / shared expert
+```
+
+默认路径写在 `run_infer.sh`，不复制任何权重。
+
+## 输入
+
+支持：
+
+- source-only edit；
+- reference-only generation；
+- source + reference edit；
+- 多个 reference：用逗号分隔路径（由真实 infer 实现解析）；
+- optional target：只用于保存 comparison，不进入生成条件。
+
+## 启动
+
+Source-only：
+
+```bash
+SOURCE_IMAGE=/path/source.png \
+PROMPT='move the object to the left' \
+TASK_TYPE=object_movement \
+bash hc_test/umm/inference/run_infer.sh
+```
+
+Source + reference：
+
+```bash
+SOURCE_IMAGE=/path/source.png \
+REFERENCE_IMAGE=/path/ref.png \
+PROMPT='replace the subject using the reference identity' \
+TASK_TYPE=subject_replacement \
+bash hc_test/umm/inference/run_infer.sh
+```
+
+只打印命令：
+
+```bash
+DRY_RUN=1 SOURCE_IMAGE=/path/source.png \
+bash hc_test/umm/inference/run_infer.sh
+```
+
+覆盖 checkpoint：
+
+```bash
+CHECKPOINT=/path/checkpoint-N/trainable_weights.bin \
+SOURCE_IMAGE=/path/source.png \
+bash hc_test/umm/inference/run_infer.sh
+```
+
+## CoVT mode
+
+```text
+COVT_MODE=full    全部 anchor response
+COVT_MODE=random  随机 anchor 子集
+COVT_MODE=none    不构造 CoVT response
+```
+
+这控制 Qwen assistant 序列中的 CoVT；不代表 infer 一定把 anchor hidden 直接拼进 DiT。
+
+## 当前一致性风险
+
+当前 Stage-4 full training 中：
+
+```text
+DiT condition = anchor hidden + query hidden
+```
+
+当前 experiment-local inference 实现中：
+
+```text
+DiT condition = query hidden
+```
+
+所以它实质更接近“CoVT 间接影响 query、但不直接拼 anchor”的推理。正式使用 full-CoVT checkpoint 前应修正/统一该条件路径，再比较生成结果。
+
+## 输出
+
+```text
+outputs/
+└── single/
+    ├── generated.png
+    └── comparison.png
+```
+
+multi-case 会按 case 名分别建目录。
+
+## 建议验收
+
+固定 seed，对每个 checkpoint 至少跑：
+
+- instruction 原始 / shuffle；
+- source 原始 / remove；
+- reference 原始 / shuffle；
+- `COVT_MODE=full/none`；
+- train/infer 都 query-only 或都 full-anchor 的严格对照。

code/umm/inference/__init__.py

@@ -0,0 +1 @@
+"""Local UMM inference implementation."""

code/umm/inference/infer.py

@@ -0,0 +1,994 @@
+#!/usr/bin/env python3
+"""
+TAG-MoE (ffn_tag) inference for the quality500k-v1 experiment.
+
+That run uses (see ``configs/moe_ffn_tag.yaml`` + ``01_train.log``):
+  - ``OmniGenUndForConditionalGeneration`` (qwen_covt_gen -> qwen_omni_gen_und remap)
+  - local data.py Qwen tokenization helpers (user images + assistant CoVT/gen placeholder)
+  - ``build_omni_dit_latent_input`` + multi-bucket VAE conditioning
+  - ``inject_flux_moe_ffn_tag`` (4 experts, shared, last 10 layers)
+  - gen_down_projector from checkpoint-30000.bin (+ trainable delta in bin)
+"""
+from __future__ import annotations
+
+import argparse
+import ast
+import os
+import random
+import sys
+import types
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence
+
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from torchvision.transforms import v2
+from tqdm import tqdm
+from transformers import AutoConfig, AutoProcessor
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_RUNTIME_ROOT = Path(__file__).resolve().parents[1] / "runtime"
+sys.path.insert(0, str(_RUNTIME_ROOT))
+sys.path.insert(1, str(_REPO_ROOT))
+
+
+def _stub_model_package() -> None:
+    model_pkg_dir = _RUNTIME_ROOT / "models" / "blip3o" / "model"
+    if "models.blip3o.model" in sys.modules:
+        return
+    for name, path in (
+        ("models", _RUNTIME_ROOT / "models"),
+        ("models.blip3o", _RUNTIME_ROOT / "models" / "blip3o"),
+    ):
+        if name not in sys.modules:
+            pkg = types.ModuleType(name)
+            pkg.__path__ = [str(path)]
+            sys.modules[name] = pkg
+    stub = types.ModuleType("models.blip3o.model")
+    stub.__path__ = [str(model_pkg_dir)]
+    sys.modules["models.blip3o.model"] = stub
+
+
+_stub_model_package()
+
+import models.blip3o.model.language_model.covt_gen_qwen_stage3  # noqa: F401
+from models.blip3o.constants import DEFAULT_IM_END_TOKEN, IGNORE_INDEX, SYSTEM_MESSAGE
+from models.blip3o.data.covt_data_var import (
+    get_anchor_task_name,
+    get_anchor_token,
+    get_comt_data_in_response,
+    get_token_num,
+)
+from models.blip3o.model.language_model.omni_und_gen import (
+    OmniGenUndForConditionalGeneration,
+    QwenOmniGenUndConfig,
+)
+from models.blip3o.model.flux_moe_tag import inject_flux_moe_ffn_tag
+from models.blip3o.model.language_model.und_gen_stage3_infer import calculate_shift, retrieve_timesteps
+from hc_test.umm.stage4_tag_moe.data import (
+    TASK_TYPE_TO_ID,
+    TASK_TYPES,
+    bucket_images,
+    get_image_info,
+    get_rope_index_3,
+)
+from models.blip3o.moe_train_tag.omni_dit_latent_pack import build_omni_dit_latent_input
+
+CONDITION_NAMES = {0: "ref-only", 1: "source-only", 2: "source+ref"}
+_MIN_FREE_VRAM_GIB = 45.0
+_GEN_TRANSFORM = T.Compose([v2.ToTensor(), v2.Normalize([0.5], [0.5])])
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def _resolve_device(name: str) -> torch.device:
+    if name == "auto":
+        if not torch.cuda.is_available():
+            return torch.device("cpu")
+        best_i, best_free = 0, -1
+        for i in range(torch.cuda.device_count()):
+            free, _total = torch.cuda.mem_get_info(i)
+            if free > best_free:
+                best_free, best_i = free, i
+        free_gib = best_free / 1024**3
+        print(f"[device] auto-selected cuda:{best_i} ({free_gib:.1f} GiB free)")
+        if free_gib < _MIN_FREE_VRAM_GIB:
+            raise RuntimeError(
+                f"GPU cuda:{best_i} 仅 {free_gib:.1f} GiB 空闲，推理约需 {_MIN_FREE_VRAM_GIB:.0f} GiB+。"
+            )
+        return torch.device(f"cuda:{best_i}")
+    return torch.device(name)
+
+
+def _move_model_to_device(model, device: torch.device, dtype: torch.dtype):
+    if device.type == "cpu":
+        return model.to(device=device, dtype=dtype)
+    model = model.to(device="cpu", dtype=dtype)
+    torch.cuda.empty_cache()
+    for attr in ("dit", "gen_vision_tower", "gen_down_projector"):
+        sub = getattr(model, attr, None)
+        if sub is not None:
+            setattr(model, attr, sub.to(device=device, dtype=dtype))
+            torch.cuda.empty_cache()
+    return model.to(device=device, dtype=dtype)
+
+
+def _split_tag_checkpoint(raw: Dict[str, torch.Tensor]) -> tuple:
+    dit_w: Dict[str, torch.Tensor] = {}
+    proj_w: Dict[str, torch.Tensor] = {}
+    emb_w: Dict[str, torch.Tensor] = {}
+    for k, v in raw.items():
+        if k.startswith("stage1_model.dit."):
+            dit_w[k[len("stage1_model.dit.") :]] = v
+        elif k.startswith("dit."):
+            dit_w[k[len("dit.") :]] = v
+        elif k.startswith("stage1_model.gen_down_projector."):
+            proj_w[k[len("stage1_model.gen_down_projector.") :]] = v
+        elif k.startswith("gen_down_projector."):
+            proj_w[k[len("gen_down_projector.") :]] = v
+        elif k.startswith("tag_semantic.") or k.startswith("tag_par_head."):
+            continue
+        elif k.startswith("stage1_model."):
+            emb_w[k[len("stage1_model.") :]] = v
+        elif k.endswith("embed_tokens.weight") or k.endswith("lm_head.weight"):
+            emb_w[k] = v
+    return dit_w, proj_w, emb_w
+
+
+def _gen_placeholder(num_metaqueries: int) -> str:
+    return (
+        "\n<begin_of_img>"
+        + "".join(f"<img{i}>" for i in range(num_metaqueries))
+        + "<end_of_img>"
+    )
+
+
+def _build_gpt_response(
+    gen_content: str,
+    *,
+    inject_covt: bool,
+    covt_mode: str,
+    anchor_token_nums: Sequence[int],
+    anchor_tokens: Sequence[str],
+    anchor_task_names: Sequence[str],
+    seed: int,
+) -> str:
+    gpt_response = gen_content
+    if inject_covt and covt_mode != "none":
+        rng = random.Random(seed)
+        if covt_mode == "full":
+            gpt_response = get_comt_data_in_response(
+                gpt_response, list(anchor_token_nums), list(anchor_tokens), list(anchor_task_names)
+            )
+        elif covt_mode == "random":
+            if rng.random() < 0.5:
+                gpt_response = get_comt_data_in_response(
+                    gpt_response, list(anchor_token_nums), list(anchor_tokens), list(anchor_task_names)
+                )
+            elif rng.randint(0, 5) != 0:
+                total = len(anchor_tokens)
+                if total > 0:
+                    x = rng.randint(1, total)
+                    idxs = sorted(rng.sample(range(total), x))
+                    sel_nums = [anchor_token_nums[i] for i in idxs]
+                    sel_toks = [anchor_tokens[i] for i in idxs]
+                    sel_names = [anchor_task_names[i] for i in idxs]
+                    gpt_response = get_comt_data_in_response(gpt_response, sel_nums, sel_toks, sel_names)
+    return f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+
+
+def build_covt_encoded_input(
+    processor,
+    *,
+    prompt: str,
+    source_pil: Optional[Image.Image],
+    ref_pils: List[Image.Image],
+    num_metaqueries: int,
+    image_min_pixels: int,
+    image_max_pixels: int,
+    image_resized_width: int,
+    image_resized_height: int,
+    anchor_model_id: List[str],
+    covt_mode: str,
+    seed: int,
+    device: torch.device,
+) -> Dict[str, Any]:
+    """Mirror the experiment-local UMMMoEArrowDataset tokenization for one sample."""
+    merge_size = getattr(processor.image_processor, "merge_size", 2)
+    anchor_token_nums = get_token_num(anchor_model_id)
+    anchor_tokens = get_anchor_token(anchor_model_id)
+    anchor_task_names = get_anchor_task_name(anchor_model_id)
+
+    content: List[dict] = []
+    llm_vision_images: List[Any] = []
+    if source_pil is not None:
+        content.append({"type": "image", "image": source_pil})
+        llm_vision_images.append(
+            get_image_info(source_pil, image_min_pixels, image_max_pixels, image_resized_width, image_resized_height)
+        )
+    for ref_pil in ref_pils:
+        content.append({"type": "image", "image": ref_pil})
+        llm_vision_images.append(
+            get_image_info(ref_pil, image_min_pixels, image_max_pixels, image_resized_width, image_resized_height)
+        )
+    content.append({"type": "text", "text": prompt})
+
+    messages: List[dict] = []
+    if SYSTEM_MESSAGE:
+        messages.append({"role": "system", "content": SYSTEM_MESSAGE})
+    messages.append({"role": "user", "content": content})
+
+    user_str = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    gpt_response = _build_gpt_response(
+        _gen_placeholder(num_metaqueries),
+        inject_covt=bool(llm_vision_images),
+        covt_mode=covt_mode,
+        anchor_token_nums=anchor_token_nums,
+        anchor_tokens=anchor_tokens,
+        anchor_task_names=anchor_task_names,
+        seed=seed,
+    )
+
+    inputs = processor(
+        text=[user_str],
+        images=llm_vision_images if llm_vision_images else None,
+        padding=False,
+        return_tensors="pt",
+    )
+    prompt_input_ids = inputs["input_ids"]
+    response_input_ids = processor.tokenizer(
+        gpt_response, add_special_tokens=False, padding=False, return_tensors="pt"
+    )["input_ids"]
+    input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
+    attention_mask = (input_ids > -1000000).to(torch.long)
+    input_ids_2d = input_ids.unsqueeze(0)
+
+    encoded: Dict[str, Any] = {
+        "input_ids": input_ids.unsqueeze(0),
+        "attention_mask": attention_mask.unsqueeze(0),
+    }
+    if "pixel_values" in inputs and inputs["pixel_values"] is not None:
+        encoded["pixel_values"] = inputs["pixel_values"]
+        encoded["image_grid_thw"] = inputs["image_grid_thw"]
+
+    position_ids, _ = get_rope_index_3(
+        merge_size,
+        input_ids_2d,
+        image_grid_thw=encoded.get("image_grid_thw"),
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+    )
+    encoded["position_ids"] = position_ids
+    return {k: v.to(device) for k, v in encoded.items() if hasattr(v, "to")}
+
+
+def _zero_gen_tensor(image_size: int) -> torch.Tensor:
+    return torch.zeros(3, image_size, image_size, dtype=torch.float32)
+
+
+def prepare_dit_gen_batch(
+    *,
+    source_pil: Optional[Image.Image],
+    ref_pils: List[Image.Image],
+    image_size: int,
+    max_reference_images: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple:
+    """Build source/reference pixel values like the experiment-local collator."""
+    has_source = source_pil is not None
+    src = _GEN_TRANSFORM(source_pil.convert("RGB")) if has_source else _zero_gen_tensor(image_size)
+    source_pv = bucket_images([src], resolution=image_size).to(device=device, dtype=dtype)
+
+    ref_tensors: List[torch.Tensor] = []
+    ref_mask: List[bool] = []
+    for ref_pil in ref_pils[:max_reference_images]:
+        ref_tensors.append(_GEN_TRANSFORM(ref_pil.convert("RGB")))
+        ref_mask.append(True)
+    while len(ref_tensors) < max_reference_images:
+        ref_tensors.append(_zero_gen_tensor(image_size))
+        ref_mask.append(False)
+
+    ref_bucketed = [bucket_images([t], resolution=image_size)[0] for t in ref_tensors]
+    ref_pv = torch.stack(ref_bucketed, dim=0).unsqueeze(0).to(device=device, dtype=dtype)
+
+    height = int(source_pv.shape[-2]) if has_source else image_size
+    width = int(source_pv.shape[-1]) if has_source else image_size
+    batch = {
+        "source_pixel_values": source_pv,
+        "has_source_image": torch.tensor([has_source], device=device),
+        "reference_pixel_values": ref_pv,
+        "reference_image_mask": torch.tensor([ref_mask], device=device),
+    }
+    return batch, height, width
+
+
+def _extract_meta_hidden(model, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+    boi_id = model.gen_begin_token_idx
+    eoi_id = model.gen_end_token_idx
+    ids = input_ids[0]
+    boi_pos = int((ids == boi_id).nonzero(as_tuple=True)[0][0].item())
+    eoi_pos = int((ids == eoi_id).nonzero(as_tuple=True)[0][0].item())
+    return hidden_states[:, boi_pos + 1 : eoi_pos, :]
+
+
+@torch.no_grad()
+def _encode_meta_hidden_omni(model, encoded: dict) -> torch.Tensor:
+    """Omni training path: embed prep + ``model.model`` forward (see ``stage1_flux_wrapper``)."""
+    input_ids = encoded["input_ids"]
+    attention_mask = encoded["attention_mask"]
+    pixel_values = encoded.get("pixel_values")
+    image_grid_thw = encoded.get("image_grid_thw")
+    position_ids = encoded.get("position_ids")
+
+    inputs_embeds, visual_pos_masks, deepstack_visual_embeds = (
+        model._prepare_multimodal_inputs_embeds_for_forward(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+        )
+    )
+    qwen_kwargs: Dict[str, Any] = {
+        "input_ids": None,
+        "inputs_embeds": inputs_embeds,
+        "attention_mask": attention_mask,
+        "output_hidden_states": True,
+    }
+    if position_ids is not None:
+        qwen_kwargs["position_ids"] = position_ids
+    if visual_pos_masks is not None:
+        qwen_kwargs["visual_pos_masks"] = visual_pos_masks
+    if deepstack_visual_embeds is not None:
+        qwen_kwargs["deepstack_visual_embeds"] = deepstack_visual_embeds
+
+    qwen_out = model.model(**qwen_kwargs)
+    hidden = qwen_out.last_hidden_state if hasattr(qwen_out, "last_hidden_state") else qwen_out[0]
+    return _extract_meta_hidden(model, hidden, input_ids)
+
+
+def load_omni_tag_model(args):
+    print(f"[load] OmniGenUnd from {args.model_path}")
+    compute_dtype = torch.bfloat16
+    attn_impl = "flash_attention_2" if not args.disable_flash_attn2 else "sdpa"
+
+    cfg = AutoConfig.from_pretrained(args.model_path)
+    src_model_type = getattr(cfg, "model_type", None)
+    if src_model_type != QwenOmniGenUndConfig.model_type:
+        cfg_dict = cfg.to_dict()
+        cfg_dict["model_type"] = QwenOmniGenUndConfig.model_type
+        cfg_dict["architectures"] = [OmniGenUndForConditionalGeneration.__name__]
+        cfg = QwenOmniGenUndConfig(**cfg_dict)
+        print(f"[load] remapped config {src_model_type} -> {QwenOmniGenUndConfig.model_type}")
+
+    model = OmniGenUndForConditionalGeneration.from_pretrained(
+        args.model_path,
+        config=cfg,
+        torch_dtype=compute_dtype,
+        attn_implementation=attn_impl,
+        ignore_mismatched_sizes=True,
+    )
+    model.config.use_cache = False
+    model.config.dit_pe_mode = str(args.dit_pe_mode)
+
+    processor = AutoProcessor.from_pretrained(args.model_path, padding_side="left")
+    tokenizer = processor.tokenizer
+    tokenizer.num_metaqueries = args.num_metaqueries
+    tokenizer.system_prompt = args.system_prompt
+
+    boi_id = int(tokenizer.convert_tokens_to_ids("<begin_of_img>"))
+    eoi_id = int(tokenizer.convert_tokens_to_ids("<end_of_img>"))
+    model.gen_begin_token_idx = boi_id
+    model.gen_end_token_idx = eoi_id
+    model.num_metaqueries = args.num_metaqueries
+    model.get_anchor_token_idx(args.num_metaqueries, boi_id, eoi_id)
+    print(f"[load] boi={boi_id} eoi={eoi_id} metaqueries={args.num_metaqueries}")
+
+    print(f"[load] initialize_gen_vision_modules from {args.flux_path}")
+    model.initialize_gen_vision_modules(args.flux_path, from_config=False)
+
+    if args.gen_down_projector_path:
+        print(f"[load] gen_down_projector from {args.gen_down_projector_path}")
+        model.load_gen_down_projector(args.gen_down_projector_path)
+
+    print(
+        f"[load] inject MoE-FFN-TAG experts={args.dit_moe_num_routed_experts} "
+        f"topk={args.dit_moe_num_activated_experts} shared={args.dit_moe_ffn_shared_expert} "
+        f"last_n={args.dit_moe_ffn_last_n_double_layers}"
+    )
+    inject_flux_moe_ffn_tag(
+        model.dit,
+        num_routed_experts=args.dit_moe_num_routed_experts,
+        num_activated_experts=args.dit_moe_num_activated_experts,
+        gate_aux_alpha=args.dit_moe_gate_aux_alpha,
+        use_shared_expert=args.dit_moe_ffn_shared_expert,
+        moe_last_n_double_layers=args.dit_moe_ffn_last_n_double_layers,
+    )
+
+    if not args.skip_checkpoint:
+        ckpt_path = args.checkpoint
+        if os.path.isdir(ckpt_path):
+            ckpt_path = os.path.join(ckpt_path, "trainable_weights.bin")
+        print(f"[load] trainable weights from {ckpt_path}")
+        raw = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        dit_w, proj_w, emb_w = _split_tag_checkpoint(raw)
+
+        missing, unexpected = model.dit.load_state_dict(dit_w, strict=False)
+        print(
+            f"[load] DiT+TAG: dit_keys={len(dit_w)} missing={len(missing)} unexpected={len(unexpected)}"
+        )
+        if proj_w:
+            missing_p, unexpected_p = model.gen_down_projector.load_state_dict(proj_w, strict=False)
+            print(
+                f"[load] gen_down_projector: keys={len(proj_w)} missing={len(missing_p)} "
+                f"unexpected={len(unexpected_p)}"
+            )
+        if emb_w:
+            missing_e, unexpected_e = model.load_state_dict(emb_w, strict=False)
+            print(f"[load] stage1 embed: keys={len(emb_w)} missing={len(missing_e)} unexpected={len(unexpected_e)}")
+    else:
+        print("[load] skip_checkpoint")
+
+    model.tokenizer = tokenizer
+    device = _resolve_device(args.device)
+    model = _move_model_to_device(model, device, compute_dtype)
+    model.eval()
+    print(f"[load] ready on {device} dtype={compute_dtype}")
+    return model, processor, device
+
+
+@torch.no_grad()
+def tag_covt_sample_images(
+    model,
+    encoded: dict,
+    dit_batch: dict,
+    *,
+    height: int,
+    width: int,
+    guidance_scale: float,
+    num_inference_steps: int,
+    generator: Optional[torch.Generator],
+    device: torch.device,
+) -> List[Image.Image]:
+    dtype = next(model.parameters()).dtype
+    vae_scale_factor = 2 ** (len(model.gen_vision_tower.vae.config.block_out_channels) - 1)
+    multiple_of = vae_scale_factor * 2
+    height = height // multiple_of * multiple_of
+    width = width // multiple_of * multiple_of
+
+    meta_hidden = _encode_meta_hidden_omni(model, encoded)
+
+    num_channels_latents = model.dit.config.in_channels // 4
+    noise = torch.randn(
+        1,
+        num_channels_latents,
+        height // vae_scale_factor,
+        width // vae_scale_factor,
+        device=device,
+        dtype=dtype,
+        generator=generator,
+    )
+
+    latents, img_ids, l_noise = build_omni_dit_latent_input(
+        model,
+        dit_batch,
+        noisy_latents=noise,
+        B=1,
+        device=device,
+        dtype=dtype,
+    )
+    print(f"  DiT tokens: total={latents.shape[1]} target_noise={l_noise}  gen={height}x{width}")
+
+    scheduler = model.noise_scheduler
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+    if getattr(scheduler.config, "use_flow_sigmas", False):
+        sigmas = None
+    mu = calculate_shift(
+        l_noise,
+        scheduler.config.get("base_image_seq_len", 256),
+        scheduler.config.get("max_image_seq_len", 4096),
+        scheduler.config.get("base_shift", 0.5),
+        scheduler.config.get("max_shift", 1.15),
+    )
+    timesteps, num_inference_steps = retrieve_timesteps(
+        scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu
+    )
+    text_ids = torch.zeros(meta_hidden.shape[1], 3, device=device, dtype=dtype)
+    guidance = (
+        torch.full([1], guidance_scale, device=device, dtype=dtype)
+        if model.dit.config.guidance_embeds
+        else None
+    )
+
+    last_routing: List[torch.Tensor] = []
+    for t in tqdm(timesteps, desc="denoising"):
+        timestep = t.expand(latents.shape[0]).to(dtype)
+        joint_attention_kwargs = {
+            "num_target_tokens": l_noise,
+            "moe_aux_list": [],
+            "moe_routing_probs_list": last_routing,
+        }
+        prompt_embeds, pooled = model.gen_down_projector(meta_hidden, timestep / 1000)
+        noise_pred = model.dit(
+            hidden_states=latents,
+            timestep=timestep / 1000,
+            guidance=guidance,
+            pooled_projections=pooled,
+            encoder_hidden_states=prompt_embeds,
+            txt_ids=text_ids,
+            img_ids=img_ids,
+            joint_attention_kwargs=joint_attention_kwargs,
+            return_dict=False,
+        )[0]
+        noise_pred_tgt = noise_pred[:, :l_noise]
+        latents_tgt = latents[:, :l_noise]
+        latents_tgt = scheduler.step(noise_pred_tgt, t, latents_tgt, return_dict=False)[0]
+        if latents.shape[1] > l_noise:
+            latents = torch.cat([latents_tgt, latents[:, l_noise:]], dim=1)
+        else:
+            latents = latents_tgt
+
+    if last_routing:
+        g = torch.stack(last_routing, dim=0).mean(0)[0]
+        print("  routing: " + "  ".join(f"e{i}={p:.3f}" for i, p in enumerate(g.tolist())))
+
+    latents_out = model._unpack_latents(latents[:, :l_noise], height, width, vae_scale_factor)
+    return model.decode_latents(latents_out.to(model.gen_vision_tower.vae.dtype))
+
+
+def save_comparison(images: dict, out_path: str) -> None:
+    from PIL import ImageDraw, ImageFont
+
+    pils, labels = [], []
+    w_target = 512
+    for label, pil in images.items():
+        if pil is None:
+            continue
+        ratio = w_target / pil.width
+        pil = pil.resize((w_target, int(pil.height * ratio)), Image.LANCZOS)
+        pils.append(pil)
+        labels.append(label)
+    if not pils:
+        return
+    max_h = max(p.height for p in pils)
+    canvas = Image.new("RGB", (w_target * len(pils), max_h + 30), (240, 240, 240))
+    draw = ImageDraw.Draw(canvas)
+    font = ImageFont.load_default()
+    for i, (pil, label) in enumerate(zip(pils, labels)):
+        canvas.paste(pil, (i * w_target, 30))
+        draw.text((i * w_target + 5, 5), label, fill=(50, 50, 50), font=font)
+    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
+    canvas.save(out_path)
+    print(f"  saved → {out_path}")
+
+
+# TUNABLE: x2edit / multi_ref 默认多 sample 评测路径
+_X2EDIT_SRC = "/rczhang/rczhang/zhangrch/video_gen/gen_dataset/processed/x2edit/images/source"
+_X2EDIT_TGT = "/rczhang/rczhang/zhangrch/video_gen/gen_dataset/processed/x2edit/images/target"
+_MULTI_REF_ROOT = (
+    "/rczhang/rczhang/zhangrch/video_gen/gen_dataset/processed/"
+    "multi_ref_edit_2_extracted/data/samples/20260421_165536_8way/shard_00_of_08"
+)
+# TUNABLE: 与 configs/moe_ffn_tag.yaml data 一致
+_DEFAULT_ARROW_DIR = (
+    "/rczhang/rczhang/zhangrch/video_gen/gen_dataset/processed/"
+    "instruction_only_quality_500k/x2edit/arrow"
+)
+_DEFAULT_IMAGE_ROOT = None
+
+
+def _x2edit_case(idx: int, task_type: str, prompt: str) -> Dict[str, Any]:
+    sid = f"{idx:08d}"
+    return {
+        "name": f"{sid}_{task_type}",
+        "task_type": task_type,
+        "prompt": prompt,
+        "source": f"{_X2EDIT_SRC}/x2edit_train_{sid}.jpg",
+        "target": f"{_X2EDIT_TGT}/x2edit_train_{sid}.jpg",
+    }
+
+
+# TUNABLE: 默认固定 case（5 类 task + 更多 index + source+ref）
+DEFAULT_X2EDIT_MULTI_CASES = [
+    _x2edit_case(0, "camera_movement", "Rotate the bowl 45 degrees counterclockwise."),
+    _x2edit_case(1, "action_change", "a judge's gavel on a wooden table"),
+    _x2edit_case(2, "reasoning_edit", "Can you show me an image with more details about the ring's band?"),
+    _x2edit_case(5, "portrait_editing", "an elegant ring with a round diamond center"),
+    _x2edit_case(16, "text_change", "Can you remove the text 'Fashionthings' from the bracelet?"),
+    _x2edit_case(50, "camera_movement", "a man is sitting in a car on a stage"),
+    _x2edit_case(100, "action_change", "two stuffed eggplant halves on a plate with vegetables"),
+    _x2edit_case(200, "portrait_editing", "a pair of french style sofas with white upholstered cushions"),
+    _x2edit_case(500, "text_change", "Remove the text 'BOSCH' from the stove hood."),
+    _x2edit_case(1000, "action_change", "a man wearing a t - shirt with a monkey on it"),
+]
+
+DEFAULT_MULTI_REF_CASES = [
+    {
+        "name": "multi_ref_subject_deletion",
+        "task_type": "subject_deletion",
+        "prompt": "Remove the Person in Image 1.",
+        "source": f"{_MULTI_REF_ROOT}/s00_267878_0d7bad31_original.jpg",
+        "reference": f"{_MULTI_REF_ROOT}/s00_267878_0d7bad31_reference_Person.jpg",
+        "target": f"{_MULTI_REF_ROOT}/s00_267878_0d7bad31_inpainted.jpg",
+    },
+    {
+        "name": "multi_ref_subject_addition",
+        "task_type": "subject_addition",
+        "prompt": "Add the Person in Image 1.",
+        "source": f"{_MULTI_REF_ROOT}/s00_267878_0d7bad31_inpainted.jpg",
+        "reference": f"{_MULTI_REF_ROOT}/s00_267878_0d7bad31_reference_Person.jpg",
+        "target": f"{_MULTI_REF_ROOT}/s00_267878_0d7bad31_original.jpg",
+    },
+]
+
+DEFAULT_ALL_MULTI_CASES = DEFAULT_X2EDIT_MULTI_CASES + DEFAULT_MULTI_REF_CASES
+
+
+def _resolve_image_path(path: str, image_root: Optional[str]) -> str:
+    path = str(path or "").strip()
+    if not path:
+        return ""
+    if os.path.isabs(path) or not image_root:
+        return path
+    return os.path.join(image_root, path)
+
+
+def _row_to_case(row: Dict[str, Any], idx: int, image_root: Optional[str]) -> Optional[Dict[str, Any]]:
+    target = _resolve_image_path(str(row.get("target_image") or ""), image_root)
+    source = _resolve_image_path(str(row.get("source_image") or ""), image_root)
+    if not target or not os.path.isfile(target):
+        return None
+
+    refs_raw = row.get("reference_images") or []
+    if isinstance(refs_raw, str):
+        refs_raw = [refs_raw] if refs_raw else []
+    ref_paths = [
+        p for p in (_resolve_image_path(str(r), image_root) for r in refs_raw) if p and os.path.isfile(p)
+    ]
+
+    has_source = bool(source) and os.path.isfile(source)
+    has_ref = bool(ref_paths)
+    if not has_source and not has_ref:
+        return None
+
+    task_type = str(row.get("task_type") or "other").strip() or "other"
+    if task_type not in TASK_TYPE_TO_ID:
+        task_type = "other"
+    instruction = str(row.get("instruction") or "").strip()
+    if not instruction:
+        return None
+
+    sample_id = str(row.get("sample_id") or row.get("id") or f"idx{idx:08d}")
+    case: Dict[str, Any] = {
+        "name": f"arrow_{idx:08d}_{task_type}",
+        "task_type": task_type,
+        "prompt": instruction,
+        "target": target,
+    }
+    if has_source:
+        case["source"] = source
+    if has_ref:
+        case["reference"] = ref_paths[0] if len(ref_paths) == 1 else ",".join(ref_paths)
+    return case
+
+
+def load_cases_from_arrow_indices(
+    arrow_dir: str,
+    image_root: Optional[str],
+    indices: Sequence[int],
+) -> List[Dict[str, Any]]:
+    from datasets import load_from_disk
+
+    ds = load_from_disk(arrow_dir)
+    cases: List[Dict[str, Any]] = []
+    for idx in indices:
+        if idx < 0 or idx >= len(ds):
+            print(f"[warn] arrow index {idx} out of range (len={len(ds)}), skip")
+            continue
+        case = _row_to_case(ds[int(idx)], int(idx), image_root)
+        if case is None:
+            print(f"[warn] arrow index {idx} invalid (missing image/instruction), skip")
+            continue
+        cases.append(case)
+    return cases
+
+
+def sample_cases_from_arrow(
+    arrow_dir: str,
+    image_root: Optional[str],
+    num_samples: int,
+    seed: int,
+    *,
+    source_only: bool = True,
+) -> List[Dict[str, Any]]:
+    from datasets import load_from_disk
+
+    ds = load_from_disk(arrow_dir)
+    rng = np.random.default_rng(seed)
+    order = rng.permutation(len(ds))
+    cases: List[Dict[str, Any]] = []
+
+    for idx in order:
+        row = ds[int(idx)]
+        source = _resolve_image_path(str(row.get("source_image") or ""), image_root)
+        has_source = bool(source) and os.path.isfile(source)
+        refs_raw = row.get("reference_images") or []
+        if isinstance(refs_raw, str):
+            refs_raw = [refs_raw] if refs_raw else []
+        has_ref = bool(refs_raw)
+        if source_only and not has_source:
+            continue
+        case = _row_to_case(row, int(idx), image_root)
+        if case is None:
+            continue
+        cases.append(case)
+        if len(cases) >= num_samples:
+            break
+
+    if len(cases) < num_samples:
+        raise RuntimeError(
+            f"Only found {len(cases)} valid arrow samples (requested {num_samples}). "
+            "Check arrow_dir / image_root or relax source_only."
+        )
+    return cases
+
+
+def resolve_multi_cases(args) -> List[Dict[str, Any]]:
+    cases = list(DEFAULT_ALL_MULTI_CASES)
+    if args.arrow_indices:
+        indices = [int(x.strip()) for x in args.arrow_indices.split(",") if x.strip()]
+        cases.extend(load_cases_from_arrow_indices(args.arrow_dir, args.image_root, indices))
+    if args.multi_from_arrow > 0:
+        cases.extend(
+            sample_cases_from_arrow(
+                args.arrow_dir,
+                args.image_root,
+                args.multi_from_arrow,
+                args.seed,
+                source_only=not args.multi_include_ref,
+            )
+        )
+    if args.multi_limit > 0:
+        cases = cases[: args.multi_limit]
+    return cases
+
+
+def run_single_case(
+    model,
+    processor,
+    device,
+    dtype,
+    args,
+    *,
+    case_name: str,
+    task_type: str,
+    prompt: str,
+    source_image: Optional[str],
+    target_image: Optional[str],
+    reference_image: Optional[str] = None,
+    case_output_dir: Path,
+    seed: int,
+) -> None:
+    anchor_model_id = ast.literal_eval(args.anchor_model_id)
+    case_output_dir.mkdir(parents=True, exist_ok=True)
+
+    source_pil = Image.open(source_image).convert("RGB") if source_image else None
+    ref_paths = [p.strip() for p in reference_image.split(",") if p.strip()] if reference_image else []
+    if len(ref_paths) > args.max_reference_images:
+        print(
+            f"[warn] truncating references from {len(ref_paths)} "
+            f"to {args.max_reference_images}"
+        )
+        ref_paths = ref_paths[: args.max_reference_images]
+    ref_pils = [Image.open(p).convert("RGB") for p in ref_paths]
+    gt_pil = Image.open(target_image).convert("RGB") if target_image else None
+
+    has_source = source_pil is not None
+    has_ref = len(ref_pils) > 0
+    cond_id = 2 if has_source and has_ref else 1 if has_source else 0 if has_ref else -1
+    task_type_id = TASK_TYPE_TO_ID.get(task_type, TASK_TYPE_TO_ID["other"])
+    print(
+        f"\n=== case={case_name} ===\n"
+        f"task_type={task_type}({task_type_id})  "
+        f"condition={CONDITION_NAMES.get(cond_id, 't2i')}({cond_id})  covt_mode={args.covt_mode}\n"
+        f"prompt: {prompt}"
+    )
+
+    encoded = build_covt_encoded_input(
+        processor,
+        prompt=prompt,
+        source_pil=source_pil,
+        ref_pils=ref_pils,
+        num_metaqueries=args.num_metaqueries,
+        image_min_pixels=args.image_min_pixels,
+        image_max_pixels=args.image_max_pixels,
+        image_resized_width=args.image_resized_width,
+        image_resized_height=args.image_resized_height,
+        anchor_model_id=anchor_model_id,
+        covt_mode=args.covt_mode,
+        seed=seed,
+        device=device,
+    )
+
+    dit_batch, gen_h, gen_w = prepare_dit_gen_batch(
+        source_pil=source_pil,
+        ref_pils=ref_pils,
+        image_size=args.image_size,
+        max_reference_images=args.max_reference_images,
+        device=device,
+        dtype=dtype,
+    )
+    if not has_source:
+        gen_h, gen_w = args.height, args.width
+
+    gen = torch.Generator(device=device).manual_seed(seed)
+    images = tag_covt_sample_images(
+        model,
+        encoded,
+        dit_batch,
+        height=gen_h,
+        width=gen_w,
+        guidance_scale=args.guidance_scale,
+        num_inference_steps=args.steps,
+        generator=gen,
+        device=device,
+    )
+
+    out_pil = images[0] if images else None
+    if out_pil is not None:
+        out_path = case_output_dir / "generated.png"
+        out_pil.save(out_path)
+        print(f"  saved → {out_path}")
+
+    save_comparison(
+        {"source": source_pil, "generated": out_pil, "target": gt_pil},
+        str(case_output_dir / "comparison.png"),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="TAG-MoE quality500k-v1 inference")
+    # TUNABLE: 与 configs/moe_ffn_tag.yaml model 段一致
+    parser.add_argument(
+        "--model_path",
+        default="/rczhang/rczhang/zhangrch/video_gen/gen_checkpoint/stage3_dino_depth_siglip_pidinet_lora16_zero2_trainable_embed_lm_head_downpro_web18m_gen_llmlora_v2/lora_merged/lora_stage3_und_gen_merged_30k",
+    )
+    parser.add_argument(
+        "--flux_path",
+        default="/rczhang/rczhang/zhangrch/video_gen/pretrined_models/FLUX.1-dev",
+    )
+    parser.add_argument(
+        "--gen_down_projector_path",
+        default="/rczhang/rczhang/zhangrch/video_gen/gen_checkpoint/stage3_dino_depth_siglip_pidinet_lora16_zero2_trainable_embed_lm_head_downpro_web18m_gen_llmlora_v2/lora_vision_test/lora_stage3_und_gen/embeddings_connector/checkpoint-30000.bin",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        default="/rczhang/rczhang/zhangrch/video_gen/gen_checkpoint/moe_ffn_tag_covt_quality500k_v1_moe_fast/checkpoint-1000/trainable_weights.bin",
+    )
+    parser.add_argument(
+        "--anchor_model_id",
+        default="['dino', 'depth', 'siglip', 'pidinet']",
+        help="CoVT anchor 列表，与实验本地 data.py 一致",
+    )
+    parser.add_argument("--prompt", default="Rotate the bowl 45 degrees counterclockwise.")
+    parser.add_argument("--source_image", default=None)
+    parser.add_argument("--reference_image", default=None)
+    parser.add_argument("--target_image", default=None)
+    parser.add_argument("--task_type", default="camera_movement", choices=TASK_TYPES)
+    parser.add_argument("--output_dir", default="./tag_covt_infer_out")
+    parser.add_argument("--num_metaqueries", type=int, default=256)
+    parser.add_argument("--system_prompt", default="You are a helpful assistant.")
+    parser.add_argument("--disable_flash_attn2", action="store_true")
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--dit_moe_num_routed_experts", type=int, default=4)
+    parser.add_argument("--dit_moe_num_activated_experts", type=int, default=1)
+    parser.add_argument("--dit_moe_gate_aux_alpha", type=float, default=0.01)
+    parser.add_argument("--dit_moe_ffn_shared_expert", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--dit_moe_ffn_last_n_double_layers", type=int, default=10)
+    parser.add_argument("--dit_pe_mode", default="d")
+    parser.add_argument("--image_size", type=int, default=1024)
+    parser.add_argument("--image_min_pixels", type=int, default=200704)
+    parser.add_argument("--image_max_pixels", type=int, default=802816)
+    parser.add_argument("--image_resized_width", type=int, default=448)
+    parser.add_argument("--image_resized_height", type=int, default=448)
+    parser.add_argument("--max_reference_images", type=int, default=4)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=30)
+    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--covt_mode",
+        choices=("full", "random", "none"),
+        default="full",
+        help="assistant CoVT 注入：full=全 anchor（I2I 推荐）；random=复现训练随机；none=仅 gen placeholder",
+    )
+    parser.add_argument("--skip_checkpoint", action="store_true")
+    parser.add_argument(
+        "--multi_x2edit",
+        action="store_true",
+        help="一次加载模型，跑默认多样本（10 x2edit + 2 source+ref，可用 --multi_limit 截断）",
+    )
+    # TUNABLE: 多样本数量与 arrow 采样
+    parser.add_argument(
+        "--multi_limit",
+        type=int,
+        default=0,
+        help=">0 时只跑前 N 条 multi case（默认 0=全部）",
+    )
+    parser.add_argument(
+        "--multi_from_arrow",
+        type=int,
+        default=0,
+        help=">0 时从 arrow 随机再抽 N 条（需 --arrow_dir）",
+    )
+    parser.add_argument(
+        "--multi_include_ref",
+        action="store_true",
+        help="--multi_from_arrow 时允许 source+ref 样本（默认仅 source-only）",
+    )
+    parser.add_argument("--arrow_dir", default=_DEFAULT_ARROW_DIR)
+    parser.add_argument("--image_root", default=_DEFAULT_IMAGE_ROOT)
+    parser.add_argument(
+        "--arrow_indices",
+        default="",
+        help="逗号分隔 arrow index，追加到默认 case 列表（如 300,800,1200）",
+    )
+    args = parser.parse_args()
+
+    set_seed(args.seed)
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+    model, processor, device = load_omni_tag_model(args)
+    dtype = next(model.parameters()).dtype
+
+    if args.multi_x2edit:
+        multi_cases = resolve_multi_cases(args)
+        print(f"[multi] running {len(multi_cases)} cases → {args.output_dir}")
+        for i, case in enumerate(multi_cases):
+            run_single_case(
+                model,
+                processor,
+                device,
+                dtype,
+                args,
+                case_name=case["name"],
+                task_type=case["task_type"],
+                prompt=case["prompt"],
+                source_image=case.get("source"),
+                target_image=case.get("target"),
+                reference_image=case.get("reference"),
+                case_output_dir=Path(args.output_dir) / case["name"],
+                seed=args.seed + i,
+            )
+        print(f"\n[multi] all done → {args.output_dir}")
+        return
+
+    if not args.source_image and not args.reference_image:
+        raise SystemExit("请指定 --source_image/--reference_image，或使用 --multi_x2edit")
+
+    run_single_case(
+        model,
+        processor,
+        device,
+        dtype,
+        args,
+        case_name="single",
+        task_type=args.task_type,
+        prompt=args.prompt,
+        source_image=args.source_image,
+        target_image=args.target_image,
+        reference_image=args.reference_image,
+        case_output_dir=Path(args.output_dir),
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    main()

code/umm/inference/run_infer.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+STAGE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${STAGE_DIR}/../../.." && pwd)"
+cd "${REPO_ROOT}"
+
+export PYTHONPATH="${REPO_ROOT}/hc_test/umm/runtime:${REPO_ROOT}:${PYTHONPATH:-}"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
+export HF_HOME="${HF_HOME:-/rczhang/rczhang/zhangrch/video_gen/omnigen/cache/huggingface}"
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY
+
+PYTHON_BIN="${PYTHON_BIN:-/rczhang/rczhang/conda_packs/unvideo/bin/python}"
+MODEL_PATH="${MODEL_PATH:-/rczhang/rczhang/zhangrch/video_gen/gen_checkpoint/stage3_dino_depth_siglip_pidinet_lora16_zero2_trainable_embed_lm_head_downpro_web18m_gen_llmlora_v2/lora_merged/lora_stage3_und_gen_merged_30k}"
+FLUX_PATH="${FLUX_PATH:-/rczhang/rczhang/zhangrch/video_gen/pretrined_models/FLUX.1-dev}"
+PROJECTOR_PATH="${PROJECTOR_PATH:-/rczhang/rczhang/zhangrch/video_gen/gen_checkpoint/stage3_dino_depth_siglip_pidinet_lora16_zero2_trainable_embed_lm_head_downpro_web18m_gen_llmlora_v2/lora_vision_test/lora_stage3_und_gen/embeddings_connector/checkpoint-30000.bin}"
+CHECKPOINT="${CHECKPOINT:-/rczhang/rczhang/zhangrch/video_gen/gen_checkpoint/moe_ffn_tag_covt_quality500k_v1_moe_fast/checkpoint-1000/trainable_weights.bin}"
+SOURCE_IMAGE="${SOURCE_IMAGE:-}"
+REFERENCE_IMAGE="${REFERENCE_IMAGE:-}"
+TARGET_IMAGE="${TARGET_IMAGE:-}"
+PROMPT="${PROMPT:-Rotate the bowl 45 degrees counterclockwise.}"
+TASK_TYPE="${TASK_TYPE:-camera_movement}"
+OUTPUT_DIR="${OUTPUT_DIR:-${STAGE_DIR}/outputs}"
+IMAGE_SIZE="${IMAGE_SIZE:-512}"
+STEPS="${STEPS:-30}"
+SEED="${SEED:-42}"
+COVT_MODE="${COVT_MODE:-full}"
+
+CMD=(
+  "${PYTHON_BIN}" "${STAGE_DIR}/infer.py"
+  --model_path "${MODEL_PATH}"
+  --flux_path "${FLUX_PATH}"
+  --gen_down_projector_path "${PROJECTOR_PATH}"
+  --checkpoint "${CHECKPOINT}"
+  --prompt "${PROMPT}"
+  --task_type "${TASK_TYPE}"
+  --output_dir "${OUTPUT_DIR}"
+  --covt_mode "${COVT_MODE}"
+  --image_size "${IMAGE_SIZE}"
+  --height "${IMAGE_SIZE}"
+  --width "${IMAGE_SIZE}"
+  --steps "${STEPS}"
+  --seed "${SEED}"
+)
+
+if [[ -n "${SOURCE_IMAGE}" ]]; then
+  CMD+=(--source_image "${SOURCE_IMAGE}")
+fi
+if [[ -n "${REFERENCE_IMAGE}" ]]; then
+  CMD+=(--reference_image "${REFERENCE_IMAGE}")
+fi
+if [[ -n "${TARGET_IMAGE}" ]]; then
+  CMD+=(--target_image "${TARGET_IMAGE}")
+fi
+if [[ "${MULTI_X2EDIT:-0}" == "1" ]]; then
+  CMD+=(--multi_x2edit)
+fi
+CMD+=("$@")
+
+if [[ "${DRY_RUN:-0}" == "1" ]]; then
+  printf '%q ' "${CMD[@]}"
+  printf '\n'
+  exit 0
+fi
+
+if [[ -z "${SOURCE_IMAGE}" && -z "${REFERENCE_IMAGE}" && "${MULTI_X2EDIT:-0}" != "1" ]]; then
+  echo "请通过 SOURCE_IMAGE=/path 或 REFERENCE_IMAGE=/path 指定输入；也可设置 MULTI_X2EDIT=1。" >&2
+  exit 2
+fi
+
+exec "${CMD[@]}"

code/umm/runtime/README.md

@@ -0,0 +1,50 @@
+# Runtime：UMM 本地共享源码
+
+这里是各阶段真正导入的共享源码副本，不是跳转层。
+
+```text
+runtime/
+├── models/blip3o/             Qwen/CoVT/FLUX、dataset、trainer、MoE 实现
+└── deepspeed_configs/zero2.json
+```
+
+所有 `run_*.sh` 都把本目录放在 `PYTHONPATH` 第一位：
+
+```bash
+export PYTHONPATH=".../hc_test/umm/runtime:.../unigen_ar_v1:${PYTHONPATH:-}"
+```
+
+因此 `import models.blip3o...` 实际命中本目录。可以用下面的命令确认：
+
+```bash
+/rczhang/rczhang/conda_packs/unvideo/bin/python \
+  hc_test/umm/runtime/check_imports.py
+```
+
+输出的三个模型路径都应位于 `hc_test/umm/runtime/models/blip3o/`。
+
+## 哪些内容仍在目录外
+
+只保留大文件的路径引用，不重复复制：
+
+- Qwen、FLUX、DINO、Depth、PiDiNet、SigLIP 等预训练权重；
+- JSON/JSONL、Arrow 与图片数据；
+- 已有 dense、LoRA、projector、MoE checkpoint；
+- Hugging Face、datasets、W&B 与临时缓存。
+
+本目录中没有 `.pth`、`.bin`、`.safetensors` 或 `.ckpt` 权重文件。修改模型、dataset 或 trainer 时，应修改这里的副本；阶段入口的专属逻辑则修改相应 `stage*/train.py`。
+
+## 本地修正
+
+复制后做了让本工作区自洽并保证 Stage-B 训练正确性的修正：
+
+- Stage 1/2/3 使用显式模型类导入，避免包级通配导入加载无关旧模块；
+- Stage 3 FLUX 模型显式导入 Qwen3-VL 的两个输出类型；
+- 各阶段增加本地 YAML 配置解析；
+- Stage 4 与 inference 的 experiment-local import 改为 `hc_test.umm` 路径。
+- Stage 2 加入 frozen SAM ViT-H 自动伪 mask、8-query decoder prompt，以及论文定义的 Hungarian Dice+Focal segmentation loss；
+- CoVT curriculum 统一使用 Trainer optimizer `global_step`，避免 worker/梯度累积改变阶段边界；
+- 修正真实 anchor token 注册、Stage-1 query 行加载、专家头确定性初始化、非有限 loss 硬检查；
+- checkpoint 单 rank 写入并显式恢复 non-LoRA trainables 与专家 sidecar，保证断点续训不重置视觉能力头。
+
+这些修改都只发生在 `hc_test/umm` 内，没有改写原模型源码。

code/umm/runtime/check_imports.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+"""Confirm that UMM resolves shared model code from this local runtime copy."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from models.blip3o.model.language_model import (
+    blip3o_qwen_stage1,
+    covt_qwen_stage2_van,
+    moe_qwen_stage3_flux,
+)
+
+
+def main() -> None:
+    runtime_root = Path(__file__).resolve().parent
+    modules = {
+        "stage1_model": blip3o_qwen_stage1,
+        "stage2_model": covt_qwen_stage2_van,
+        "stage3_model": moe_qwen_stage3_flux,
+    }
+    for name, module in modules.items():
+        path = Path(module.__file__).resolve()
+        print(f"{name}={path}")
+        if runtime_root not in path.parents:
+            raise RuntimeError(f"{name} did not resolve from UMM runtime: {path}")
+    print("status=ok")
+
+
+if __name__ == "__main__":
+    main()

code/umm/runtime/deepspeed_configs/zero2.json

@@ -0,0 +1,42 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "betas": "auto",
+            "eps": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 5e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 5e8,
+        "contiguous_gradients": true,
+        "ignore_unused_parameters": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 100,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

code/umm/runtime/models/blip3o/constants.py

@@ -0,0 +1,82 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+IGNORE_INDEX = -100
+
+# DEFAULT_IMAGE_TOKEN = "<image>"
+# DEFAULT_IM_START_TOKEN = "[IMG]"
+# DEFAULT_IM_END_TOKEN = "[/IMG]"
+
+
+# IMAGE_TOKEN_IDX = 151667
+T2I_TOKEN_IDX = 151666
+I2I_TOKEN_IDX = 151667
+
+
+UND_IMAGE_TOKEN = "<|image_pad|>"
+UND_IMAGE_TOKEN_IDX = 151655
+
+
+DEFAULT_IM_START_TOKEN = "<|im_start|>"
+DEFAULT_IM_END_TOKEN = "<|im_end|>"
+DEFAULT_IMAGE_TOKEN = "<|image_pad|>"
+DEFAULT_VIDEO_TOKEN = "<|video_pad|>"
+LLAVA_IMAGE_TOKEN = "<image>"
+LLAVA_VIDEO_TOKEN = "<video>"
+VISION_START_TOKEN = "<|vision_start|>"
+VISION_END_TOKEN = "<|vision_end|>"
+
+ANCHOR_START_TOKEN = "<|anchor_start|>"
+ANCHOR_END_TOKEN = "<|anchor_end|>"
+SAM_PAD_TOKEN = "<|sam_pad|>"
+DINO_PAD_TOKEN = "<|dino_pad|>"
+DEPTH_PAD_TOKEN = "<|depth_pad|>"
+SD_PAD_TOKEN = "<|sd_pad|>"
+INTERN_PAD_TOKEN = "<|intern_pad|>"
+PIDINET_PAD_TOKEN = "<|pidinet_pad|>"
+SIGLIP_PAD_TOKEN = "<|siglip_pad|>"
+METACLIP_PAD_TOKEN = "<|metaclip_pad|>"
+
+# Stage-B CoVT semantic-token counts.  Keep these counts in one place because
+# the dataset, Stage-2 geometry projector, and Stage-3 I2I suffix must agree.
+COVT_SAM_TOKEN_COUNT = 8
+COVT_DINO_TOKEN_COUNT = 4
+COVT_DEPTH_TOKEN_COUNT = 4
+COVT_PIDINET_TOKEN_COUNT = 4
+COVT_SIGLIP_TOKEN_COUNT = 4
+
+
+# --- TUNABLE BLOCK: Stage-3 gen anchor strings (one tokenizer id each after add_tokens; keep in sync with train_stage3_covt_gen / covt_data_var) ---
+VIS_DINO_ANCHOR_TOKENS = [f"<|vis_dino_pad_{i}|>" for i in range(4)]
+VIS_DEPTH_ANCHOR_TOKENS = [f"<|vis_depth_pad_{i}|>" for i in range(4)]
+VIS_PIDINET_ANCHOR_TOKENS = [f"<|vis_pidinet_pad_{i}|>" for i in range(4)]
+VIS_SIGLIP_ANCHOR_TOKENS = [f"<|vis_siglip_pad_{i}|>" for i in range(4)]
+
+# Optional I2I: append CoVT anchor pads after the generation placeholder — **same layout as stage-2** covt_data:
+# ``<|anchor_start|><|sam_pad|>*8<|anchor_end|>`` + ``<|anchor_start|><|dino_pad|>*4<|anchor_end|>`` + ``<|anchor_start|><|depth_pad|>*4<|anchor_end|>``
+# (counts from ``covt_data.get_token_num`` for sam / dino / depth). Do **not** spell ``<|*_pad|>`` in the human
+# turn, or those strings tokenize to duplicate ids and break parity with stage 2.
+def _i2i_covt_anchor_blocks_stage2_sam_dino_depth() -> str:
+    return (
+        ANCHOR_START_TOKEN
+        + SAM_PAD_TOKEN * COVT_SAM_TOKEN_COUNT
+        + ANCHOR_END_TOKEN
+        + ANCHOR_START_TOKEN
+        + DINO_PAD_TOKEN * COVT_DINO_TOKEN_COUNT
+        + ANCHOR_END_TOKEN
+        + ANCHOR_START_TOKEN
+        + DEPTH_PAD_TOKEN * COVT_DEPTH_TOKEN_COUNT
+        + ANCHOR_END_TOKEN
+    )
+
+
+I2I_COVT_ANCHOR_HUMAN_EXTRA = (
+    "\n\n[Format requirement] After the generation block (<begin_of_img>…<end_of_img>), the assistant reply must "
+    "include the same CoVT anchor pad blocks as stage-2 training (SAM / DINO / depth counts between "
+    "anchor start and end markers). Do not omit or reorder them."
+)
+I2I_COVT_ANCHOR_GPT_SUFFIX = "\n" + _i2i_covt_anchor_blocks_stage2_sam_dino_depth() + "\n"
+
+SYSTEM_MESSAGE = "You are a helpful assistant."

code/umm/runtime/models/blip3o/conversation.py

@@ -0,0 +1,479 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+import base64
+from io import BytesIO
+from PIL import Image
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+    CHATML = auto()
+    QWEN = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0]
+            if "mmtag" in self.version:
+                init_msg = init_msg.replace("<image>", "").strip()
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            elif not init_msg.startswith("<image>"):
+                init_msg = init_msg.replace("<image>", "").strip()
+                messages[0] = (init_role, "<image>\n" + init_msg)
+            else:
+                messages[0] = (init_role, init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = "" if self.system == "" else self.system + self.sep + "\n"
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, images, _ = message
+                        message = "<image>" * len(images) + message
+                    ret += role + "\n" + message + self.sep + "\n"
+                else:
+                    ret += role + "\n"
+            return ret
+
+        elif self.sep_style == SeparatorStyle.LLAMA_3:
+            if self.tokenizer is None:
+                raise ValueError("Llama 3 tokenizer is not available. Make sure you have the necessary permissions.")
+            chat_template_messages = [{"role": "system", "content": self.system}]
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, images = message
+                        message = "<image>" * len(images) + message
+                    chat_template_messages.append({"role": role, "content": message})
+
+            # print(chat_template_messages)
+            return self.tokenizer.apply_chat_template(chat_template_messages, tokenize=False, add_generation_prompt=True)
+            # ret = "" if self.system == "" else self.system + self.sep + "\n"
+            # for role, message in messages:
+            #     if message:
+            #         if type(message) is tuple:
+            #             message, images = message
+            #             message = "<image>" * len(images) + message
+            #         ret += role + "\n" + message + self.sep + "\n"
+            #     else:
+            #         ret += role + "\n"
+            # return ret
+
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+
+        elif self.sep_style == SeparatorStyle.GEMMA:
+            ret = ""
+            for i, (role, message) in enumerate(messages):
+                assert role == self.roles[i % 2], "Conversation should alternate user/assistant/user/assistant/..."
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n" if len(msg) > 0 else msg
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0:
+                        message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def process_image(self, image, image_process_mode, return_pil=False, image_format='PNG', max_len=1344, min_len=672):
+        if image_process_mode == "Pad":
+            def expand2square(pil_img, background_color=(122, 116, 104)):
+                width, height = pil_img.size
+                if width == height:
+                    return pil_img
+                elif width > height:
+                    result = Image.new(pil_img.mode, (width, width), background_color)
+                    result.paste(pil_img, (0, (width - height) // 2))
+                    return result
+                else:
+                    result = Image.new(pil_img.mode, (height, height), background_color)
+                    result.paste(pil_img, ((height - width) // 2, 0))
+                    return result
+
+            image = expand2square(image)
+        elif image_process_mode in ["Default", "Crop"]:
+            pass
+        elif image_process_mode == "Resize":
+            image = image.resize((336, 336))
+        else:
+            raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+        if max(image.size) > max_len:
+            max_hw, min_hw = max(image.size), min(image.size)
+            aspect_ratio = max_hw / min_hw
+            shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+            longest_edge = int(shortest_edge * aspect_ratio)
+            W, H = image.size
+            if H > W:
+                H, W = longest_edge, shortest_edge
+            else:
+                H, W = shortest_edge, longest_edge
+            image = image.resize((W, H))
+        if return_pil:
+            return image
+        else:
+            buffered = BytesIO()
+            image.save(buffered, format=image_format)
+            img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+            return img_b64_str
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg, image, image_process_mode = msg
+                    image = self.process_image(image, image_process_mode, return_pil=return_pil)
+                    images.append(image)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg, image, image_process_mode = msg
+                    img_b64_str = self.process_image(
+                        image, "Default", return_pil=False,
+                        image_format='JPEG')
+                    img_str = f'<img src="data:image/jpeg;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace('<image>', '').strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+         "Renewable energy sources are those that can be replenished naturally in a relatively "
+         "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+         "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+         "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+         "renewable and non-renewable energy sources:\n"
+         "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+         "energy sources are finite and will eventually run out.\n"
+         "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+         "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+         "and other negative effects.\n"
+         "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+         "have lower operational costs than non-renewable sources.\n"
+         "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+         "locations than non-renewable sources.\n"
+         "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+         "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+         "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+         "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+
+conv_blip3o_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+           "You are able to understand the visual content that the user provides, "
+           "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_blip3o_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_blip3o_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_blip3o_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_blip3o_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_blip3o_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+conv_mistral_instruct = Conversation(
+    system="",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="",
+    sep2="</s>",
+)
+
+conv_chatml_direct = Conversation(
+    system="""<|im_start|>system
+Answer the questions.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llama3 = Conversation(
+    system="""<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nA chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.""",
+    roles=("<|start_header_id|>user<|end_header_id|>\n\n", "<|start_header_id|>assistant<|end_header_id|>\n\n"),
+    version="llama3",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|eot_id|>",
+)
+
+conv_qwen = Conversation(
+    system="""<|im_start|>system
+You are a helpful assistant.""",
+    roles=("<|im_start|>user", "<|im_start|>assistant"),
+    version="qwen",
+    messages=[],
+    offset=0,
+    sep_style=SeparatorStyle.CHATML,
+    sep="<|im_end|>",
+)
+
+
+default_conversation = conv_llama3
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+    "mistral_instruct": conv_mistral_instruct,
+    "chatml_direct": conv_chatml_direct,
+    "mistral_direct": conv_chatml_direct,
+
+    "plain": conv_blip3o_plain,
+    "v0_plain": conv_blip3o_plain,
+    "blip3o_v0": conv_blip3o_v0,
+    "v0_mmtag": conv_blip3o_v0_mmtag,
+    "blip3o_v1": conv_blip3o_v1,
+    "v1_mmtag": conv_blip3o_v1_mmtag,
+    "blip3o_llama_2": conv_blip3o_llama_2,
+    "llama3": conv_llama3,
+    "qwen": conv_qwen, 
+
+    "mpt": conv_mpt,
+}
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

code/umm/runtime/models/blip3o/data/covt_data.py

@@ -0,0 +1,827 @@
+import copy
+import os
+from dataclasses import dataclass, field
+from typing import Dict
+import torch
+import transformers
+import ujson as json
+from torch.utils.data import Dataset
+from qwen_vl_utils import process_vision_info
+from PIL import Image
+from transformers import AutoImageProcessor
+import re
+import numpy as np
+import cv2
+from torchvision import transforms
+import random
+
+from qwen_vl_utils import process_vision_info
+from models.blip3o.model.anchors.segment_anything import build_sam_vit_h, sam_model_registry, SamPredictor
+from models.blip3o.model.anchors.DepthAnything.depth_anything_v2.dpt import DepthAnythingV2
+from diffusers import AutoencoderKL
+from transformers import AutoModel, CLIPImageProcessor
+
+from models.blip3o.params import DataArguments
+from models.blip3o.constants import *
+
+
+def rank0_print(*args):
+    if os.environ.get("LOCAL_RANK", "0") in ("0", "-1"):
+        print(*args)
+
+
+def truncate_sequence(input_ids, labels, max_length, eos_token_id):
+    if input_ids.size(0) > max_length:
+        input_ids = input_ids[:max_length-1]
+        labels = labels[:max_length-1]
+
+    if eos_token_id is not None:
+        input_ids = torch.cat([input_ids, torch.tensor([eos_token_id])])
+        labels = torch.cat([labels, torch.tensor([eos_token_id])])
+
+    return input_ids, labels
+
+def _mm_type_ids_for_processor_batch(processor, inputs):
+    """Qwen3-VL M-RoPE expects mm_token_type_ids aligned with input_ids (0=text, 1=image, 2=video)."""
+    if "mm_token_type_ids" in inputs:
+        mm = inputs["mm_token_type_ids"]
+        if mm.dim() == 2:
+            mm = mm.squeeze(0)
+        return mm.to(torch.long)
+    if hasattr(processor, "create_mm_token_type_ids"):
+        row = inputs["input_ids"].tolist()
+        mm_list = processor.create_mm_token_type_ids(row)
+        return torch.tensor(mm_list[0], dtype=torch.long)
+    return torch.zeros(inputs["input_ids"].shape[-1], dtype=torch.long)
+
+
+def pad_sequence(sequences, padding_side='right', padding_value=0):
+    """
+    Pad a list of sequences to the same length.
+    sequences: list of tensors in [seq_len, *] shape
+    """
+    assert padding_side in ['right', 'left']
+    max_size = sequences[0].size()
+    trailing_dims = max_size[1:]
+    max_len = max(len(seq) for seq in sequences)
+    batch_size = len(sequences)
+    output = sequences[0].new_full((batch_size, max_len) + trailing_dims, padding_value)
+    for i, seq in enumerate(sequences):
+        length = seq.size(0)
+        if padding_side == 'right':
+            output.data[i, :length] = seq
+        else:
+            output.data[i, -length:] = seq
+    return output
+
+def get_image_info(image_path, min_pixel, max_pixel, width, height):
+    # Using this because of process_vision_info function
+    # Need to fix this in the future
+
+
+    content = {
+        "type": "image", 
+        "image": image_path,
+        "min_pixel": min_pixel,
+        "max_pixel": max_pixel
+    }
+
+    if width is not None and height is not None:
+        content["resized_width"] = width
+        content["resized_height"] = height
+
+    messages = [
+        {"role": "user", 
+         "content": [content]
+        }
+    ]
+
+    image_input, _ = process_vision_info(messages)
+
+    return image_input[0]
+
+def get_video_info(video_path, min_pixels, max_pixels, fps):
+    # Using this because of process_vision_info function
+    # Need to fix this in the future
+
+    messages = [
+        {"role": "user", 
+         "content": [
+             {
+                "type": "video", 
+                "video": video_path,
+                "min_pixels": min_pixels,
+                "max_pixels": max_pixels,
+                "fps": fps
+            }
+            ]
+        }
+    ]
+
+    _, video_input, video_kwargs = process_vision_info(messages, return_video_kwargs=True)
+
+    return video_input[0], video_kwargs
+
+def add_anchor_pad(user_input, anchor_nums, anchor_tokens):
+    # add anchor pad after VISION_END_TOKEN or ANCHOR_END_TOKEN
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    anchor_pads = "".join(anchor_pads)
+    if VISION_END_TOKEN in user_input:
+        user_input = user_input.replace(VISION_END_TOKEN, VISION_END_TOKEN + anchor_pads)
+    return user_input
+
+def add_cot_anchor_pad_in_user_input(user_input, anchor_nums, anchor_tokens,anchor_names):
+    if len(anchor_nums) == 0:
+        return user_input
+    
+    anchor_pads = []
+    for anchor_num, anchor_token, _ in zip(anchor_nums, anchor_tokens, anchor_names):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_pad = ""
+    if len(anchor_pads) == 1:
+        CoT_pad = f"The {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
+            if i == 0:
+                CoT_pad += f"The {anchor_name} of the image is {anchor_pad}, "
+            elif i == len(anchor_names) - 1:
+                CoT_pad += f"and the {anchor_name} of the image is {anchor_pad}. "
+            else:
+                CoT_pad += f"the {anchor_name} of the image is {anchor_pad}, "
+    user_input = CoT_pad + user_input
+    return user_input
+
+def get_cot_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_start = "Because "
+    if len(anchor_names) == 1:
+        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
+            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
+            if anchor_name == anchor_names[-2]:
+                CoT_start += ", and "
+            elif anchor_name == anchor_names[-1]:
+                CoT_start += ". "
+            else:
+                CoT_start += ", "
+    response = CoT_start + response
+    return response
+
+
+COT_TEMPLATES = [
+    {
+        "name": "basic_causal",
+        "single": "Because the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Because the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "observational",
+        "single": "I can observe that the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "I can observe that the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "analytical",
+        "single": "After analyzing the image, the {anchor_name} is {anchor_pad}. ",
+        "multiple": "After analyzing the image, the {anchor_name} is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "descriptive",
+        "single": "The image shows that the {anchor_name} is {anchor_pad}. ",
+        "multiple": "The image shows that the {anchor_name} is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "conditional",
+        "single": "Given that the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Given that the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "evidence_based",
+        "single": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    }
+]
+
+def get_random_cot_template():
+    return random.choice(COT_TEMPLATES)
+
+def apply_cot_template(template, anchor_names, anchor_pads):
+    if len(anchor_names) == 1:
+        return template["single"].format(
+            anchor_name=anchor_names[0], 
+            anchor_pad=anchor_pads[0]
+        )
+    else:
+        result = ""
+        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
+            if i == len(anchor_names) - 1: 
+                connector = template["connectors"]["last"]
+            elif i == len(anchor_names) - 2: 
+                connector = template["connectors"]["second_last"]
+            else: 
+                connector = template["connectors"]["middle"]
+            
+            result += template["multiple"].format(
+                anchor_name=anchor_name,
+                anchor_pad=anchor_pad,
+                connector=connector
+            )
+        return result
+
+def get_templates_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    
+    template = get_random_cot_template()
+    
+    cot_text = apply_cot_template(template, anchor_names, anchor_pads)
+    
+    response = "<think>" + cot_text + "</think>" + "<answer>" + response + "</answer>"
+    return response
+        
+def get_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_start = "<think> Because "
+    if len(anchor_names) == 1:
+        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
+            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
+            if anchor_name == anchor_names[-2]:
+                CoT_start += ", and "
+            elif anchor_name == anchor_names[-1]:
+                CoT_start += ". "
+            else:
+                CoT_start += ", "
+    response = CoT_start + " </think>\n" + "<answer> " + response + " </answer>"
+    return response
+    
+def get_feature_data(user_input, gpt_response, anchor_nums, anchor_tokens, anchor_names):
+    anchor_pads = []
+    for anchor_num, anchor_token, anchor_name in zip(anchor_nums, anchor_tokens, anchor_names):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    anchor_name = ", ".join(anchor_names)
+    anchor_pads = "".join(anchor_pads)
+    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}What is the {anchor_name} of the image?\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+    gpt_response = f"{anchor_pads}\n{DEFAULT_IM_END_TOKEN}\n"
+    return user_input, gpt_response
+
+def replace_pad_with_anchor_tokens(gpt_response):
+    token_dict = {
+        "<segmentation>": SAM_PAD_TOKEN * 8,
+        "<depth>": DEPTH_PAD_TOKEN * 4,
+        "<dino>": DINO_PAD_TOKEN * 4,
+        "<pidinet>": PIDINET_PAD_TOKEN * 4,
+        "<siglip>": SIGLIP_PAD_TOKEN * 4,
+        "<metaclip>": METACLIP_PAD_TOKEN * 4,
+    }
+    for token, anchor_token in token_dict.items():
+        gpt_response = gpt_response.replace(token, anchor_token)
+    return gpt_response
+    
+def get_token_num(anchor_model_id):
+    token_nums = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            token_nums.append(8)
+        elif anchor_model == "dino":
+            token_nums.append(4)
+        elif anchor_model == "depth":
+            token_nums.append(4)
+        elif anchor_model == "InternViT":
+            token_nums.append(4)
+        elif anchor_model == "pidinet":
+            token_nums.append(4)
+        elif anchor_model == "siglip":
+            token_nums.append(4)
+        elif anchor_model == "metaclip":
+            token_nums.append(4)
+    return token_nums
+
+def get_anchor_token(anchor_model_id):
+    anchor_tokens = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            anchor_tokens.append(SAM_PAD_TOKEN)
+        elif anchor_model == "dino":
+            anchor_tokens.append(DINO_PAD_TOKEN)
+        elif anchor_model == "depth":
+            anchor_tokens.append(DEPTH_PAD_TOKEN)
+        elif anchor_model == "InternViT":
+            anchor_tokens.append(INTERN_PAD_TOKEN)
+        elif anchor_model == "pidinet":
+            anchor_tokens.append(PIDINET_PAD_TOKEN)
+        elif anchor_model == "siglip":
+            anchor_tokens.append(SIGLIP_PAD_TOKEN)
+        elif anchor_model == "metaclip":
+            anchor_tokens.append(METACLIP_PAD_TOKEN)
+    return anchor_tokens
+
+def get_anchor_task_name(anchor_model_id):
+    anchor_task_names = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            anchor_task_names.append("segmentation")
+        elif anchor_model == "dino":
+            anchor_task_names.append("perception feature")
+        elif anchor_model == "depth":
+            anchor_task_names.append("depth map")
+        elif anchor_model == "InternViT":
+            anchor_task_names.append("caption")
+        elif anchor_model == "pidinet":
+            anchor_task_names.append("edge map")
+        elif anchor_model == "siglip":
+            anchor_task_names.append("clip feature")
+        elif anchor_model == "metaclip":
+            anchor_task_names.append("metaclip feature")
+    return anchor_task_names
+    
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(
+        self,
+        data_path: str | list,
+        processor: transformers.ProcessorMixin,
+        data_args: DataArguments,
+        model_id,
+        padding=True,
+        shuffle=True,
+        random_seed=42,
+        anchor_model_id=None,
+    ):
+        super(SupervisedDataset, self).__init__()
+        if isinstance(data_path, str):
+            list_data_dict = json.load(open(data_path, "r"))
+        else:
+            list_data_dict = data_path
+
+        self.model_id = model_id
+        self.processor = processor
+        self.list_data_dict = list_data_dict
+        self.data_args = data_args
+        self.padding = padding
+        self.image_min_pixel = data_args.image_min_pixels
+        self.image_max_pixel = data_args.image_max_pixels
+        self.image_resized_w = data_args.image_resized_width
+        self.image_resized_h = data_args.image_resized_height
+        self.video_min_pixel = data_args.video_min_pixels
+        self.video_max_pixel = data_args.video_max_pixels
+        self.fps = data_args.fps
+        self.anchor_model_id = anchor_model_id
+        self.anchor_token_nums = get_token_num(anchor_model_id)
+        self.anchor_tokens = get_anchor_token(anchor_model_id)
+        self.anchor_task_names = get_anchor_task_name(anchor_model_id)
+        
+        self.cur_step = 0
+        self.stage_0_step = data_args.stage_0_step
+        self.stage_1_step = data_args.stage_1_step
+        self.stage_2_step = data_args.stage_2_step
+        
+        # for shuffle
+        self.rng = np.random.default_rng(seed=random_seed)
+        
+        if shuffle:
+            self.rng.shuffle(self.list_data_dict)
+    
+    def set_cur_step(self, step: int):
+        self.cur_step = step
+        print(f"[Dataset] cur_step has been set to {step}")
+        
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        
+        # import ipdb; ipdb.set_trace()
+        
+        sources = self.list_data_dict[i]
+        
+        is_video = False
+
+        processor = self.processor
+        if "image" in sources:
+            videos = None
+            grid_key = "image_grid_thw"
+            pixel_key = "pixel_values"
+            
+            image_files = sources["image"]
+            image_folder = self.data_args.image_folder
+
+            if isinstance(image_files, str):
+                image_files = Image.open(image_files).convert("RGB")
+                image_files = [image_files]
+            else:
+                image_files = [Image.open(image_file).convert("RGB") for image_file in image_files]
+            images = []
+            
+            for image_file in image_files:
+                # if not os.path.exists(image_file):
+                #     if not image_file.startswith("http"):
+                #         image_file = os.path.join(image_folder, image_file)
+                #         images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+                # else:
+                images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+
+        elif "video" in sources:
+            is_video = True
+            images=None
+            grid_key = "video_grid_thw"
+            pixel_key = "pixel_values_videos"
+
+            video_files = sources["video"]
+            video_folder = self.data_args.image_folder
+
+            if isinstance(video_files, str):
+                video_files = [video_files]
+
+            videos = []
+            for video_file in video_files:
+                if not os.path.exists(video_file):
+                    if not video_file.startswith("http"):
+                        video_file = os.path.join(video_folder, video_file)
+                video_input, video_kwargs = get_video_info(video_file, self.video_min_pixel, self.video_max_pixel, self.data_args.fps)
+                videos.append(video_input)
+        else:
+            grid_key = None
+            pixel_key = None
+            images=None
+            videos=None
+           
+        if images is None:
+            print("No image or video found in the data.")
+            images = []
+            # Create a black image as a placeholder
+            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+
+        elif len(images) == 0:
+            print("No image or video found in the data.")
+            # Create a black image as a placeholder
+            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+        
+        if videos is not None:
+            
+            # import ipdb; ipdb.set_trace()
+            pass
+            
+        sources = copy.deepcopy(llava_to_openai(sources['conversations'], is_video=is_video))
+
+        all_input_ids = []
+        all_mm_token_type_ids = []
+        all_labels = []
+        all_pixel_values = []
+        all_image_grid_thw = []
+        all_second_gird = []
+        
+        # all_dino_encoded_values = []
+
+        # Qwen2-VL uses a default system message so I've added this.
+        if len(SYSTEM_MESSAGE) > 0:
+            system_message = f"{DEFAULT_IM_START_TOKEN}system\n{SYSTEM_MESSAGE}\n{DEFAULT_IM_END_TOKEN}\n"
+            system_message_input_ids = processor.tokenizer(system_message, add_special_tokens=False, return_tensors='pt')['input_ids']
+            system_labels = torch.full_like(system_message_input_ids, IGNORE_INDEX) 
+            
+            all_input_ids.append(system_message_input_ids.squeeze(0))
+            all_mm_token_type_ids.append(
+                torch.zeros(system_message_input_ids.shape[1], dtype=torch.long)
+            )
+            all_labels.append(system_labels.squeeze(0))
+        
+        # import ipdb; ipdb.set_trace()
+        for _, j in enumerate(range(0, len(sources), 2)):
+            
+            if j >= 2:
+                break
+            
+            user_input = sources[j]
+            gpt_response = sources[j + 1]
+                        
+            if (DEFAULT_IMAGE_TOKEN not in user_input['content']) and (DEFAULT_VIDEO_TOKEN not in user_input['content']) and (LLAVA_IMAGE_TOKEN in user_input['content']):
+                user_input = f"{DEFAULT_IM_START_TOKEN}{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                user_input = add_cot_anchor_pad_in_user_input(user_input, self.anchor_token_nums, self.anchor_tokens,self.anchor_task_names)
+                gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                raise ValueError('Every man is a poet when he is in love')
+            else:
+                if self.cur_step < self.stage_0_step:
+                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    user_input = add_cot_anchor_pad_in_user_input(user_input, self.anchor_token_nums, self.anchor_tokens,self.anchor_task_names)
+                    gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                elif self.cur_step < self.stage_1_step:
+                    user_input, gpt_response = get_feature_data(user_input, gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
+                elif self.cur_step < self.stage_2_step:
+                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    gpt_response = f"{gpt_response['content']}"
+                    if DEFAULT_IMAGE_TOKEN in user_input:
+                        gpt_response = get_comt_data_in_response(gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
+                    gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # print(f"\033[92m gpt_response: {gpt_response}\033[0m")
+                else:
+                    # user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    # gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # gpt_response = replace_pad_with_anchor_tokens(gpt_response)
+                    import random
+                    xxx = random.randint(0, 5)
+                    if xxx == 0:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    else:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}"
+                        if DEFAULT_IMAGE_TOKEN in user_input:
+                            # INSERT_YOUR_CODE
+                            total = len(self.anchor_tokens)
+                            if total == 0:
+                                selected_anchor_token_nums = []
+                                selected_anchor_tokens = []
+                                selected_anchor_task_names = []
+                            else:
+                                x = random.randint(1, total)
+                                idxs = sorted(random.sample(range(total), x)) if x > 0 else []
+                                selected_anchor_token_nums = [self.anchor_token_nums[i] for i in idxs]
+                                selected_anchor_tokens = [self.anchor_tokens[i] for i in idxs]
+                                selected_anchor_task_names = [self.anchor_task_names[i] for i in idxs]
+                            gpt_response = get_comt_data_in_response(gpt_response, selected_anchor_token_nums, selected_anchor_tokens, selected_anchor_task_names)
+                        gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+            
+            # print(f'the user_input is {user_input}')
+            # print(f'the gpt_response is {gpt_response}')
+                    
+
+            # print("-----------------")
+            # print(user_input, gpt_response)
+            # print("-----------------")
+
+            # import ipdb; ipdb.set_trace()
+            if DEFAULT_IMAGE_TOKEN in user_input:
+                inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
+                prompt_input_ids = inputs['input_ids']
+                all_mm_token_type_ids.append(_mm_type_ids_for_processor_batch(processor, inputs))
+                # raise ValueError('Every man is a poet when he is in love')
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+                
+                # del dino_val
+                torch.cuda.empty_cache()
+                
+            elif DEFAULT_VIDEO_TOKEN in user_input:
+                if "Qwen2.5" in self.model_id:
+                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt', **video_kwargs)
+                    all_second_gird.extend(inputs["second_per_grid_ts"])
+                else:
+                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
+                prompt_input_ids = inputs['input_ids']
+                all_mm_token_type_ids.append(_mm_type_ids_for_processor_batch(processor, inputs))
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+
+            else:
+                prompt_input_ids = processor.tokenizer(user_input, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+                all_mm_token_type_ids.append(
+                    torch.zeros(prompt_input_ids.shape[1], dtype=torch.long)
+                )
+
+            response_input_ids = processor.tokenizer(gpt_response, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+            all_mm_token_type_ids.append(
+                torch.zeros(response_input_ids.shape[1], dtype=torch.long)
+            )
+
+            input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
+            labels = torch.cat(
+                [
+                    torch.tensor([IGNORE_INDEX] * len(prompt_input_ids[0])),  
+                    response_input_ids.squeeze(0),
+                ],
+                dim=0,
+            )
+
+            all_input_ids.append(input_ids)
+            all_labels.append(labels)
+        
+        # There is no need for eos or bos tokens in the input_ids
+        # Qwen2-VL does not use them
+        input_ids = torch.cat(all_input_ids, dim=0).to(torch.long)
+        labels = torch.cat(all_labels, dim=0).to(torch.long)
+        mm_token_type_ids = torch.cat(all_mm_token_type_ids, dim=0).to(torch.long)
+        if mm_token_type_ids.shape[0] != input_ids.shape[0]:
+            raise ValueError(
+                f"mm_token_type_ids length {mm_token_type_ids.shape[0]} != input_ids {input_ids.shape[0]}"
+            )
+
+        # eos_token_id = processor.tokenizer.convert_tokens_to_ids(DEFAULT_IM_END_TOKEN)
+        # input_ids, labels = truncate_sequence(input_ids, labels, self.max_length, eos_token_id)
+
+        attention_mask = (input_ids > -1000000).to(torch.long)
+        task_type = "und_image"
+        data_dict = dict(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            mm_token_type_ids=mm_token_type_ids,
+            task_types=task_type
+        )
+
+        if pixel_key and grid_key:
+            pixel_values = torch.cat(all_pixel_values, dim=0)
+            image_thw = torch.cat(all_image_grid_thw, dim=0)
+            
+            data_dict[pixel_key] = pixel_values
+            data_dict[grid_key] = image_thw
+            data_dict["image_files"] = image_files
+
+        if len(all_second_gird) > 0:
+            second_gird = all_second_gird
+            data_dict["second_per_grid_ts"] = second_gird
+        
+        self.cur_step += 1
+        return data_dict
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning; truncates to tokenizer.model_max_length (same as train_stage1_algin)."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __post_init__(self):
+        self.pad_token_id = self.tokenizer.pad_token_id
+
+    def __call__(self, examples):
+        max_len = self.tokenizer.model_max_length
+        batch_input_ids = []
+        batch_label_ids = []
+        batch_pixel_values = []
+        batch_pixel_video_values = []
+        batch_video_thw = []
+        batch_image_thw = []
+        batch_second_per_grid_ts = []
+        batch_mm_token_type_ids = []
+
+        batch_image_files = []
+        
+        for example in examples:
+            keys = example.keys()
+            if "pixel_values_videos" in keys:
+                batch_pixel_video_values.append(example["pixel_values_videos"])
+                batch_video_thw.append(example["video_grid_thw"])
+            elif "pixel_values" in keys:
+                batch_pixel_values.append(example["pixel_values"])
+                batch_image_thw.append(example["image_grid_thw"])
+            
+            if "image_files" in keys:
+                batch_image_files.append(example["image_files"])
+            
+            iid = example["input_ids"]
+            lab = example["labels"]
+            if iid.shape[0] > max_len:
+                iid = iid[:max_len]
+                lab = lab[:max_len]
+            batch_input_ids.append(iid)
+            batch_label_ids.append(lab)
+            if "mm_token_type_ids" in keys:
+                mm = example["mm_token_type_ids"]
+                if mm.shape[0] > max_len:
+                    mm = mm[:max_len]
+                batch_mm_token_type_ids.append(mm)
+
+            if "second_per_grid_ts" in keys:
+                batch_second_per_grid_ts.extend(example["second_per_grid_ts"])
+        
+        input_ids = pad_sequence(
+            batch_input_ids, padding_side='right', padding_value=self.pad_token_id
+        )
+
+        labels = pad_sequence(batch_label_ids, padding_side='right', padding_value=IGNORE_INDEX)
+
+        if input_ids.shape[1] > max_len:
+            rank0_print(
+                f"Warning: input with length {input_ids.shape[1]} is longer than max length {max_len}. Truncating."
+            )
+        input_ids = input_ids[:, :max_len]
+        labels = labels[:, :max_len]
+        attention_mask = input_ids != self.pad_token_id
+
+        data_dict = {
+            'input_ids': input_ids,
+            'labels': labels,
+            'attention_mask': attention_mask,
+        }
+
+        if len(batch_mm_token_type_ids) > 0:
+            mm_pad = pad_sequence(
+                batch_mm_token_type_ids, padding_side='right', padding_value=0
+            )
+            data_dict["mm_token_type_ids"] = mm_pad[:, :max_len]
+
+        if len(batch_pixel_values) > 0:
+            pixel_values = torch.cat(batch_pixel_values, dim=0)
+            image_thw = torch.cat(batch_image_thw, dim=0)
+            data_dict["pixel_values"] = pixel_values
+            data_dict["image_grid_thw"] = image_thw
+
+        if len(batch_pixel_video_values) > 0:
+            pixel_video_values = torch.cat(batch_pixel_video_values, dim=0)
+            video_thw = torch.cat(batch_video_thw, dim=0)
+            data_dict["pixel_values_videos"] = pixel_video_values
+            data_dict["video_grid_thw"] = video_thw
+
+        if len(batch_second_per_grid_ts) > 0:
+            data_dict["second_per_grid_ts"] = batch_second_per_grid_ts
+            
+        if len(batch_image_files) > 0:
+            data_dict["image_files"] = batch_image_files
+        return data_dict
+
+def replace_image_tokens(input_string, is_video=False):
+    if is_video:
+        pattern = r'\n?' + re.escape(LLAVA_VIDEO_TOKEN) + r'\n?'
+        replacement = VISION_START_TOKEN + DEFAULT_VIDEO_TOKEN + VISION_END_TOKEN
+    else:
+        pattern = r'\n?' + re.escape(LLAVA_IMAGE_TOKEN) + r'\n?'
+        replacement = VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN
+
+    return re.sub(pattern, replacement, input_string)
+
+def llava_to_openai(conversations, is_video=False):
+    role_mapping = {"human": "user", "gpt": "assistant"}
+
+    transformed_data = []
+    for conversation in conversations:
+        transformed_content = replace_image_tokens(conversation["value"], is_video=is_video)
+        transformed_entry = {
+            "role": role_mapping.get(conversation["from"], conversation["from"]),
+            "content": transformed_content,
+        }
+        transformed_data.append(transformed_entry)
+
+    return transformed_data
+
+def make_supervised_data_module(model_id, processor, data_args, anchor_model_id):
+    """Make dataset and collator for supervised fine-tuning."""
+    sft_dataset = SupervisedDataset(
+        data_path=data_args.data_path, processor=processor, data_args=data_args, model_id=model_id, anchor_model_id=anchor_model_id
+    )
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=processor.tokenizer)
+
+    return dict(train_dataset=sft_dataset,
+                eval_dataset=None,
+                data_collator=data_collator)

code/umm/runtime/models/blip3o/data/covt_data_van.py

@@ -0,0 +1,1056 @@
+import copy
+import os
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Sequence, List, Tuple, Any
+from collections.abc import Sequence
+import torch
+import transformers
+import ujson as json
+from torch.utils.data import Dataset
+from qwen_vl_utils import process_vision_info
+from PIL import Image
+from transformers import AutoImageProcessor
+import re
+import numpy as np
+import cv2
+from torchvision import transforms
+import random
+
+from models.blip3o.model.anchors.segment_anything import build_sam_vit_h, sam_model_registry, SamPredictor
+from diffusers import AutoencoderKL
+from transformers import AutoModel, CLIPImageProcessor
+
+from models.blip3o.params import DataArguments
+from models.blip3o.constants import *
+
+
+def rank0_print(*args):
+    if os.environ.get("LOCAL_RANK", "0") in ("0", "-1"):
+        print(*args)
+
+def get_rope_index_3(
+    spatial_merge_size: Optional[int] = 2,
+    input_ids: Optional[torch.LongTensor] = None,
+    image_grid_thw: Optional[torch.LongTensor] = None,
+    video_grid_thw: Optional[torch.LongTensor] = None,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids."""
+    if input_ids is not None:
+        while input_ids.dim() < 2:
+            input_ids = input_ids.unsqueeze(0)
+    if attention_mask is not None:
+        while attention_mask.dim() < 2:
+            attention_mask = attention_mask.unsqueeze(0)
+    # Since we use timestamps to seperate videos, like <t1> <vision_start> <frame1> <vision_end> <t2> <vision_start> <frame2> <vision_end>, the video_grid_thw should also be split
+    if video_grid_thw is not None:
+        video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+        video_grid_thw[:, 0] = 1
+
+    image_token_id = 151655
+    video_token_id = 151656
+    vision_start_token_id = 151652
+    mrope_position_deltas = []
+    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+        total_input_ids = input_ids
+        if attention_mask is None:
+            attention_mask = torch.ones_like(total_input_ids)
+        position_ids = torch.ones(
+            3,
+            input_ids.shape[0],
+            input_ids.shape[1],
+            dtype=input_ids.dtype,
+            device=input_ids.device,
+        )
+        image_index, video_index = 0, 0
+        attention_mask = attention_mask.to(total_input_ids.device)
+        for i, input_ids in enumerate(total_input_ids):
+            input_ids = input_ids[attention_mask[i] == 1]
+            image_nums, video_nums = 0, 0
+            vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+            vision_tokens = input_ids[vision_start_indices + 1]
+            image_nums = (vision_tokens == image_token_id).sum()
+            video_nums = (vision_tokens == video_token_id).sum()
+            input_tokens = input_ids.tolist()
+            llm_pos_ids_list: list = []
+            st = 0
+            remain_images, remain_videos = image_nums, video_nums
+            for _ in range(image_nums + video_nums):
+                if image_token_id in input_tokens and remain_images > 0:
+                    ed_image = input_tokens.index(image_token_id, st)
+                else:
+                    ed_image = len(input_tokens) + 1
+                if video_token_id in input_tokens and remain_videos > 0:
+                    ed_video = input_tokens.index(video_token_id, st)
+                else:
+                    ed_video = len(input_tokens) + 1
+                if ed_image < ed_video:
+                    t, h, w = (
+                        image_grid_thw[image_index][0],
+                        image_grid_thw[image_index][1],
+                        image_grid_thw[image_index][2],
+                    )
+                    image_index += 1
+                    remain_images -= 1
+                    ed = ed_image
+
+                else:
+                    t, h, w = (
+                        video_grid_thw[video_index][0],
+                        video_grid_thw[video_index][1],
+                        video_grid_thw[video_index][2],
+                    )
+                    video_index += 1
+                    remain_videos -= 1
+                    ed = ed_video
+                llm_grid_t, llm_grid_h, llm_grid_w = (
+                    t.item(),
+                    h.item() // spatial_merge_size,
+                    w.item() // spatial_merge_size,
+                )
+                text_len = ed - st
+
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos)
+                t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+                llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+            if st < len(input_tokens):
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                text_len = len(input_tokens) - st
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+            position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+            mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+        return position_ids, mrope_position_deltas
+    else:
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+            max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+            mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+        else:
+            position_ids = (
+                torch.arange(input_ids.shape[1], device=input_ids.device)
+                .view(1, 1, -1)
+                .expand(3, input_ids.shape[0], -1)
+            )
+            mrope_position_deltas = torch.zeros(
+                [input_ids.shape[0], 1],
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+
+        return position_ids, mrope_position_deltas
+
+
+def truncate_sequence(input_ids, labels, max_length, eos_token_id):
+    if input_ids.size(0) > max_length:
+        input_ids = input_ids[:max_length-1]
+        labels = labels[:max_length-1]
+
+    if eos_token_id is not None:
+        input_ids = torch.cat([input_ids, torch.tensor([eos_token_id])])
+        labels = torch.cat([labels, torch.tensor([eos_token_id])])
+
+    return input_ids, labels
+
+def pad_sequence(sequences, padding_side='right', padding_value=0):
+    """
+    Pad a list of sequences to the same length.
+    sequences: list of tensors in [seq_len, *] shape
+    """
+    assert padding_side in ['right', 'left']
+    max_size = sequences[0].size()
+    trailing_dims = max_size[1:]
+    max_len = max(len(seq) for seq in sequences)
+    batch_size = len(sequences)
+    output = sequences[0].new_full((batch_size, max_len) + trailing_dims, padding_value)
+    for i, seq in enumerate(sequences):
+        length = seq.size(0)
+        if padding_side == 'right':
+            output.data[i, :length] = seq
+        else:
+            output.data[i, -length:] = seq
+    return output
+
+def get_image_info(image_path, min_pixel, max_pixel, width, height):
+    # Using this because of process_vision_info function
+    # Need to fix this in the future
+
+
+    content = {
+        "type": "image", 
+        "image": image_path,
+        "min_pixel": min_pixel,
+        "max_pixel": max_pixel
+    }
+
+    if width is not None and height is not None:
+        content["resized_width"] = width
+        content["resized_height"] = height
+
+    messages = [
+        {"role": "user", 
+         "content": [content]
+        }
+    ]
+
+    image_input, _ = process_vision_info(messages)
+
+    return image_input[0]
+
+def get_video_info(video_path, min_pixels, max_pixels, fps):
+    # Using this because of process_vision_info function
+    # Need to fix this in the future
+
+    messages = [
+        {"role": "user", 
+         "content": [
+             {
+                "type": "video", 
+                "video": video_path,
+                "min_pixels": min_pixels,
+                "max_pixels": max_pixels,
+                "fps": fps
+            }
+            ]
+        }
+    ]
+
+    _, video_input, video_kwargs = process_vision_info(messages, return_video_kwargs=True)
+
+    return video_input[0], video_kwargs
+
+def add_anchor_pad(user_input, anchor_nums, anchor_tokens):
+    # add anchor pad after VISION_END_TOKEN or ANCHOR_END_TOKEN
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    anchor_pads = "".join(anchor_pads)
+    if VISION_END_TOKEN in user_input:
+        user_input = user_input.replace(VISION_END_TOKEN, VISION_END_TOKEN + anchor_pads)
+    return user_input
+
+def add_cot_anchor_pad_in_user_input(user_input, anchor_nums, anchor_tokens,anchor_names):
+    if len(anchor_nums) == 0:
+        return user_input
+    
+    anchor_pads = []
+    for anchor_num, anchor_token, _ in zip(anchor_nums, anchor_tokens, anchor_names):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_pad = ""
+    if len(anchor_pads) == 1:
+        CoT_pad = f"The {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
+            if i == 0:
+                CoT_pad += f"The {anchor_name} of the image is {anchor_pad}, "
+            elif i == len(anchor_names) - 1:
+                CoT_pad += f"and the {anchor_name} of the image is {anchor_pad}. "
+            else:
+                CoT_pad += f"the {anchor_name} of the image is {anchor_pad}, "
+    user_input = CoT_pad + user_input
+    return user_input
+
+def get_cot_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_start = "Because "
+    if len(anchor_names) == 1:
+        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
+            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
+            if anchor_name == anchor_names[-2]:
+                CoT_start += ", and "
+            elif anchor_name == anchor_names[-1]:
+                CoT_start += ". "
+            else:
+                CoT_start += ", "
+    response = CoT_start + response
+    return response
+
+
+COT_TEMPLATES = [
+    {
+        "name": "basic_causal",
+        "single": "Because the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Because the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "observational",
+        "single": "I can observe that the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "I can observe that the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "analytical",
+        "single": "After analyzing the image, the {anchor_name} is {anchor_pad}. ",
+        "multiple": "After analyzing the image, the {anchor_name} is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "descriptive",
+        "single": "The image shows that the {anchor_name} is {anchor_pad}. ",
+        "multiple": "The image shows that the {anchor_name} is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "conditional",
+        "single": "Given that the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Given that the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "evidence_based",
+        "single": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    }
+]
+
+def get_random_cot_template():
+    return random.choice(COT_TEMPLATES)
+
+def apply_cot_template(template, anchor_names, anchor_pads):
+    if len(anchor_names) == 1:
+        return template["single"].format(
+            anchor_name=anchor_names[0], 
+            anchor_pad=anchor_pads[0]
+        )
+    else:
+        result = ""
+        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
+            if i == len(anchor_names) - 1: 
+                connector = template["connectors"]["last"]
+            elif i == len(anchor_names) - 2: 
+                connector = template["connectors"]["second_last"]
+            else: 
+                connector = template["connectors"]["middle"]
+            
+            result += template["multiple"].format(
+                anchor_name=anchor_name,
+                anchor_pad=anchor_pad,
+                connector=connector
+            )
+        return result
+
+def get_templates_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    
+    template = get_random_cot_template()
+    
+    cot_text = apply_cot_template(template, anchor_names, anchor_pads)
+    
+    response = "<think>" + cot_text + "</think>" + "<answer>" + response + "</answer>"
+    return response
+        
+def get_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    # 3DThinker's placement ablation favors geometry latents at the response
+    # boundary.  Keep the user's CoVT requirement (inside <think>) while moving
+    # only the VGGT block to the very beginning of that region, before prose.
+    geometry_items = [
+        (name, pad)
+        for name, pad in zip(anchor_names, anchor_pads)
+        if name == "3D geometry"
+    ]
+    other_items = [
+        (name, pad)
+        for name, pad in zip(anchor_names, anchor_pads)
+        if name != "3D geometry"
+    ]
+    geometry_prefix = "".join(pad for _, pad in geometry_items)
+    CoT_start = "<think> " + geometry_prefix
+    if geometry_items:
+        CoT_start += " Using the 3D geometry above, "
+    else:
+        CoT_start += " Because "
+
+    if len(other_items) == 0:
+        CoT_start += "the answer is grounded in the image geometry. "
+    elif len(other_items) == 1:
+        CoT_start += f"the {other_items[0][0]} of the image is {other_items[0][1]}. "
+    else:
+        other_names = [name for name, _ in other_items]
+        for anchor_name, anchor_pad in other_items:
+            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
+            if anchor_name == other_names[-2]:
+                CoT_start += ", and "
+            elif anchor_name == other_names[-1]:
+                CoT_start += ". "
+            else:
+                CoT_start += ", "
+    response = CoT_start + " </think>\n" + "<answer> " + response + " </answer>"
+    return response
+    
+def get_feature_data(user_input, gpt_response, anchor_nums, anchor_tokens, anchor_names):
+    anchor_pads = []
+    for anchor_num, anchor_token, anchor_name in zip(anchor_nums, anchor_tokens, anchor_names):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    anchor_name = ", ".join(anchor_names)
+    anchor_pads = "".join(anchor_pads)
+    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}What is the {anchor_name} of the image?\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+    gpt_response = f"{anchor_pads}\n{DEFAULT_IM_END_TOKEN}\n"
+    return user_input, gpt_response
+
+def replace_pad_with_anchor_tokens(gpt_response):
+    token_dict = {
+        "<segmentation>": SAM_PAD_TOKEN * COVT_SAM_TOKEN_COUNT,
+        "<depth>": DEPTH_PAD_TOKEN * COVT_DEPTH_TOKEN_COUNT,
+        "<dino>": DINO_PAD_TOKEN * COVT_DINO_TOKEN_COUNT,
+        "<pidinet>": PIDINET_PAD_TOKEN * COVT_PIDINET_TOKEN_COUNT,
+        "<siglip>": SIGLIP_PAD_TOKEN * COVT_SIGLIP_TOKEN_COUNT,
+        "<metaclip>": METACLIP_PAD_TOKEN * 4,
+    }
+    for token, anchor_token in token_dict.items():
+        gpt_response = gpt_response.replace(token, anchor_token)
+    return gpt_response
+    
+def get_token_num(anchor_model_id):
+    token_nums = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            token_nums.append(COVT_SAM_TOKEN_COUNT)
+        elif anchor_model == "dino":
+            token_nums.append(COVT_DINO_TOKEN_COUNT)
+        elif anchor_model == "depth":
+            token_nums.append(COVT_DEPTH_TOKEN_COUNT)
+        elif anchor_model == "InternViT":
+            token_nums.append(4)
+        elif anchor_model == "pidinet":
+            token_nums.append(COVT_PIDINET_TOKEN_COUNT)
+        elif anchor_model == "siglip":
+            token_nums.append(COVT_SIGLIP_TOKEN_COUNT)
+        elif anchor_model == "metaclip":
+            token_nums.append(4)
+    return token_nums
+
+def get_anchor_token(anchor_model_id):
+    anchor_tokens = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            anchor_tokens.append(SAM_PAD_TOKEN)
+        elif anchor_model == "dino":
+            anchor_tokens.append(DINO_PAD_TOKEN)
+        elif anchor_model == "depth":
+            anchor_tokens.append(DEPTH_PAD_TOKEN)
+        elif anchor_model == "InternViT":
+            anchor_tokens.append(INTERN_PAD_TOKEN)
+        elif anchor_model == "pidinet":
+            anchor_tokens.append(PIDINET_PAD_TOKEN)
+        elif anchor_model == "siglip":
+            anchor_tokens.append(SIGLIP_PAD_TOKEN)
+        elif anchor_model == "metaclip":
+            anchor_tokens.append(METACLIP_PAD_TOKEN)
+    return anchor_tokens
+
+def get_anchor_task_name(anchor_model_id):
+    anchor_task_names = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            anchor_task_names.append("segmentation")
+        elif anchor_model == "dino":
+            anchor_task_names.append("perception feature")
+        elif anchor_model == "depth":
+            anchor_task_names.append("3D geometry")
+        elif anchor_model == "InternViT":
+            anchor_task_names.append("caption")
+        elif anchor_model == "pidinet":
+            anchor_task_names.append("edge map")
+        elif anchor_model == "siglip":
+            anchor_task_names.append("clip feature")
+        elif anchor_model == "metaclip":
+            anchor_task_names.append("metaclip feature")
+    return anchor_task_names
+    
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(
+        self,
+        data_path: str | list,
+        processor: transformers.ProcessorMixin,
+        data_args: DataArguments,
+        model_id,
+        padding=True,
+        shuffle=True,
+        random_seed=42,
+        anchor_model_id=None,
+    ):
+        super(SupervisedDataset, self).__init__()
+        if isinstance(data_path, str):
+            list_data_dict = json.load(open(data_path, "r"))
+        else:
+            list_data_dict = data_path
+
+        self.model_id = model_id
+        self.processor = processor
+        self.list_data_dict = list_data_dict
+        self.data_args = data_args
+        self.padding = padding
+        self.image_min_pixel = data_args.image_min_pixels
+        self.image_max_pixel = data_args.image_max_pixels
+        self.image_resized_w = data_args.image_resized_width
+        self.image_resized_h = data_args.image_resized_height
+        self.video_min_pixel = data_args.video_min_pixels
+        self.video_max_pixel = data_args.video_max_pixels
+        self.fps = data_args.fps
+        self.merge_size = getattr(processor.image_processor, "merge_size", 2)
+        self.get_rope_index = get_rope_index_3
+        self.anchor_model_id = anchor_model_id
+        self.anchor_token_nums = get_token_num(anchor_model_id)
+        self.anchor_tokens = get_anchor_token(anchor_model_id)
+        self.anchor_task_names = get_anchor_task_name(anchor_model_id)
+        
+        self.cur_step = 0
+        self.stage_0_step = data_args.stage_0_step
+        self.stage_1_step = data_args.stage_1_step
+        self.stage_2_step = data_args.stage_2_step
+        
+        # for shuffle
+        self.rng = np.random.default_rng(seed=random_seed)
+        
+        if shuffle:
+            self.rng.shuffle(self.list_data_dict)
+    
+    def set_cur_step(self, step: int):
+        self.cur_step = step
+        print(f"[Dataset] cur_step has been set to {step}")
+        
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        
+        # import ipdb; ipdb.set_trace()
+        
+        sources = self.list_data_dict[i]
+        
+        is_video = False
+        image_files = None
+
+        processor = self.processor
+        if "image" in sources:
+            videos = None
+            grid_key = "image_grid_thw"
+            pixel_key = "pixel_values"
+            
+            image_files = sources["image"]
+            image_folder = self.data_args.image_folder
+
+            if isinstance(image_files, str):
+                image_path = image_files
+                if not os.path.isabs(image_path):
+                    image_path = os.path.join(image_folder, image_path)
+                image_files = Image.open(image_path).convert("RGB")
+                image_files = [image_files]
+            else:
+                resolved_image_files = [
+                    image_file
+                    if os.path.isabs(image_file)
+                    else os.path.join(image_folder, image_file)
+                    for image_file in image_files
+                ]
+                image_files = [
+                    Image.open(image_file).convert("RGB")
+                    for image_file in resolved_image_files
+                ]
+            images = []
+            
+            for image_file in image_files:
+                # if not os.path.exists(image_file):
+                #     if not image_file.startswith("http"):
+                #         image_file = os.path.join(image_folder, image_file)
+                #         images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+                # else:
+                images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+
+        elif "video" in sources:
+            is_video = True
+            images=None
+            grid_key = "video_grid_thw"
+            pixel_key = "pixel_values_videos"
+
+            video_files = sources["video"]
+            video_folder = self.data_args.image_folder
+
+            if isinstance(video_files, str):
+                video_files = [video_files]
+
+            videos = []
+            for video_file in video_files:
+                if not os.path.exists(video_file):
+                    if not video_file.startswith("http"):
+                        video_file = os.path.join(video_folder, video_file)
+                video_input, video_kwargs = get_video_info(video_file, self.video_min_pixel, self.video_max_pixel, self.data_args.fps)
+                videos.append(video_input)
+        else:
+            grid_key = None
+            pixel_key = None
+            images=None
+            videos=None
+           
+        if images is None:
+            print("No image or video found in the data.")
+            images = []
+            # Create a black image as a placeholder
+            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+
+        elif len(images) == 0:
+            print("No image or video found in the data.")
+            # Create a black image as a placeholder
+            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+        
+        if videos is not None:
+            
+            # import ipdb; ipdb.set_trace()
+            pass
+            
+        sources = copy.deepcopy(llava_to_openai(sources['conversations'], is_video=is_video))
+
+        all_input_ids = []
+        all_labels = []
+        all_pixel_values = []
+        all_image_grid_thw = []
+        all_second_gird = []
+        
+        # all_dino_encoded_values = []
+
+        # Qwen2-VL uses a default system message so I've added this.
+        if len(SYSTEM_MESSAGE) > 0:
+            system_message = f"{DEFAULT_IM_START_TOKEN}system\n{SYSTEM_MESSAGE}\n{DEFAULT_IM_END_TOKEN}\n"
+            system_message_input_ids = processor.tokenizer(system_message, add_special_tokens=False, return_tensors='pt')['input_ids']
+            system_labels = torch.full_like(system_message_input_ids, IGNORE_INDEX) 
+            
+            all_input_ids.append(system_message_input_ids.squeeze(0))
+            all_labels.append(system_labels.squeeze(0))
+        
+        # import ipdb; ipdb.set_trace()
+        for _, j in enumerate(range(0, len(sources), 2)):
+            
+            if j >= 2:
+                break
+            
+            user_input = sources[j]
+            gpt_response = sources[j + 1]
+                        
+            if (DEFAULT_IMAGE_TOKEN not in user_input['content']) and (DEFAULT_VIDEO_TOKEN not in user_input['content']) and (LLAVA_IMAGE_TOKEN in user_input['content']):
+                user_input = f"{DEFAULT_IM_START_TOKEN}{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                # user_input = add_cot_anchor_pad_in_user_input(user_input, self.anchor_token_nums, self.anchor_tokens,self.anchor_task_names)
+                user_input = add_anchor_pad(user_input, self.anchor_token_nums, self.anchor_tokens)
+                gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                raise ValueError('Every man is a poet when he is in love')
+            else:
+                if self.cur_step < self.stage_0_step:
+                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    # user_input = add_cot_anchor_pad_in_user_input(user_input, self.anchor_token_nums, self.anchor_tokens,self.anchor_task_names)
+                    user_input = add_anchor_pad(user_input, self.anchor_token_nums, self.anchor_tokens)
+                    gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                elif self.cur_step < self.stage_1_step:
+                    user_input, gpt_response = get_feature_data(user_input, gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
+                elif self.cur_step < self.stage_2_step:
+                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    gpt_response = f"{gpt_response['content']}"
+                    if DEFAULT_IMAGE_TOKEN in user_input:
+                        gpt_response = get_comt_data_in_response(gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
+                    gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # print(f"\033[92m gpt_response: {gpt_response}\033[0m")
+                else:
+                    # user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    # gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # gpt_response = replace_pad_with_anchor_tokens(gpt_response)
+                    import random
+                    xxx = random.randint(0, 5)
+                    if xxx == 0:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    else:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}"
+                        if DEFAULT_IMAGE_TOKEN in user_input:
+                            # INSERT_YOUR_CODE
+                            total = len(self.anchor_tokens)
+                            if total == 0:
+                                selected_anchor_token_nums = []
+                                selected_anchor_tokens = []
+                                selected_anchor_task_names = []
+                            else:
+                                x = random.randint(1, total)
+                                idxs = sorted(random.sample(range(total), x)) if x > 0 else []
+                                selected_anchor_token_nums = [self.anchor_token_nums[i] for i in idxs]
+                                selected_anchor_tokens = [self.anchor_tokens[i] for i in idxs]
+                                selected_anchor_task_names = [self.anchor_task_names[i] for i in idxs]
+                            gpt_response = get_comt_data_in_response(gpt_response, selected_anchor_token_nums, selected_anchor_tokens, selected_anchor_task_names)
+                        gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+            
+            # print(f'the user_input is {user_input}')
+            # print(f'the gpt_response is {gpt_response}')
+                    
+
+            # print("-----------------")
+            # print(user_input, gpt_response)
+            # print("-----------------")
+
+            # import ipdb; ipdb.set_trace()
+            if DEFAULT_IMAGE_TOKEN in user_input:
+                inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
+                prompt_input_ids = inputs['input_ids']
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+                
+            elif DEFAULT_VIDEO_TOKEN in user_input:
+                if "Qwen3" in self.model_id:
+                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt', **video_kwargs)
+                    all_second_gird.extend(inputs["second_per_grid_ts"])
+                else:
+                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
+                prompt_input_ids = inputs['input_ids']
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+
+            else:
+                prompt_input_ids = processor.tokenizer(user_input, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+        
+
+            response_input_ids = processor.tokenizer(gpt_response, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+
+            input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
+            labels = torch.cat(
+                [
+                    torch.tensor([IGNORE_INDEX] * len(prompt_input_ids[0])),  
+                    response_input_ids.squeeze(0),
+                ],
+                dim=0,
+            )
+            all_input_ids.append(input_ids)
+            all_labels.append(labels)
+        
+        # There is no need for eos or bos tokens in the input_ids
+        # Qwen2-VL does not use them
+        input_ids = torch.cat(all_input_ids, dim=0).to(torch.long)
+        labels = torch.cat(all_labels, dim=0).to(torch.long)
+
+        # eos_token_id = processor.tokenizer.convert_tokens_to_ids(DEFAULT_IM_END_TOKEN)
+        # input_ids, labels = truncate_sequence(input_ids, labels, self.max_length, eos_token_id)
+
+        attention_mask = (input_ids > -1000000).to(torch.long)
+        task_type = "und_image"
+
+        x_ids, x_mask = input_ids, attention_mask
+        while x_ids.dim() < 2:
+            x_ids = x_ids.unsqueeze(0)
+    
+        input_ids_2d= x_ids
+        data_dict = dict(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            task_types=task_type,
+        )
+
+        if pixel_key and grid_key:
+            pixel_values = torch.cat(all_pixel_values, dim=0)
+            image_thw = torch.cat(all_image_grid_thw, dim=0)
+            
+            data_dict[pixel_key] = pixel_values
+            data_dict[grid_key] = image_thw
+            data_dict["image_files"] = image_files
+
+        if len(all_second_gird) > 0:
+            second_gird = all_second_gird
+            data_dict["second_per_grid_ts"] = second_gird
+
+        position_ids, _ = self.get_rope_index(
+            self.merge_size,
+            input_ids_2d,
+            image_grid_thw=torch.cat(all_image_grid_thw, dim=0) if images else None,
+            video_grid_thw=(
+                torch.cat(all_image_grid_thw, dim=0) if videos else None
+            ),
+            second_per_grid_ts= all_second_gird if  len(all_second_gird) > 0 else None,
+           
+        )
+        data_dict["position_ids"] = position_ids
+        
+        return data_dict
+
+
+def pad_and_cat(tensor_list):
+    max_length = max(tensor.shape[2] for tensor in tensor_list)
+
+    padded_tensors = []
+    for tensor in tensor_list:
+        pad_length = max_length - tensor.shape[2]
+        padded_tensor = torch.nn.functional.pad(tensor, (0, pad_length), "constant", 1)
+        padded_tensors.append(padded_tensor)
+
+    stacked_tensor = torch.cat(padded_tensors, dim=1)
+
+    return stacked_tensor
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels, position_ids = tuple(
+            [instance[key] for instance in instances]
+            for key in ("input_ids", "labels", "position_ids")
+        )
+        input_ids = [ids.squeeze(0) for ids in input_ids]
+        labels = [ids.squeeze(0) for ids in labels]
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX
+        )
+        position_ids = pad_and_cat(position_ids)
+        input_ids = input_ids[:, : self.tokenizer.model_max_length]
+        labels = labels[:, : self.tokenizer.model_max_length]
+        position_ids = position_ids[:, :, : self.tokenizer.model_max_length]
+        batch = dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+        images = list(
+            instance["pixel_values"]
+            for instance in instances
+            if "pixel_values" in instance
+        )
+        videos = list(
+            instance["pixel_values_videos"]
+            for instance in instances
+            if "pixel_values_videos" in instance
+        )
+
+        # Align with batch per sample; use None for missing keys or placeholder samples to avoid image_files[i] being misaligned as None in forward if the length is shorter than batch.
+        images_files = [
+            instance["image_files"] if "image_files" in instance else None
+            for instance in instances
+        ]
+
+        if len(images) != 0:
+            concat_images = torch.cat([image for image in images], dim=0)
+            grid_thw = [
+                instance["image_grid_thw"]
+                for instance in instances
+                if "image_grid_thw" in instance
+            ]
+            grid_thw = torch.cat(grid_thw, dim=0)
+        else:
+            concat_images = None
+            grid_thw = None
+
+        if len(videos) != 0:
+            concat_videos = torch.cat([video for video in videos], dim=0)
+            video_grid_thw = [
+                instance["video_grid_thw"]
+                for instance in instances
+                if "video_grid_thw" in instance
+            ]
+            video_grid_thw = torch.cat(video_grid_thw, dim=0)
+        else:
+            concat_videos = None
+            video_grid_thw = None
+
+        batch["pixel_values"] = concat_images
+        batch["image_grid_thw"] = grid_thw
+        batch["pixel_values_videos"] = concat_videos
+        batch["video_grid_thw"] = video_grid_thw
+        batch["position_ids"] = position_ids
+        batch["image_files"] = images_files
+        return batch
+
+# @dataclass
+# class DataCollatorForSupervisedDataset(object):
+#     """Collate examples for supervised fine-tuning; truncates to tokenizer.model_max_length (same as train_stage1_algin)."""
+
+#     tokenizer: transformers.PreTrainedTokenizer
+
+#     def __post_init__(self):
+#         self.pad_token_id = self.tokenizer.pad_token_id
+
+#     def __call__(self, examples):
+#         max_len = self.tokenizer.model_max_length
+#         batch_input_ids = []
+#         batch_label_ids = []
+#         batch_pixel_values = []
+#         batch_pixel_video_values = []
+#         batch_video_thw = []
+#         batch_image_thw = []
+#         batch_second_per_grid_ts = []
+#         batch_mm_token_type_ids = []
+
+#         batch_image_files = []
+        
+#         for example in examples:
+#             keys = example.keys()
+#             if "pixel_values_videos" in keys:
+#                 batch_pixel_video_values.append(example["pixel_values_videos"])
+#                 batch_video_thw.append(example["video_grid_thw"])
+#             elif "pixel_values" in keys:
+#                 batch_pixel_values.append(example["pixel_values"])
+#                 batch_image_thw.append(example["image_grid_thw"])
+            
+#             if "image_files" in keys:
+#                 batch_image_files.append(example["image_files"])
+            
+#             iid = example["input_ids"]
+#             lab = example["labels"]
+#             if iid.shape[0] > max_len:
+#                 iid = iid[:max_len]
+#                 lab = lab[:max_len]
+#             batch_input_ids.append(iid)
+#             batch_label_ids.append(lab)
+#             if "mm_token_type_ids" in keys:
+#                 mm = example["mm_token_type_ids"]
+#                 if mm.shape[0] > max_len:
+#                     mm = mm[:max_len]
+#                 batch_mm_token_type_ids.append(mm)
+
+#             if "second_per_grid_ts" in keys:
+#                 batch_second_per_grid_ts.extend(example["second_per_grid_ts"])
+        
+#         input_ids = pad_sequence(
+#             batch_input_ids, padding_side='right', padding_value=self.pad_token_id
+#         )
+
+#         labels = pad_sequence(batch_label_ids, padding_side='right', padding_value=IGNORE_INDEX)
+
+#         if input_ids.shape[1] > max_len:
+#             rank0_print(
+#                 f"Warning: input with length {input_ids.shape[1]} is longer than max length {max_len}. Truncating."
+#             )
+#         input_ids = input_ids[:, :max_len]
+#         labels = labels[:, :max_len]
+#         attention_mask = input_ids != self.pad_token_id
+
+#         data_dict = {
+#             'input_ids': input_ids,
+#             'labels': labels,
+#             'attention_mask': attention_mask,
+#         }
+
+#         if len(batch_pixel_values) > 0:
+#             pixel_values = torch.cat(batch_pixel_values, dim=0)
+#             image_thw = torch.cat(batch_image_thw, dim=0)
+#             data_dict["pixel_values"] = pixel_values
+#             data_dict["image_grid_thw"] = image_thw
+
+#         if len(batch_pixel_video_values) > 0:
+#             pixel_video_values = torch.cat(batch_pixel_video_values, dim=0)
+#             video_thw = torch.cat(batch_video_thw, dim=0)
+#             data_dict["pixel_values_videos"] = pixel_video_values
+#             data_dict["video_grid_thw"] = video_thw
+
+#         if len(batch_second_per_grid_ts) > 0:
+#             data_dict["second_per_grid_ts"] = batch_second_per_grid_ts
+            
+#         if len(batch_image_files) > 0:
+#             data_dict["image_files"] = batch_image_files
+#         return data_dict
+
+def replace_image_tokens(input_string, is_video=False):
+    if is_video:
+        pattern = r'\n?' + re.escape(LLAVA_VIDEO_TOKEN) + r'\n?'
+        replacement = VISION_START_TOKEN + DEFAULT_VIDEO_TOKEN + VISION_END_TOKEN
+    else:
+        pattern = r'\n?' + re.escape(LLAVA_IMAGE_TOKEN) + r'\n?'
+        replacement = VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN
+
+    return re.sub(pattern, replacement, input_string)
+
+def llava_to_openai(conversations, is_video=False):
+    role_mapping = {"human": "user", "gpt": "assistant"}
+
+    transformed_data = []
+    for conversation in conversations:
+        transformed_content = replace_image_tokens(conversation["value"], is_video=is_video)
+        transformed_entry = {
+            "role": role_mapping.get(conversation["from"], conversation["from"]),
+            "content": transformed_content,
+        }
+        transformed_data.append(transformed_entry)
+
+    return transformed_data
+
+def make_supervised_data_module(model_id, processor, data_args, anchor_model_id):
+    """Make dataset and collator for supervised fine-tuning."""
+    sft_dataset = SupervisedDataset(
+        data_path=data_args.data_path, processor=processor, data_args=data_args, model_id=model_id, anchor_model_id=anchor_model_id
+    )
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=processor.tokenizer)
+
+    return dict(train_dataset=sft_dataset,
+                eval_dataset=None,
+                data_collator=data_collator)

code/umm/runtime/models/blip3o/data/covt_data_var.py

@@ -0,0 +1,1039 @@
+import copy
+import os
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Sequence, List, Tuple, Any
+from collections.abc import Sequence
+import torch
+import transformers
+import ujson as json
+from torch.utils.data import Dataset
+from qwen_vl_utils import process_vision_info
+from PIL import Image
+from transformers import AutoImageProcessor
+import re
+import numpy as np
+import cv2
+from torchvision import transforms
+import random
+
+from models.blip3o.model.anchors.segment_anything import build_sam_vit_h, sam_model_registry, SamPredictor
+from models.blip3o.model.anchors.DepthAnything.depth_anything_v2.dpt import DepthAnythingV2
+from diffusers import AutoencoderKL
+from transformers import AutoModel, CLIPImageProcessor
+
+from models.blip3o.params import DataArguments
+from models.blip3o.constants import *
+
+
+def rank0_print(*args):
+    if os.environ.get("LOCAL_RANK", "0") in ("0", "-1"):
+        print(*args)
+
+
+def _anchor_placeholder_body(anchor_num: int, anchor_token) -> str:
+    """Body between <|anchor_start|> and <|anchor_end|> (single repeated token or tuple of distinct VIS_* strings)."""
+    if isinstance(anchor_token, (list, tuple)):
+        seq = list(anchor_token)
+        if len(seq) != anchor_num:
+            raise ValueError(
+                f"anchor_token sequence length {len(seq)} != anchor_num {anchor_num}"
+            )
+        return "".join(seq)
+    return anchor_token * anchor_num
+
+
+def get_rope_index_3(
+    spatial_merge_size: Optional[int] = 2,
+    input_ids: Optional[torch.LongTensor] = None,
+    image_grid_thw: Optional[torch.LongTensor] = None,
+    video_grid_thw: Optional[torch.LongTensor] = None,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids."""
+    if input_ids is not None:
+        while input_ids.dim() < 2:
+            input_ids = input_ids.unsqueeze(0)
+    if attention_mask is not None:
+        while attention_mask.dim() < 2:
+            attention_mask = attention_mask.unsqueeze(0)
+    # Since we use timestamps to seperate videos, like <t1> <vision_start> <frame1> <vision_end> <t2> <vision_start> <frame2> <vision_end>, the video_grid_thw should also be split
+    if video_grid_thw is not None:
+        video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+        video_grid_thw[:, 0] = 1
+
+    image_token_id = 151655
+    video_token_id = 151656
+    vision_start_token_id = 151652
+    mrope_position_deltas = []
+    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+        total_input_ids = input_ids
+        if attention_mask is None:
+            attention_mask = torch.ones_like(total_input_ids)
+        position_ids = torch.ones(
+            3,
+            input_ids.shape[0],
+            input_ids.shape[1],
+            dtype=input_ids.dtype,
+            device=input_ids.device,
+        )
+        image_index, video_index = 0, 0
+        attention_mask = attention_mask.to(total_input_ids.device)
+        for i, input_ids in enumerate(total_input_ids):
+            input_ids = input_ids[attention_mask[i] == 1]
+            image_nums, video_nums = 0, 0
+            vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+            vision_tokens = input_ids[vision_start_indices + 1]
+            image_nums = (vision_tokens == image_token_id).sum()
+            video_nums = (vision_tokens == video_token_id).sum()
+            input_tokens = input_ids.tolist()
+            llm_pos_ids_list: list = []
+            st = 0
+            remain_images, remain_videos = image_nums, video_nums
+            for _ in range(image_nums + video_nums):
+                if image_token_id in input_tokens and remain_images > 0:
+                    ed_image = input_tokens.index(image_token_id, st)
+                else:
+                    ed_image = len(input_tokens) + 1
+                if video_token_id in input_tokens and remain_videos > 0:
+                    ed_video = input_tokens.index(video_token_id, st)
+                else:
+                    ed_video = len(input_tokens) + 1
+                if ed_image < ed_video:
+                    t, h, w = (
+                        image_grid_thw[image_index][0],
+                        image_grid_thw[image_index][1],
+                        image_grid_thw[image_index][2],
+                    )
+                    image_index += 1
+                    remain_images -= 1
+                    ed = ed_image
+
+                else:
+                    t, h, w = (
+                        video_grid_thw[video_index][0],
+                        video_grid_thw[video_index][1],
+                        video_grid_thw[video_index][2],
+                    )
+                    video_index += 1
+                    remain_videos -= 1
+                    ed = ed_video
+                llm_grid_t, llm_grid_h, llm_grid_w = (
+                    t.item(),
+                    h.item() // spatial_merge_size,
+                    w.item() // spatial_merge_size,
+                )
+                text_len = ed - st
+
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos)
+                t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+                llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+            if st < len(input_tokens):
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                text_len = len(input_tokens) - st
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+            position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+            mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+        return position_ids, mrope_position_deltas
+    else:
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+            max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+            mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+        else:
+            position_ids = (
+                torch.arange(input_ids.shape[1], device=input_ids.device)
+                .view(1, 1, -1)
+                .expand(3, input_ids.shape[0], -1)
+            )
+            mrope_position_deltas = torch.zeros(
+                [input_ids.shape[0], 1],
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+
+        return position_ids, mrope_position_deltas
+
+
+def truncate_sequence(input_ids, labels, max_length, eos_token_id):
+    if input_ids.size(0) > max_length:
+        input_ids = input_ids[:max_length-1]
+        labels = labels[:max_length-1]
+
+    if eos_token_id is not None:
+        input_ids = torch.cat([input_ids, torch.tensor([eos_token_id])])
+        labels = torch.cat([labels, torch.tensor([eos_token_id])])
+
+    return input_ids, labels
+
+def pad_sequence(sequences, padding_side='right', padding_value=0):
+    """
+    Pad a list of sequences to the same length.
+    sequences: list of tensors in [seq_len, *] shape
+    """
+    assert padding_side in ['right', 'left']
+    max_size = sequences[0].size()
+    trailing_dims = max_size[1:]
+    max_len = max(len(seq) for seq in sequences)
+    batch_size = len(sequences)
+    output = sequences[0].new_full((batch_size, max_len) + trailing_dims, padding_value)
+    for i, seq in enumerate(sequences):
+        length = seq.size(0)
+        if padding_side == 'right':
+            output.data[i, :length] = seq
+        else:
+            output.data[i, -length:] = seq
+    return output
+
+def get_image_info(image_path, min_pixel, max_pixel, width, height):
+    # Using this because of process_vision_info function
+    # Need to fix this in the future
+
+    content = {
+        "type": "image", 
+        "image": image_path,
+        "min_pixel": min_pixel,
+        "max_pixel": max_pixel
+    }
+
+    if width is not None and height is not None:
+        content["resized_width"] = width
+        content["resized_height"] = height
+
+    messages = [
+        {"role": "user", 
+         "content": [content]
+        }
+    ]
+
+    image_input, _ = process_vision_info(messages)
+
+    return image_input[0]
+
+def get_video_info(video_path, min_pixels, max_pixels, fps):
+    # Using this because of process_vision_info function
+    # Need to fix this in the future
+
+    messages = [
+        {"role": "user", 
+         "content": [
+             {
+                "type": "video", 
+                "video": video_path,
+                "min_pixels": min_pixels,
+                "max_pixels": max_pixels,
+                "fps": fps
+            }
+            ]
+        }
+    ]
+
+    _, video_input, video_kwargs = process_vision_info(messages, return_video_kwargs=True)
+
+    return video_input[0], video_kwargs
+
+def add_anchor_pad(user_input, anchor_nums, anchor_tokens):
+    # add anchor pad after VISION_END_TOKEN or ANCHOR_END_TOKEN
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + _anchor_placeholder_body(anchor_num, anchor_token) + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    anchor_pads = "".join(anchor_pads)
+    if VISION_END_TOKEN in user_input:
+        user_input = user_input.replace(VISION_END_TOKEN, VISION_END_TOKEN + anchor_pads)
+    return user_input
+
+def add_cot_anchor_pad_in_user_input(user_input, anchor_nums, anchor_tokens,anchor_names):
+    if len(anchor_nums) == 0:
+        return user_input
+    
+    anchor_pads = []
+    for anchor_num, anchor_token, _ in zip(anchor_nums, anchor_tokens, anchor_names):
+        anchor_pad = ANCHOR_START_TOKEN + _anchor_placeholder_body(anchor_num, anchor_token) + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_pad = ""
+    if len(anchor_pads) == 1:
+        CoT_pad = f"The {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
+            if i == 0:
+                CoT_pad += f"The {anchor_name} of the image is {anchor_pad}, "
+            elif i == len(anchor_names) - 1:
+                CoT_pad += f"and the {anchor_name} of the image is {anchor_pad}. "
+            else:
+                CoT_pad += f"the {anchor_name} of the image is {anchor_pad}, "
+    user_input = CoT_pad + user_input
+    return user_input
+
+def get_cot_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + _anchor_placeholder_body(anchor_num, anchor_token) + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_start = "Because "
+    if len(anchor_names) == 1:
+        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
+            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
+            if anchor_name == anchor_names[-2]:
+                CoT_start += ", and "
+            elif anchor_name == anchor_names[-1]:
+                CoT_start += ". "
+            else:
+                CoT_start += ", "
+    response = CoT_start + response
+    return response
+
+
+COT_TEMPLATES = [
+    {
+        "name": "basic_causal",
+        "single": "Because the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Because the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "observational",
+        "single": "I can observe that the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "I can observe that the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "analytical",
+        "single": "After analyzing the image, the {anchor_name} is {anchor_pad}. ",
+        "multiple": "After analyzing the image, the {anchor_name} is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "descriptive",
+        "single": "The image shows that the {anchor_name} is {anchor_pad}. ",
+        "multiple": "The image shows that the {anchor_name} is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "conditional",
+        "single": "Given that the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Given that the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    },
+    
+    {
+        "name": "evidence_based",
+        "single": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}. ",
+        "multiple": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}{connector}",
+        "connectors": {
+            "middle": ", ",
+            "second_last": ", and ",
+            "last": ". "
+        }
+    }
+]
+
+def get_random_cot_template():
+    return random.choice(COT_TEMPLATES)
+
+def apply_cot_template(template, anchor_names, anchor_pads):
+    if len(anchor_names) == 1:
+        return template["single"].format(
+            anchor_name=anchor_names[0], 
+            anchor_pad=anchor_pads[0]
+        )
+    else:
+        result = ""
+        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
+            if i == len(anchor_names) - 1: 
+                connector = template["connectors"]["last"]
+            elif i == len(anchor_names) - 2: 
+                connector = template["connectors"]["second_last"]
+            else: 
+                connector = template["connectors"]["middle"]
+            
+            result += template["multiple"].format(
+                anchor_name=anchor_name,
+                anchor_pad=anchor_pad,
+                connector=connector
+            )
+        return result
+
+def get_templates_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + _anchor_placeholder_body(anchor_num, anchor_token) + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    
+    template = get_random_cot_template()
+    
+    cot_text = apply_cot_template(template, anchor_names, anchor_pads)
+    
+    response = "<think>" + cot_text + "</think>" + "<answer>" + response + "</answer>"
+    return response
+        
+def get_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + _anchor_placeholder_body(anchor_num, anchor_token) + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_start = "<think> Because "
+    if len(anchor_names) == 1:
+        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
+            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
+            if anchor_name == anchor_names[-2]:
+                CoT_start += ", and "
+            elif anchor_name == anchor_names[-1]:
+                CoT_start += ". "
+            else:
+                CoT_start += ", "
+    response = CoT_start + " </think>\n" + "<answer> " + response + " </answer>"
+    return response
+    
+def get_feature_data(user_input, gpt_response, anchor_nums, anchor_tokens, anchor_names):
+    anchor_pads = []
+    for anchor_num, anchor_token, anchor_name in zip(anchor_nums, anchor_tokens, anchor_names):
+        anchor_pad = ANCHOR_START_TOKEN + _anchor_placeholder_body(anchor_num, anchor_token) + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    anchor_name = ", ".join(anchor_names)
+    anchor_pads = "".join(anchor_pads)
+    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}What is the {anchor_name} of the image?\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+    gpt_response = f"{anchor_pads}\n{DEFAULT_IM_END_TOKEN}\n"
+    return user_input, gpt_response
+
+def replace_pad_with_anchor_tokens(gpt_response):
+    token_dict = {
+        "<segmentation>": SAM_PAD_TOKEN * 8,
+        "<depth>": "".join(VIS_DEPTH_ANCHOR_TOKENS),
+        "<dino>": "".join(VIS_DINO_ANCHOR_TOKENS),
+        "<pidinet>": "".join(VIS_PIDINET_ANCHOR_TOKENS),
+        "<siglip>": "".join(VIS_SIGLIP_ANCHOR_TOKENS),
+        "<metaclip>": METACLIP_PAD_TOKEN * 4,
+    }
+    for token, anchor_token in token_dict.items():
+        gpt_response = gpt_response.replace(token, anchor_token)
+    return gpt_response
+    
+def get_token_num(anchor_model_id):
+    token_nums = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            token_nums.append(8)
+        elif anchor_model == "dino":
+            token_nums.append(4)
+        elif anchor_model == "depth":
+            token_nums.append(4)
+        elif anchor_model == "InternViT":
+            token_nums.append(4)
+        elif anchor_model == "pidinet":
+            token_nums.append(4)
+        elif anchor_model == "siglip":
+            token_nums.append(4)
+        elif anchor_model == "metaclip":
+            token_nums.append(4)
+    return token_nums
+
+def get_anchor_token(anchor_model_id):
+    anchor_tokens = []
+    for anchor_model in anchor_model_id:
+        # if anchor_model == "sam":
+        #     anchor_tokens.append((SAM_PAD_TOKEN,) * 8)
+        if anchor_model == "dino":
+            anchor_tokens.append(tuple(VIS_DINO_ANCHOR_TOKENS))
+        elif anchor_model == "depth":
+            anchor_tokens.append(tuple(VIS_DEPTH_ANCHOR_TOKENS))
+        # elif anchor_model == "InternViT":
+        #     anchor_tokens.append((INTERN_PAD_TOKEN,) * 4)
+        elif anchor_model == "pidinet":
+            anchor_tokens.append(tuple(VIS_PIDINET_ANCHOR_TOKENS))
+        elif anchor_model == "siglip":
+            anchor_tokens.append(tuple(VIS_SIGLIP_ANCHOR_TOKENS))
+        # elif anchor_model == "metaclip":
+        #     anchor_tokens.append((METACLIP_PAD_TOKEN,) * 4)
+    return anchor_tokens
+
+def get_anchor_task_name(anchor_model_id):
+    anchor_task_names = []
+    for anchor_model in anchor_model_id:
+        if anchor_model == "sam":
+            anchor_task_names.append("segmentation")
+        elif anchor_model == "dino":
+            anchor_task_names.append("perception feature")
+        elif anchor_model == "depth":
+            anchor_task_names.append("depth map")
+        elif anchor_model == "InternViT":
+            anchor_task_names.append("caption")
+        elif anchor_model == "pidinet":
+            anchor_task_names.append("edge map")
+        elif anchor_model == "siglip":
+            anchor_task_names.append("semantic feature")
+        elif anchor_model == "metaclip":
+            anchor_task_names.append("metaclip feature")
+    return anchor_task_names
+    
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(
+        self,
+        data_path: str | list,
+        processor: transformers.ProcessorMixin,
+        data_args: DataArguments,
+        model_id,
+        padding=True,
+        shuffle=True,
+        random_seed=42,
+        anchor_model_id=None,
+    ):
+        super(SupervisedDataset, self).__init__()
+        if isinstance(data_path, str):
+            list_data_dict = json.load(open(data_path, "r"))
+        else:
+            list_data_dict = data_path
+
+        self.model_id = model_id
+        self.processor = processor
+        self.list_data_dict = list_data_dict
+        self.data_args = data_args
+        self.padding = padding
+        self.image_min_pixel = data_args.image_min_pixels
+        self.image_max_pixel = data_args.image_max_pixels
+        self.image_resized_w = data_args.image_resized_width
+        self.image_resized_h = data_args.image_resized_height
+        self.video_min_pixel = data_args.video_min_pixels
+        self.video_max_pixel = data_args.video_max_pixels
+        self.fps = data_args.fps
+        self.merge_size = getattr(processor.image_processor, "merge_size", 2)
+        self.get_rope_index = get_rope_index_3
+        self.anchor_model_id = anchor_model_id
+        self.anchor_token_nums = get_token_num(anchor_model_id)
+        self.anchor_tokens = get_anchor_token(anchor_model_id)
+        self.anchor_task_names = get_anchor_task_name(anchor_model_id)
+        
+        self.cur_step = 0
+        self.stage_0_step = data_args.stage_0_step
+        self.stage_1_step = data_args.stage_1_step
+        self.stage_2_step = data_args.stage_2_step
+        
+        # for shuffle
+        self.rng = np.random.default_rng(seed=random_seed)
+        
+        if shuffle:
+            self.rng.shuffle(self.list_data_dict)
+    
+    def set_cur_step(self, step: int):
+        self.cur_step = step
+        print(f"[Dataset] cur_step has been set to {step}")
+        
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        
+        # import ipdb; ipdb.set_trace()
+        
+        sources = self.list_data_dict[i]
+        
+        is_video = False
+        image_files = None
+
+        processor = self.processor
+        if "image" in sources:
+            videos = None
+            grid_key = "image_grid_thw"
+            pixel_key = "pixel_values"
+            
+            image_files = sources["image"]
+            image_folder = self.data_args.image_folder
+
+            if isinstance(image_files, str):
+                image_files = Image.open(image_files).convert("RGB")
+                image_files = [image_files]
+            else:
+                image_files = [Image.open(image_file).convert("RGB") for image_file in image_files]
+            images = []
+            
+            for image_file in image_files:
+                # if not os.path.exists(image_file):
+                #     if not image_file.startswith("http"):
+                #         image_file = os.path.join(image_folder, image_file)
+                #         images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+                # else:
+                images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+
+        elif "video" in sources:
+            is_video = True
+            images=None
+            grid_key = "video_grid_thw"
+            pixel_key = "pixel_values_videos"
+
+            video_files = sources["video"]
+            video_folder = self.data_args.image_folder
+
+            if isinstance(video_files, str):
+                video_files = [video_files]
+
+            videos = []
+            for video_file in video_files:
+                if not os.path.exists(video_file):
+                    if not video_file.startswith("http"):
+                        video_file = os.path.join(video_folder, video_file)
+                video_input, video_kwargs = get_video_info(video_file, self.video_min_pixel, self.video_max_pixel, self.data_args.fps)
+                videos.append(video_input)
+        else:
+            grid_key = None
+            pixel_key = None
+            images=None
+            videos=None
+           
+        if images is None:
+            print("No image or video found in the data.")
+            images = []
+            # Create a black image as a placeholder
+            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+
+        elif len(images) == 0:
+            print("No image or video found in the data.")
+            # Create a black image as a placeholder
+            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
+        
+        if videos is not None:
+            
+            # import ipdb; ipdb.set_trace()
+            pass
+            
+        sources = copy.deepcopy(llava_to_openai(sources['conversations'], is_video=is_video))
+
+        all_input_ids = []
+        all_labels = []
+        all_pixel_values = []
+        all_image_grid_thw = []
+        all_second_gird = []
+        
+        # all_dino_encoded_values = []
+
+        # Qwen2-VL uses a default system message so I've added this.
+        if len(SYSTEM_MESSAGE) > 0:
+            system_message = f"{DEFAULT_IM_START_TOKEN}system\n{SYSTEM_MESSAGE}\n{DEFAULT_IM_END_TOKEN}\n"
+            system_message_input_ids = processor.tokenizer(system_message, add_special_tokens=False, return_tensors='pt')['input_ids']
+            system_labels = torch.full_like(system_message_input_ids, IGNORE_INDEX) 
+            
+            all_input_ids.append(system_message_input_ids.squeeze(0))
+            all_labels.append(system_labels.squeeze(0))
+        
+        # import ipdb; ipdb.set_trace()
+        for _, j in enumerate(range(0, len(sources), 2)):
+            
+            if j >= 2:
+                break
+            
+            user_input = sources[j]
+            gpt_response = sources[j + 1]
+                        
+            if (DEFAULT_IMAGE_TOKEN not in user_input['content']) and (DEFAULT_VIDEO_TOKEN not in user_input['content']) and (LLAVA_IMAGE_TOKEN in user_input['content']):
+                user_input = f"{DEFAULT_IM_START_TOKEN}{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                # user_input = add_cot_anchor_pad_in_user_input(user_input, self.anchor_token_nums, self.anchor_tokens,self.anchor_task_names)
+                user_input = add_anchor_pad(user_input, self.anchor_token_nums, self.anchor_tokens)
+                gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                raise ValueError('Every man is a poet when he is in love')
+            else:
+                if self.cur_step < self.stage_0_step:
+                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    # user_input = add_cot_anchor_pad_in_user_input(user_input, self.anchor_token_nums, self.anchor_tokens,self.anchor_task_names)
+                    user_input = add_anchor_pad(user_input, self.anchor_token_nums, self.anchor_tokens)
+                    gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                elif self.cur_step < self.stage_1_step:
+                    user_input, gpt_response = get_feature_data(user_input, gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
+                elif self.cur_step < self.stage_2_step:
+                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    gpt_response = f"{gpt_response['content']}"
+                    if DEFAULT_IMAGE_TOKEN in user_input:
+                        gpt_response = get_comt_data_in_response(gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
+                    gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # print(f"\033[92m gpt_response: {gpt_response}\033[0m")
+                else:
+                    # user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    # gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # gpt_response = replace_pad_with_anchor_tokens(gpt_response)
+                    import random
+                    xxx = random.randint(0, 5)
+                    if xxx == 0:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    else:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}"
+                        if DEFAULT_IMAGE_TOKEN in user_input:
+                            # INSERT_YOUR_CODE
+                            total = len(self.anchor_tokens)
+                            if total == 0:
+                                selected_anchor_token_nums = []
+                                selected_anchor_tokens = []
+                                selected_anchor_task_names = []
+                            else:
+                                x = random.randint(1, total)
+                                idxs = sorted(random.sample(range(total), x)) if x > 0 else []
+                                selected_anchor_token_nums = [self.anchor_token_nums[i] for i in idxs]
+                                selected_anchor_tokens = [self.anchor_tokens[i] for i in idxs]
+                                selected_anchor_task_names = [self.anchor_task_names[i] for i in idxs]
+                            gpt_response = get_comt_data_in_response(gpt_response, selected_anchor_token_nums, selected_anchor_tokens, selected_anchor_task_names)
+                        gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+            
+            # print(f'the user_input is {user_input}')
+            # print(f'the gpt_response is {gpt_response}')
+                    
+
+            # print("-----------------")
+            # print(user_input, gpt_response)
+            # print("-----------------")
+
+            # import ipdb; ipdb.set_trace()
+            if DEFAULT_IMAGE_TOKEN in user_input:
+                inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
+                prompt_input_ids = inputs['input_ids']
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+                
+                # del dino_val
+                torch.cuda.empty_cache()
+                
+            elif DEFAULT_VIDEO_TOKEN in user_input:
+                if "Qwen3" in self.model_id:
+                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt', **video_kwargs)
+                    all_second_gird.extend(inputs["second_per_grid_ts"])
+                else:
+                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
+                prompt_input_ids = inputs['input_ids']
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+
+            else:
+                prompt_input_ids = processor.tokenizer(user_input, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+        
+
+            response_input_ids = processor.tokenizer(gpt_response, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+
+            input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
+            labels = torch.cat(
+                [
+                    torch.tensor([IGNORE_INDEX] * len(prompt_input_ids[0])),  
+                    response_input_ids.squeeze(0),
+                ],
+                dim=0,
+            )
+            all_input_ids.append(input_ids)
+            all_labels.append(labels)
+        
+        # There is no need for eos or bos tokens in the input_ids
+        # Qwen2-VL does not use them
+        input_ids = torch.cat(all_input_ids, dim=0).to(torch.long)
+        labels = torch.cat(all_labels, dim=0).to(torch.long)
+
+        # eos_token_id = processor.tokenizer.convert_tokens_to_ids(DEFAULT_IM_END_TOKEN)
+        # input_ids, labels = truncate_sequence(input_ids, labels, self.max_length, eos_token_id)
+
+        attention_mask = (input_ids > -1000000).to(torch.long)
+        task_type = "und_image"
+
+        x_ids, x_mask = input_ids, attention_mask
+        while x_ids.dim() < 2:
+            x_ids = x_ids.unsqueeze(0)
+    
+        input_ids_2d= x_ids
+        data_dict = dict(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            task_types=task_type,
+        )
+
+        if pixel_key and grid_key:
+            pixel_values = torch.cat(all_pixel_values, dim=0)
+            image_thw = torch.cat(all_image_grid_thw, dim=0)
+            
+            data_dict[pixel_key] = pixel_values
+            data_dict[grid_key] = image_thw
+            data_dict["image_files"] = image_files
+
+        if len(all_second_gird) > 0:
+            second_gird = all_second_gird
+            data_dict["second_per_grid_ts"] = second_gird
+
+        position_ids, _ = self.get_rope_index(
+            self.merge_size,
+            input_ids_2d,
+            image_grid_thw=torch.cat(all_image_grid_thw, dim=0) if images else None,
+            video_grid_thw=(
+                torch.cat(all_image_grid_thw, dim=0) if videos else None
+            ),
+            second_per_grid_ts= all_second_gird if  len(all_second_gird) > 0 else None,
+           
+        )
+        data_dict["position_ids"] = position_ids
+        
+        self.cur_step += 1
+        return data_dict
+
+
+def pad_and_cat(tensor_list):
+    max_length = max(tensor.shape[2] for tensor in tensor_list)
+
+    padded_tensors = []
+    for tensor in tensor_list:
+        pad_length = max_length - tensor.shape[2]
+        padded_tensor = torch.nn.functional.pad(tensor, (0, pad_length), "constant", 1)
+        padded_tensors.append(padded_tensor)
+
+    stacked_tensor = torch.cat(padded_tensors, dim=1)
+
+    return stacked_tensor
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels, position_ids = tuple(
+            [instance[key] for instance in instances]
+            for key in ("input_ids", "labels", "position_ids")
+        )
+        input_ids = [ids.squeeze(0) for ids in input_ids]
+        labels = [ids.squeeze(0) for ids in labels]
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX
+        )
+        position_ids = pad_and_cat(position_ids)
+        input_ids = input_ids[:, : self.tokenizer.model_max_length]
+        labels = labels[:, : self.tokenizer.model_max_length]
+        position_ids = position_ids[:, :, : self.tokenizer.model_max_length]
+        batch = dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+        images = list(
+            instance["pixel_values"]
+            for instance in instances
+            if "pixel_values" in instance
+        )
+        videos = list(
+            instance["pixel_values_videos"]
+            for instance in instances
+            if "pixel_values_videos" in instance
+        )
+
+        # Align with batch per sample; use None for missing keys or placeholder samples to avoid image_files[i] being misaligned as None in forward if the length is shorter than batch.
+        images_files = [
+            instance["image_files"] if "image_files" in instance else None
+            for instance in instances
+        ]
+
+        if len(images) != 0:
+            concat_images = torch.cat([image for image in images], dim=0)
+            grid_thw = [
+                instance["image_grid_thw"]
+                for instance in instances
+                if "image_grid_thw" in instance
+            ]
+            grid_thw = torch.cat(grid_thw, dim=0)
+        else:
+            concat_images = None
+            grid_thw = None
+
+        if len(videos) != 0:
+            concat_videos = torch.cat([video for video in videos], dim=0)
+            video_grid_thw = [
+                instance["video_grid_thw"]
+                for instance in instances
+                if "video_grid_thw" in instance
+            ]
+            video_grid_thw = torch.cat(video_grid_thw, dim=0)
+        else:
+            concat_videos = None
+            video_grid_thw = None
+
+        batch["pixel_values"] = concat_images
+        batch["image_grid_thw"] = grid_thw
+        batch["pixel_values_videos"] = concat_videos
+        batch["video_grid_thw"] = video_grid_thw
+        batch["position_ids"] = position_ids
+        batch["image_files"] = images_files
+        return batch
+
+# @dataclass
+# class DataCollatorForSupervisedDataset(object):
+#     """Collate examples for supervised fine-tuning; truncates to tokenizer.model_max_length (same as train_stage1_algin)."""
+
+#     tokenizer: transformers.PreTrainedTokenizer
+
+#     def __post_init__(self):
+#         self.pad_token_id = self.tokenizer.pad_token_id
+
+#     def __call__(self, examples):
+#         max_len = self.tokenizer.model_max_length
+#         batch_input_ids = []
+#         batch_label_ids = []
+#         batch_pixel_values = []
+#         batch_pixel_video_values = []
+#         batch_video_thw = []
+#         batch_image_thw = []
+#         batch_second_per_grid_ts = []
+#         batch_mm_token_type_ids = []
+
+#         batch_image_files = []
+        
+#         for example in examples:
+#             keys = example.keys()
+#             if "pixel_values_videos" in keys:
+#                 batch_pixel_video_values.append(example["pixel_values_videos"])
+#                 batch_video_thw.append(example["video_grid_thw"])
+#             elif "pixel_values" in keys:
+#                 batch_pixel_values.append(example["pixel_values"])
+#                 batch_image_thw.append(example["image_grid_thw"])
+            
+#             if "image_files" in keys:
+#                 batch_image_files.append(example["image_files"])
+            
+#             iid = example["input_ids"]
+#             lab = example["labels"]
+#             if iid.shape[0] > max_len:
+#                 iid = iid[:max_len]
+#                 lab = lab[:max_len]
+#             batch_input_ids.append(iid)
+#             batch_label_ids.append(lab)
+#             if "mm_token_type_ids" in keys:
+#                 mm = example["mm_token_type_ids"]
+#                 if mm.shape[0] > max_len:
+#                     mm = mm[:max_len]
+#                 batch_mm_token_type_ids.append(mm)
+
+#             if "second_per_grid_ts" in keys:
+#                 batch_second_per_grid_ts.extend(example["second_per_grid_ts"])
+        
+#         input_ids = pad_sequence(
+#             batch_input_ids, padding_side='right', padding_value=self.pad_token_id
+#         )
+
+#         labels = pad_sequence(batch_label_ids, padding_side='right', padding_value=IGNORE_INDEX)
+
+#         if input_ids.shape[1] > max_len:
+#             rank0_print(
+#                 f"Warning: input with length {input_ids.shape[1]} is longer than max length {max_len}. Truncating."
+#             )
+#         input_ids = input_ids[:, :max_len]
+#         labels = labels[:, :max_len]
+#         attention_mask = input_ids != self.pad_token_id
+
+#         data_dict = {
+#             'input_ids': input_ids,
+#             'labels': labels,
+#             'attention_mask': attention_mask,
+#         }
+
+#         if len(batch_pixel_values) > 0:
+#             pixel_values = torch.cat(batch_pixel_values, dim=0)
+#             image_thw = torch.cat(batch_image_thw, dim=0)
+#             data_dict["pixel_values"] = pixel_values
+#             data_dict["image_grid_thw"] = image_thw
+
+#         if len(batch_pixel_video_values) > 0:
+#             pixel_video_values = torch.cat(batch_pixel_video_values, dim=0)
+#             video_thw = torch.cat(batch_video_thw, dim=0)
+#             data_dict["pixel_values_videos"] = pixel_video_values
+#             data_dict["video_grid_thw"] = video_thw
+
+#         if len(batch_second_per_grid_ts) > 0:
+#             data_dict["second_per_grid_ts"] = batch_second_per_grid_ts
+            
+#         if len(batch_image_files) > 0:
+#             data_dict["image_files"] = batch_image_files
+#         return data_dict
+
+def replace_image_tokens(input_string, is_video=False):
+    if is_video:
+        pattern = r'\n?' + re.escape(LLAVA_VIDEO_TOKEN) + r'\n?'
+        replacement = VISION_START_TOKEN + DEFAULT_VIDEO_TOKEN + VISION_END_TOKEN
+    else:
+        pattern = r'\n?' + re.escape(LLAVA_IMAGE_TOKEN) + r'\n?'
+        replacement = VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN
+
+    return re.sub(pattern, replacement, input_string)
+
+def llava_to_openai(conversations, is_video=False):
+    role_mapping = {"human": "user", "gpt": "assistant"}
+
+    transformed_data = []
+    for conversation in conversations:
+        transformed_content = replace_image_tokens(conversation["value"], is_video=is_video)
+        transformed_entry = {
+            "role": role_mapping.get(conversation["from"], conversation["from"]),
+            "content": transformed_content,
+        }
+        transformed_data.append(transformed_entry)
+
+    return transformed_data
+
+def make_supervised_data_module(model_id, processor, data_args, anchor_model_id):
+    """Make dataset and collator for supervised fine-tuning."""
+    sft_dataset = SupervisedDataset(
+        data_path=data_args.data_path, processor=processor, data_args=data_args, model_id=model_id, anchor_model_id=anchor_model_id
+    )
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=processor.tokenizer)
+
+    return dict(train_dataset=sft_dataset,
+                eval_dataset=None,
+                data_collator=data_collator)

code/umm/runtime/models/blip3o/data/mix_und_t2i_data.py

@@ -0,0 +1,274 @@
+"""
+Mixed per-batch sampling: understanding (SupervisedDataset / covt_data) + T2I (LazySupervisedMixDataset).
+
+Each batch contains a fixed number of UND samples and T2I samples so both tasks train together.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, Iterator, List, Optional, Sequence
+
+import torch
+from torch.utils.data import ConcatDataset, Sampler
+
+import transformers
+
+from models.blip3o.constants import IGNORE_INDEX
+from models.blip3o.data.covt_data import SupervisedDataset
+from models.blip3o.data.t2i_covt_data import LazySupervisedMixDataset
+
+
+class MixedUndT2IConcatDataset(ConcatDataset):
+    """ConcatDataset with `set_cur_step` forwarded to child datasets that support it."""
+
+    def set_cur_step(self, step: int) -> None:
+        for d in self.datasets:
+            if hasattr(d, "set_cur_step"):
+                d.set_cur_step(step)
+
+
+class MixedUndT2IBatchSampler(Sampler[List[int]]):
+    """
+    Yields batches of indices into ConcatDataset: first part = UND, second part = T2I.
+    Each batch has exactly `num_und` indices from [0, und_len) and `num_t2i` from [und_len, und_len+t2i_len).
+
+    If ``mix_start_global_step > 0``, batches before that HF trainer ``global_step`` use only UND
+    indices (batch size ``num_und + num_t2i`` from the UND pool); from ``global_step >= mix_start_global_step``
+    on, batches alternate UND + T2I as usual. Requires ``global_step_ref`` updated each step (see
+    ``MixedUndT2IStepSyncCallback``); safe with ``dataloader_num_workers=0``.
+    """
+
+    def __init__(
+        self,
+        und_len: int,
+        t2i_len: int,
+        num_und: int,
+        num_t2i: int,
+        seed: int = 42,
+        rank: int = 0,
+        world_size: int = 1,
+        mix_start_global_step: int = 0,
+        global_step_ref: Optional[List[int]] = None,
+    ):
+        if und_len <= 0 or t2i_len <= 0:
+            raise ValueError(f"und_len and t2i_len must be positive, got {und_len}, {t2i_len}")
+        if num_und <= 0 or num_t2i <= 0:
+            raise ValueError(f"num_und and num_t2i must be positive, got {num_und}, {num_t2i}")
+        if num_und > und_len or num_t2i > t2i_len:
+            raise ValueError(
+                f"Batch larger than dataset: num_und={num_und} > und_len={und_len} or "
+                f"num_t2i={num_t2i} > t2i_len={t2i_len}"
+            )
+        if mix_start_global_step < 0:
+            raise ValueError(f"mix_start_global_step must be >= 0, got {mix_start_global_step}")
+
+        self.und_len = und_len
+        self.t2i_len = t2i_len
+        self.t2i_offset = und_len
+        self.num_und = num_und
+        self.num_t2i = num_t2i
+        self.seed = seed
+        self.rank = rank
+        self.world_size = world_size
+        self.epoch = 0
+        self.mix_start_global_step = mix_start_global_step
+        self.global_step_ref = global_step_ref if global_step_ref is not None else [0]
+
+        n_mixed = min(und_len // num_und, t2i_len // num_t2i)
+        und_bs = num_und + num_t2i
+        n_und_only = und_len // und_bs if und_bs > 0 else 0
+        self._num_batches = max(n_mixed, n_und_only)
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = int(epoch)
+
+    def __iter__(self) -> Iterator[List[int]]:
+        g = torch.Generator()
+        g.manual_seed(self.seed + self.epoch * 100_003 + self.rank * 17)
+        und_perm = torch.randperm(self.und_len, generator=g).tolist()
+        t2i_perm = torch.randperm(self.t2i_len, generator=g).tolist()
+
+        iu = 0
+        it = 0
+        batch_idx = 0
+
+        while True:
+            gs = self.global_step_ref[0]
+            if self.mix_start_global_step == 0 or gs >= self.mix_start_global_step:
+                if iu + self.num_und > self.und_len or it + self.num_t2i > self.t2i_len:
+                    break
+                batch = und_perm[iu : iu + self.num_und] + [
+                    self.t2i_offset + t2i_perm[it + k] for k in range(self.num_t2i)
+                ]
+                iu += self.num_und
+                it += self.num_t2i
+            else:
+                und_bs = self.num_und + self.num_t2i
+                if iu + und_bs > self.und_len:
+                    break
+                batch = und_perm[iu : iu + und_bs]
+                iu += und_bs
+
+            if batch_idx % self.world_size != self.rank:
+                batch_idx += 1
+                continue
+            yield list(batch)
+            batch_idx += 1
+
+    def __len__(self) -> int:
+        return math.ceil(self._num_batches / self.world_size)
+
+
+class DataCollatorForUndT2IMixed:
+    """Pad text/mm fields; concat vision patches for UND rows only; stack T2I gen_image + row indices."""
+
+    def __init__(self, pad_token_id: int):
+        self.pad_token_id = pad_token_id
+
+    def __call__(self, instances: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        from models.blip3o.data.covt_data import pad_sequence as covt_pad_sequence
+
+        batch_input_ids: List[torch.Tensor] = []
+        batch_labels: List[torch.Tensor] = []
+        batch_mm: List[torch.Tensor] = []
+        batch_pixel: List[torch.Tensor] = []
+        batch_image_thw: List[torch.Tensor] = []
+        batch_pixel_video: List[torch.Tensor] = []
+        batch_video_thw: List[torch.Tensor] = []
+        batch_second: List = []
+        batch_image_files: List[Any] = []
+        batch_gen_image: List[torch.Tensor] = []
+        t2i_row_indices: List[int] = []
+
+        for i, ex in enumerate(instances):
+            batch_input_ids.append(ex["input_ids"])
+            batch_labels.append(ex["labels"])
+            if "mm_token_type_ids" in ex:
+                batch_mm.append(ex["mm_token_type_ids"])
+
+            keys = ex.keys()
+            if "pixel_values_videos" in keys:
+                batch_pixel_video.append(ex["pixel_values_videos"])
+                batch_video_thw.append(ex["video_grid_thw"])
+            elif "pixel_values" in keys:
+                batch_pixel.append(ex["pixel_values"])
+                batch_image_thw.append(ex["image_grid_thw"])
+
+            if "image_files" in ex:
+                batch_image_files.append(ex["image_files"])
+            else:
+                batch_image_files.append(None)
+
+            if "second_per_grid_ts" in ex:
+                batch_second.extend(ex["second_per_grid_ts"])
+
+            tt = ex.get("task_types", "und_image")
+            if tt == "T2I" and "gen_image" in ex: 
+                gi = ex["gen_image"]
+                if gi.dim() == 3:
+                    gi = gi.unsqueeze(0)
+                batch_gen_image.append(gi)
+                t2i_row_indices.append(i)
+
+        input_ids = covt_pad_sequence(batch_input_ids, padding_side="right", padding_value=self.pad_token_id)
+        attention_mask = input_ids != self.pad_token_id
+        labels = covt_pad_sequence(batch_labels, padding_side="right", padding_value=IGNORE_INDEX)
+
+        out: Dict[str, Any] = {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": attention_mask,
+        }
+
+        if batch_mm:
+            out["mm_token_type_ids"] = covt_pad_sequence(batch_mm, padding_side="right", padding_value=0)
+
+        if batch_pixel:
+            out["pixel_values"] = torch.cat(batch_pixel, dim=0)
+            out["image_grid_thw"] = torch.cat(batch_image_thw, dim=0)
+
+        if batch_pixel_video:
+            out["pixel_values_videos"] = torch.cat(batch_pixel_video, dim=0)
+            out["video_grid_thw"] = torch.cat(batch_video_thw, dim=0)
+
+        if batch_second:
+            out["second_per_grid_ts"] = batch_second
+
+        out["image_files"] = batch_image_files
+
+        if batch_gen_image:
+            out["gen_image"] = torch.cat(batch_gen_image, dim=0)
+            out["t2i_row_indices"] = torch.tensor(t2i_row_indices, dtype=torch.long)
+        else:
+            out["gen_image"] = None
+            out["t2i_row_indices"] = None
+
+        return out
+
+
+def make_mixed_und_t2i_data_module(
+    *,
+    model_id: str,
+    processor: transformers.ProcessorMixin,
+    data_args: Any,
+    anchor_model_id: Optional[List[str]],
+    und_data_path: str,
+    t2i_data_path: str,
+    num_und_per_batch: int,
+    num_t2i_per_batch: int,
+    mix_start_global_step: int = 0,
+    seed: int = 42,
+    rank: int = 0,
+    world_size: int = 1,
+) -> Dict[str, Any]:
+    """
+    Build ConcatDataset( SupervisedDataset(und), LazySupervisedMixDataset(t2i) ) with a mixed BatchSampler.
+    """
+    und_ds = SupervisedDataset(
+        data_path=und_data_path,
+        processor=processor,
+        data_args=data_args,
+        model_id=model_id,
+        anchor_model_id=anchor_model_id,
+    )
+    t2i_ds = LazySupervisedMixDataset(
+        data_path=t2i_data_path,
+        processor=processor,
+        data_args=data_args,
+        model_id=model_id,
+        anchor_model_id=anchor_model_id,
+        random_seed=seed,
+    )
+
+    train_dataset = MixedUndT2IConcatDataset([und_ds, t2i_ds])
+    und_len = len(und_ds)
+    t2i_len = len(t2i_ds)
+
+    global_step_ref: List[int] = [0]
+    batch_sampler = MixedUndT2IBatchSampler(
+        und_len=und_len,
+        t2i_len=t2i_len,
+        num_und=num_und_per_batch,
+        num_t2i=num_t2i_per_batch,
+        seed=seed,
+        rank=rank,
+        world_size=world_size,
+        mix_start_global_step=mix_start_global_step,
+        global_step_ref=global_step_ref,
+    )
+
+    train_dataset._mixed_batch_sampler = batch_sampler
+    train_dataset._mixed_global_step_ref = global_step_ref
+    train_dataset._mixed_num_und = num_und_per_batch
+    train_dataset._mixed_num_t2i = num_t2i_per_batch
+
+    data_collator = DataCollatorForUndT2IMixed(pad_token_id=processor.tokenizer.pad_token_id)
+
+    return {
+        "train_dataset": train_dataset,
+        "eval_dataset": None,
+        "data_collator": data_collator,
+        "mixed_train_batch_sampler": batch_sampler,
+        "mixed_global_step_ref": global_step_ref,
+    }

code/umm/runtime/models/blip3o/data/t2I_covt_gen_data.py

@@ -0,0 +1,894 @@
+"""
+T2I / I2I generation dataset for CoVT + Qwen3-VL.
+
+- Text/mm fields and ``position_ids`` follow ``covt_data_var`` (``get_rope_index_3``).
+- CoVT anchor pads / staged CoT match ``covt_data_var`` helpers.
+- Keeps ``gen_image`` (and optional I2I ref fields) for the generation head.
+"""
+
+from __future__ import annotations
+
+import copy
+import io
+import os
+import random
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import transformers
+import ujson as json
+from datasets import Image, concatenate_datasets, load_dataset
+from PIL import Image as PILImage
+from PIL import ImageFile
+from torch.utils.data import Dataset
+from torchvision.transforms import v2
+
+from models.blip3o.constants import *
+from models.blip3o.data.covt_data_var import (
+    add_anchor_pad,
+    get_anchor_task_name,
+    get_anchor_token,
+    get_comt_data_in_response,
+    get_image_info,
+    get_rope_index_3,
+    get_token_num,
+    get_video_info,
+    llava_to_openai,
+    pad_and_cat,
+)
+from models.blip3o.params import DataArguments
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+# TUNABLE: VAE / diffusion target resolution for gen_image
+target_transform = v2.Compose(
+    [
+        v2.Resize(512),
+        v2.CenterCrop(512),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+        v2.Normalize([0.5], [0.5]),
+    ]
+)
+
+
+def rank0_print(*args):
+    if os.environ.get("LOCAL_RANK", "0") in ("0", "-1"):
+        print(*args)
+
+
+def preprocess_multimodal(sources, data_args: DataArguments):
+    """Replace LLAVA <image> in assistant turn with <begin_of_img>…<end_of_img> metaquery block."""
+    gen_placeholder = (
+        f"\n<begin_of_img>"
+        + "".join([f"<img{i}>" for i in range(data_args.num_metaqueries)])
+        + "<end_of_img>"
+    )
+    inst_type = None
+    for source in sources:
+        for sentence in source:
+            if sentence["from"] == "gpt" and LLAVA_IMAGE_TOKEN in sentence["value"]:
+                sentence["value"] = sentence["value"].replace(LLAVA_IMAGE_TOKEN, gen_placeholder).strip()
+                inst_type = "gen"
+    return sources, inst_type
+
+def _pil_from_sample(img) -> PILImage.Image:
+    if isinstance(img, bytes):
+        return PILImage.open(io.BytesIO(img)).convert("RGB")
+    if isinstance(img, PILImage.Image):
+        return img.convert("RGB")
+    if isinstance(img, dict) and "bytes" in img:
+        return PILImage.open(io.BytesIO(img["bytes"])).convert("RGB")
+    raise TypeError(f"Unsupported image type: {type(img)}")
+
+def _covt_targets_gpt(user_input: str, inst_type: Optional[str], task_type: str) -> bool:
+    """Whether to inject CoVT CoT / anchor pads into the assistant turn."""
+    if DEFAULT_IMAGE_TOKEN in user_input or DEFAULT_VIDEO_TOKEN in user_input:
+        return True
+    return inst_type == "gen" or task_type in ("T2I", "I2I")
+
+
+def get_t21_data_files(data_path: str):
+    rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", 0)))
+    max_retries = 5
+    retry_delay = 2
+    for attempt in range(max_retries):
+        try:
+            if rank == 0 or attempt == max_retries - 1:
+                print(f"[Rank {rank}] Loading T2I json (attempt {attempt + 1}/{max_retries})...")
+            train_dataset = load_dataset(
+                "json",
+                data_files=data_path,
+                split="train",
+                download_mode="reuse_dataset_if_exists",
+            )
+            if rank == 0:
+                print(f"[Rank {rank}] T2I dataset loaded.")
+            break
+        except PermissionError as e:
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay)
+                retry_delay *= 2
+            else:
+                raise e
+    train_dataset = train_dataset.rename_column("image_path", "image")
+    train_dataset = train_dataset.rename_column("prompt", "txt")
+    train_dataset = train_dataset.remove_columns(
+        [col for col in train_dataset.column_names if col not in ("image", "txt")]
+    )
+    return train_dataset
+
+
+def force_align_image_schema(ds):
+    std_image = Image(decode=True)
+    if "image" in ds.features:
+        ds = ds.cast_column("image", std_image)
+    if "ref_image" in ds.features:
+        ds = ds.cast_column("ref_image", std_image)
+    else:
+        ds = ds.map(lambda x: {**x, "ref_image": None})
+        ds = ds.cast_column("ref_image", std_image)
+    return ds
+
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning (UND / VQA), aligned with ``covt_data_var.SupervisedDataset``."""
+
+    def __init__(
+        self,
+        data_path: str | list,
+        processor: transformers.ProcessorMixin,
+        data_args: DataArguments,
+        model_id,
+        padding=True,
+        shuffle=True,
+        random_seed=42,
+        anchor_model_id=None,
+    ):
+        super(SupervisedDataset, self).__init__()
+        if isinstance(data_path, str):
+            list_data_dict = json.load(open(data_path, "r"))
+        else:
+            list_data_dict = data_path
+
+        self.model_id = model_id
+        self.processor = processor
+        self.list_data_dict = list_data_dict
+        self.data_args = data_args
+        self.padding = padding
+        self.image_min_pixel = data_args.image_min_pixels
+        self.image_max_pixel = data_args.image_max_pixels
+        self.image_resized_w = data_args.image_resized_width
+        self.image_resized_h = data_args.image_resized_height
+        self.video_min_pixel = data_args.video_min_pixels
+        self.video_max_pixel = data_args.video_max_pixels
+        self.fps = data_args.fps
+        self.merge_size = getattr(processor.image_processor, "merge_size", 2)
+        self.get_rope_index = get_rope_index_3
+        self.anchor_model_id = anchor_model_id
+        self.anchor_token_nums = get_token_num(anchor_model_id)
+        self.anchor_tokens = get_anchor_token(anchor_model_id)
+        self.anchor_task_names = get_anchor_task_name(anchor_model_id)
+
+        # for shuffle
+        self.rng = np.random.default_rng(seed=random_seed)
+
+        if shuffle:
+            self.rng.shuffle(self.list_data_dict)
+
+    def set_cur_step(self, step: int) -> None:
+        """ResumeDatasetCallback hook; UND dataset has no step-dependent sampling."""
+        del step  # unused
+
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        sources = self.list_data_dict[i]
+
+        is_video = False
+        image_files = None
+
+        processor = self.processor
+        if "image" in sources:
+            videos = None
+            grid_key = "image_grid_thw"
+            pixel_key = "pixel_values"
+
+            image_files = sources["image"]
+            image_folder = self.data_args.image_folder
+
+            if isinstance(image_files, str):
+                image_files = PILImage.open(image_files).convert("RGB")
+                image_files = [image_files]
+            else:
+                image_files = [PILImage.open(image_file).convert("RGB") for image_file in image_files]
+            images = []
+
+            for image_file in image_files:
+                images.append(
+                    get_image_info(
+                        image_file,
+                        self.image_min_pixel,
+                        self.image_max_pixel,
+                        self.image_resized_w,
+                        self.image_resized_h,
+                    )
+                )
+
+        elif "video" in sources:
+            is_video = True
+            images = None
+            grid_key = "video_grid_thw"
+            pixel_key = "pixel_values_videos"
+
+            video_files = sources["video"]
+            video_folder = self.data_args.image_folder
+
+            if isinstance(video_files, str):
+                video_files = [video_files]
+
+            videos = []
+            for video_file in video_files:
+                if not os.path.exists(video_file):
+                    if not video_file.startswith("http"):
+                        video_file = os.path.join(video_folder, video_file)
+                video_input, video_kwargs = get_video_info(
+                    video_file, self.video_min_pixel, self.video_max_pixel, self.data_args.fps
+                )
+                videos.append(video_input)
+        else:
+            grid_key = None
+            pixel_key = None
+            images = None
+            videos = None
+
+        if images is None:
+            rank0_print("No image or video found in the data.")
+            images = []
+            black_image = PILImage.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(
+                get_image_info(
+                    black_image,
+                    self.image_min_pixel,
+                    self.image_max_pixel,
+                    self.image_resized_w,
+                    self.image_resized_h,
+                )
+            )
+
+        elif len(images) == 0:
+            rank0_print("No image or video found in the data.")
+            black_image = PILImage.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
+            images.append(
+                get_image_info(
+                    black_image,
+                    self.image_min_pixel,
+                    self.image_max_pixel,
+                    self.image_resized_w,
+                    self.image_resized_h,
+                )
+            )
+        sources = copy.deepcopy(llava_to_openai(sources["conversations"], is_video=is_video))
+
+        all_input_ids = []
+        all_labels = []
+        all_pixel_values = []
+        all_image_grid_thw = []
+        all_second_gird = []
+
+        if len(SYSTEM_MESSAGE) > 0:
+            system_message = f"{DEFAULT_IM_START_TOKEN}system\n{SYSTEM_MESSAGE}\n{DEFAULT_IM_END_TOKEN}\n"
+            system_message_input_ids = processor.tokenizer(system_message, add_special_tokens=False, return_tensors="pt")[
+                "input_ids"
+            ]
+
+            system_labels = torch.full_like(system_message_input_ids, IGNORE_INDEX)
+            all_input_ids.append(system_message_input_ids.squeeze(0))
+            all_labels.append(system_labels.squeeze(0))
+
+        for _, j in enumerate(range(0, len(sources), 2)):
+            if j >= 2:
+                break
+
+            user_input = sources[j]
+            gpt_response = sources[j + 1]
+            if (DEFAULT_IMAGE_TOKEN not in user_input["content"]) and (DEFAULT_VIDEO_TOKEN not in user_input["content"]) and (
+                LLAVA_IMAGE_TOKEN in user_input["content"]
+            ):
+                user_input = (
+                    f"{DEFAULT_IM_START_TOKEN}{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}"
+                    f"{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    f"{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                )
+                user_input = add_anchor_pad(user_input, self.anchor_token_nums, self.anchor_tokens)
+                gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                raise ValueError("Every man is a poet when he is in love")
+            else:
+                user_input = (
+                    f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n"
+                    f"{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                )
+                # TUNABLE: The probability boundary between the full set and the random subset anchor CoT.
+                if random.random() < 0.5:
+                    gpt_response = f"{gpt_response['content']}"
+                    if DEFAULT_IMAGE_TOKEN in user_input:
+                        gpt_response = get_comt_data_in_response(
+                            gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names
+                        )
+                    gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+                else:
+                    xxx = random.randint(0, 5)
+                    if xxx == 0:
+                        gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    else:
+                        gpt_response = f"{gpt_response['content']}"
+                        if DEFAULT_IMAGE_TOKEN in user_input:
+                            total = len(self.anchor_tokens)
+                            if total == 0:
+                                selected_anchor_token_nums = []
+                                selected_anchor_tokens = []
+                                selected_anchor_task_names = []
+                            else:
+                                x = random.randint(1, total)
+                                idxs = sorted(random.sample(range(total), x)) if x > 0 else []
+                                selected_anchor_token_nums = [self.anchor_token_nums[i] for i in idxs]
+                                selected_anchor_tokens = [self.anchor_tokens[i] for i in idxs]
+                                selected_anchor_task_names = [self.anchor_task_names[i] for i in idxs]
+                            gpt_response = get_comt_data_in_response(
+                                gpt_response,
+                                selected_anchor_token_nums,
+                                selected_anchor_tokens,
+                                selected_anchor_task_names,
+                            )
+                        gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+
+            if DEFAULT_IMAGE_TOKEN in user_input:
+                inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors="pt")
+                prompt_input_ids = inputs["input_ids"]
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+                torch.cuda.empty_cache()
+
+
+            elif DEFAULT_VIDEO_TOKEN in user_input:
+                if "Qwen3" in self.model_id:
+                    inputs = processor(
+                        text=[user_input],
+                        images=images,
+                        videos=videos,
+                        padding=False,
+                        return_tensors="pt",
+                        **video_kwargs,
+                    )
+                    all_second_gird.extend(inputs["second_per_grid_ts"])
+                else:
+                    inputs = processor(
+                        text=[user_input], images=images, videos=videos, padding=False, return_tensors="pt"
+                    )
+                prompt_input_ids = inputs["input_ids"]
+                all_pixel_values.append(inputs[pixel_key])
+                all_image_grid_thw.append(inputs[grid_key])
+
+            else:
+                prompt_input_ids = processor.tokenizer(
+                    user_input, add_special_tokens=False, padding=False, return_tensors="pt"
+                )["input_ids"]
+
+            response_input_ids = processor.tokenizer(
+                gpt_response, add_special_tokens=False, padding=False, return_tensors="pt"
+            )["input_ids"]
+
+            input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
+            labels = torch.cat(
+                [
+                    torch.tensor([IGNORE_INDEX] * len(prompt_input_ids[0])),
+                    response_input_ids.squeeze(0),
+                ],
+                dim=0,
+            )
+            all_input_ids.append(input_ids)
+            all_labels.append(labels)
+
+        input_ids = torch.cat(all_input_ids, dim=0).to(torch.long)
+        labels = torch.cat(all_labels, dim=0).to(torch.long)
+
+        attention_mask = (input_ids > -1000000).to(torch.long)
+        task_type = "und_image"
+
+        x_ids, x_mask = input_ids, attention_mask
+        while x_ids.dim() < 2:
+            x_ids = x_ids.unsqueeze(0)
+
+        input_ids_2d = x_ids
+        data_dict = dict(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            task_types=task_type,
+        )
+
+        if pixel_key and grid_key:
+            pixel_values = torch.cat(all_pixel_values, dim=0)
+            image_thw = torch.cat(all_image_grid_thw, dim=0)
+
+            data_dict[pixel_key] = pixel_values
+            data_dict[grid_key] = image_thw
+            data_dict["image_files"] = image_files
+
+        if len(all_second_gird) > 0:
+            second_gird = all_second_gird
+            data_dict["second_per_grid_ts"] = second_gird
+
+        position_ids, _ = self.get_rope_index(
+            self.merge_size,
+            input_ids_2d,
+            image_grid_thw=torch.cat(all_image_grid_thw, dim=0) if images else None,
+            video_grid_thw=(torch.cat(all_image_grid_thw, dim=0) if videos else None),
+            second_per_grid_ts=all_second_gird if len(all_second_gird) > 0 else None,
+        )
+        data_dict["position_ids"] = position_ids
+
+        # SmartEdit LLaVAMLLMSD: understanding samples mask diffusion branch
+        data_dict["is_gen_task"] = torch.zeros(1)
+
+        return data_dict
+
+
+class SupervisedT2IDataset(Dataset):
+    """T2I / I2I mix dataset: Qwen3-VL processor + CoVT anchors + gen_image."""
+
+    def __init__(
+        self,
+        data_path: str,
+        processor: transformers.ProcessorMixin,
+        data_args: DataArguments,
+        model_id: str,
+        padding: bool = True,
+        shuffle: bool = False,
+        random_seed: int = 42,
+        anchor_model_id: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        self.model_id = model_id
+        self.processor = processor
+        self.data_args = data_args
+        self.padding = padding
+        self.image_min_pixel = data_args.image_min_pixels
+        self.image_max_pixel = data_args.image_max_pixels
+        self.image_resized_w = data_args.image_resized_width
+        self.image_resized_h = data_args.image_resized_height
+        self.video_min_pixel = data_args.video_min_pixels
+        self.video_max_pixel = data_args.video_max_pixels
+        self.fps = data_args.fps
+        self.merge_size = getattr(processor.image_processor, "merge_size", 2)
+        self.get_rope_index = get_rope_index_3
+        self.anchor_model_id = anchor_model_id or []
+        self.anchor_token_nums = get_token_num(self.anchor_model_id)
+        self.anchor_tokens = get_anchor_token(self.anchor_model_id)
+        self.anchor_task_names = get_anchor_task_name(self.anchor_model_id)
+
+        t2i_dataset = get_t21_data_files(data_path)
+        t2i_dataset = (
+            t2i_dataset.add_column("type", ["T2I"] * len(t2i_dataset)).select_columns(["image", "txt", "type"])
+        )
+        t2i_dataset = force_align_image_schema(t2i_dataset)
+        rank0_print(f"Loaded T2I dataset: {len(t2i_dataset)} samples")
+        self.list_data_dict = t2i_dataset.shuffle(seed=random_seed) if shuffle else t2i_dataset
+        rank0_print(f"Total training instances: {len(self.list_data_dict)}")
+
+    def set_cur_step(self, step: int) -> None:
+        """ResumeDatasetCallback hook; T2I dataset has no step-dependent sampling."""
+        del step  # unused
+
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    def _build_user_gpt_strings(
+        self,
+        user_turn: Dict,
+        gpt_turn: Dict,
+        inst_type: Optional[str],
+        task_type: str,
+    ) -> Tuple[str, str]:
+        user_input = (
+            f"{DEFAULT_IM_START_TOKEN}{user_turn['role']}\n{user_turn['content']}\n"
+            f"{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_turn['role']}\n"
+        )
+
+        if random.random() < 0.5:
+            gpt_response = f"{gpt_turn['content']}"
+            if _covt_targets_gpt(user_input, inst_type, task_type):
+                gpt_response = get_comt_data_in_response(
+                    gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names
+                )
+            gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+        else:
+            if random.randint(0, 5) == 0:
+                gpt_response = f"{gpt_turn['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+            else:
+                gpt_response = f"{gpt_turn['content']}"
+                if _covt_targets_gpt(user_input, inst_type, task_type):
+                    total = len(self.anchor_tokens)
+                    if total == 0:
+                        sel_nums, sel_toks, sel_names = [], [], []
+                    else:
+                        x = random.randint(1, total)
+                        idxs = sorted(random.sample(range(total), x))
+                        sel_nums = [self.anchor_token_nums[i] for i in idxs]
+                        sel_toks = [self.anchor_tokens[i] for i in idxs]
+                        sel_names = [self.anchor_task_names[i] for i in idxs]
+                    gpt_response = get_comt_data_in_response(gpt_response, sel_nums, sel_toks, sel_names)
+                gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+        return user_input, gpt_response
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        processor = self.processor
+        while True:
+            row = self.list_data_dict[i]
+            task_type = row["type"]
+            if task_type == "T2I":
+                conversations = [
+                    {
+                        "from": "human",
+                        "value": (
+                            f"Please generate an image that matches the following description: {row['txt']}"
+                        ),
+                    },
+                    {"from": "gpt", "value": LLAVA_IMAGE_TOKEN},
+                ]
+                target_pils = [_pil_from_sample(row["image"])]
+                ref_pils = None
+            elif task_type == "I2I":
+                conversations = [
+                    {
+                        "from": "human",
+                        "value": (
+                            f"{LLAVA_IMAGE_TOKEN}\nPlease edit the given image based on the "
+                            f"following instruction: {row['txt']}."
+                        ),
+                    },
+                    {"from": "gpt", "value": LLAVA_IMAGE_TOKEN},
+                ]
+                try:
+                    target_pils = [_pil_from_sample(row["image"])]
+                    ref_pils = [_pil_from_sample(row["ref_image"])]
+                except Exception as e:
+                    rank0_print(f"I2I image load error: {e}; resample.")
+                    i = random.randint(0, len(self.list_data_dict) - 1)
+                    continue
+            else:
+                raise ValueError(f"Unknown task type: {task_type}")
+
+            sources, inst_type = preprocess_multimodal(copy.deepcopy([conversations]), self.data_args)
+            if inst_type != "gen":
+                i = random.randint(0, len(self.list_data_dict) - 1)
+                continue
+
+            openai_turns = copy.deepcopy(llava_to_openai(sources[0], is_video=False))
+
+            # Vision tensors for Qwen3-VL (I2I: ref image in user turn)
+            videos = None
+            grid_key = "image_grid_thw"
+            pixel_key = "pixel_values"
+            vision_images: List[Any] = []
+            if task_type == "I2I" and ref_pils is not None:
+                for ref in ref_pils:
+                    vision_images.append(
+                        get_image_info(
+                            ref,
+                            self.image_min_pixel,
+                            self.image_max_pixel,
+                            self.image_resized_w,
+                            self.image_resized_h,
+                        )
+                    )
+            image_files = target_pils
+
+            all_input_ids: List[torch.Tensor] = []
+            all_labels: List[torch.Tensor] = []
+            all_pixel_values: List[torch.Tensor] = []
+            all_image_grid_thw: List[torch.Tensor] = []
+            all_second_grid: List = []
+
+            if len(SYSTEM_MESSAGE) > 0:
+                system_message = (
+                    f"{DEFAULT_IM_START_TOKEN}system\n{SYSTEM_MESSAGE}\n{DEFAULT_IM_END_TOKEN}\n"
+                )
+                system_ids = processor.tokenizer(system_message, add_special_tokens=False, return_tensors="pt")[
+                    "input_ids"
+                ].squeeze(0)
+                all_input_ids.append(system_ids)
+                all_labels.append(torch.full_like(system_ids, IGNORE_INDEX))
+
+            for j in range(0, len(openai_turns), 2):
+                if j >= 2:
+                    break
+                user_turn = openai_turns[j]
+                gpt_turn = openai_turns[j + 1]
+                user_str, gpt_str = self._build_user_gpt_strings(user_turn, gpt_turn, inst_type, task_type)
+
+                if DEFAULT_IMAGE_TOKEN in user_str or DEFAULT_VIDEO_TOKEN in user_str:
+                    if "Qwen3" in self.model_id and DEFAULT_VIDEO_TOKEN in user_str:
+                        inputs = processor(
+                            text=[user_str],
+                            images=vision_images if vision_images else None,
+                            videos=videos,
+                            padding=False,
+                            return_tensors="pt",
+                        )
+                        if "second_per_grid_ts" in inputs:
+                            all_second_grid.extend(inputs["second_per_grid_ts"])
+                    else:
+                        inputs = processor(
+                            text=[user_str],
+                            images=vision_images if vision_images else None,
+                            videos=videos,
+                            padding=False,
+                            return_tensors="pt",
+                        )
+                    prompt_input_ids = inputs["input_ids"]
+                    if pixel_key in inputs and inputs[pixel_key] is not None:
+                        all_pixel_values.append(inputs[pixel_key])
+                        all_image_grid_thw.append(inputs[grid_key])
+                else:
+                    prompt_input_ids = processor.tokenizer(
+                        user_str, add_special_tokens=False, padding=False, return_tensors="pt"
+                    )["input_ids"]
+
+                response_input_ids = processor.tokenizer(
+                    gpt_str, add_special_tokens=False, padding=False, return_tensors="pt"
+                )["input_ids"]
+
+                input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
+                labels = torch.cat(
+                    [
+                        torch.full((prompt_input_ids.shape[1],), IGNORE_INDEX, dtype=torch.long),
+                        response_input_ids.squeeze(0),
+                    ],
+                    dim=0,
+                )
+                all_input_ids.append(input_ids)
+                all_labels.append(labels)
+
+            input_ids = torch.cat(all_input_ids, dim=0).to(torch.long)
+            labels = torch.cat(all_labels, dim=0).to(torch.long)
+            attention_mask = (input_ids > -1000000).to(torch.long)
+
+            input_ids_2d = input_ids.unsqueeze(0) if input_ids.dim() == 1 else input_ids
+
+            data_dict: Dict[str, Any] = {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels": labels,
+                "task_types": task_type,
+                "image_files": image_files,
+            }
+
+            if all_pixel_values:
+                data_dict[pixel_key] = torch.cat(all_pixel_values, dim=0)
+                data_dict[grid_key] = torch.cat(all_image_grid_thw, dim=0)
+
+            if all_second_grid:
+                data_dict["second_per_grid_ts"] = all_second_grid
+
+            position_ids, _ = self.get_rope_index(
+                self.merge_size,
+                input_ids_2d,
+                image_grid_thw=data_dict.get(grid_key),
+                video_grid_thw=None,
+                second_per_grid_ts=data_dict.get("second_per_grid_ts"),
+            )
+            data_dict["position_ids"] = position_ids
+
+            # TUNABLE: gen_image spatial size via data_args.image_size or target_transform (512)
+            if inst_type == "gen":
+                if getattr(self.data_args, "gen_image_processor", None) is not None:
+                    sz = self.data_args.gen_image_size
+                    resized = [
+                        img.resize((sz, sz), resample=PILImage.Resampling.BICUBIC) for img in target_pils
+                    ]
+                    data_dict["gen_image"] = self.data_args.gen_image_processor.preprocess(
+                        resized, return_tensors="pt"
+                    )["pixel_values"]
+                else:
+                    data_dict["gen_image"] = torch.stack(
+                        [target_transform(img) for img in target_pils], dim=0
+                    )
+
+                if task_type == "I2I" and ref_pils is not None:
+                    if getattr(self.data_args, "gen_image_processor", None) is not None:
+                        sz = self.data_args.gen_image_size
+                        ref_resized = [
+                            img.resize((sz, sz), resample=PILImage.Resampling.BICUBIC) for img in ref_pils
+                        ]
+                        data_dict["gen_ref_image"] = self.data_args.gen_image_processor.preprocess(
+                            ref_resized, return_tensors="pt"
+                        )["pixel_values"]
+                    else:
+                        data_dict["gen_ref_image"] = torch.stack(
+                            [target_transform(img) for img in ref_pils], dim=0
+                        )
+                    data_dict["ref_mask"] = torch.ones(
+                        (1, 1, 1, 1), dtype=data_dict["gen_image"].dtype, device=data_dict["gen_image"].device
+                    )
+
+            # SmartEdit ReasonEditMLLMSD: generation / editing samples enable diffusion branch
+            data_dict["is_gen_task"] = torch.ones(1)
+
+            return data_dict
+
+
+@dataclass
+class DataCollatorForSupervisedDataset:
+    """Pad text/mm tensors (Qwen3-VL position_ids) and stack gen_image for T2I rows."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels, position_ids = tuple(
+            [instance[key] for instance in instances] for key in ("input_ids", "labels", "position_ids")
+        )
+        input_ids = [ids.squeeze(0) if ids.dim() > 1 else ids for ids in input_ids]
+        labels = [ids.squeeze(0) if ids.dim() > 1 else ids for ids in labels]
+
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        position_ids = pad_and_cat(position_ids)
+
+        input_ids = input_ids[:, : self.tokenizer.model_max_length]
+        labels = labels[:, : self.tokenizer.model_max_length]
+        position_ids = position_ids[:, :, : self.tokenizer.model_max_length]
+
+        batch: Dict[str, Any] = {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": input_ids.ne(self.tokenizer.pad_token_id),
+            "position_ids": position_ids,
+        }
+
+        pixel_list = [ex["pixel_values"] for ex in instances if "pixel_values" in ex]
+        if pixel_list:
+            batch["pixel_values"] = torch.cat(pixel_list, dim=0)
+            batch["image_grid_thw"] = torch.cat(
+                [ex["image_grid_thw"] for ex in instances if "image_grid_thw" in ex], dim=0
+            )
+        else:
+            batch["pixel_values"] = None
+            batch["image_grid_thw"] = None
+
+        batch["image_files"] = [ex.get("image_files") for ex in instances]
+
+        gen_images: List[torch.Tensor] = []
+        t2i_row_indices: List[int] = []
+        gen_ref_images: List[torch.Tensor] = []
+        ref_masks: List[torch.Tensor] = []
+
+        for idx, ex in enumerate(instances):
+            if ex.get("task_types") == "T2I" and "gen_image" in ex:
+                gi = ex["gen_image"]
+                if gi.dim() == 3:
+                    gi = gi.unsqueeze(0)
+                gen_images.append(gi)
+                t2i_row_indices.append(idx)
+            if "gen_ref_image" in ex:
+                gen_ref_images.append(ex["gen_ref_image"])
+            if "ref_mask" in ex:
+                ref_masks.append(ex["ref_mask"])
+
+        if gen_images:
+            batch["gen_image"] = torch.cat(gen_images, dim=0)
+            batch["t2i_row_indices"] = torch.tensor(t2i_row_indices, dtype=torch.long)
+        else:
+            batch["gen_image"] = None
+            batch["t2i_row_indices"] = None
+
+        if gen_ref_images:
+            batch["gen_ref_image"] = torch.cat(gen_ref_images, dim=0)
+        if ref_masks:
+            batch["ref_mask"] = torch.cat(ref_masks, dim=0)
+
+        if all("is_gen_task" in ex for ex in instances):
+            batch["is_gen_task"] = torch.stack([ex["is_gen_task"] for ex in instances])
+        return batch
+
+class Merge_UndT2IDataset(Dataset):
+    """SmartEdit-style merge: each ``__getitem__`` randomly draws UND or T2I by ``und_weight:t2i_weight``."""
+
+    def __init__(
+        self,
+        und_dataset: SupervisedDataset,
+        t2i_dataset: SupervisedT2IDataset,
+        und_weight: int = 1,
+        t2i_weight: int = 3,
+    ):
+        if und_weight <= 0 or t2i_weight <= 0:
+            raise ValueError(f"und_weight and t2i_weight must be positive, got {und_weight}, {t2i_weight}")
+        self.und_dataset = und_dataset
+        self.t2i_dataset = t2i_dataset
+        self.und_len = len(und_dataset)
+        self.t2i_len = len(t2i_dataset)
+        self.total_len = self.und_len + self.t2i_len
+        self.und_weight = int(und_weight)
+        self.t2i_weight = int(t2i_weight)
+        self.und_prob = self.und_weight / (self.und_weight + self.t2i_weight)
+
+    def set_cur_step(self, step: int) -> None:
+        """Forward to child datasets that support resume step sync."""
+        if hasattr(self.und_dataset, "set_cur_step"):
+            self.und_dataset.set_cur_step(step)
+        if hasattr(self.t2i_dataset, "set_cur_step"):
+            self.t2i_dataset.set_cur_step(step)
+
+    def __getitem__(self, index):
+        if self.total_len <= 0:
+            raise IndexError("Merge_UndT2IDataset has zero total length")
+        if random.random() < self.und_prob:
+            return self.und_dataset[random.randint(0, self.und_len - 1)]
+        return self.t2i_dataset[random.randint(0, self.t2i_len - 1)]
+
+    def __len__(self):
+        return self.total_len
+
+    def __repr__(self) -> str:
+        return (
+            f"Merge_UndT2IDataset(und_len={self.und_len}, t2i_len={self.t2i_len}, total_len={self.total_len}, "
+            f"und_weight={self.und_weight}, t2i_weight={self.t2i_weight}, und_prob={self.und_prob:.4f})"
+        )
+
+
+def make_merged_und_t2i_gen_data_module(
+    model_id: str,
+    processor: transformers.ProcessorMixin,
+    data_args: DataArguments,
+    anchor_model_id: Optional[List[str]],
+    und_data_path: str,
+    t2i_data_path: str,
+) -> Dict[str, Any]:
+    """Build separate UND + T2I datasets and merge them (DS_SmartEdit_train Merge_Dataset pattern)."""
+    und_train_dataset = SupervisedDataset(
+        data_path=und_data_path,
+        processor=processor,
+        data_args=data_args,
+        model_id=model_id,
+        anchor_model_id=anchor_model_id,
+    )
+    t2i_train_dataset = SupervisedT2IDataset(
+        data_path=t2i_data_path,
+        processor=processor,
+        data_args=data_args,
+        model_id=model_id,
+        anchor_model_id=anchor_model_id,
+    )
+    und_w = int(getattr(data_args, "mix_num_und_per_device", 1) or 1)
+    t2i_w = int(getattr(data_args, "mix_num_t2i_per_device", 3) or 3)
+    merged_train_dataset = Merge_UndT2IDataset(
+        und_train_dataset, t2i_train_dataset, und_weight=und_w, t2i_weight=t2i_w
+    )
+    rank0_print(merged_train_dataset)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=processor.tokenizer)
+    return dict(train_dataset=merged_train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+def make_supervised_data_module(
+    model_id: str,
+    processor: transformers.ProcessorMixin,
+    data_args: DataArguments,
+    anchor_model_id: Optional[List[str]],
+) -> Dict[str, Any]:
+    # TUNABLE: T2I-only 时用 t2i_data_path，否则回退 data_path
+    t2i_path = getattr(data_args, "t2i_data_path", None) or data_args.data_path
+    train_dataset = SupervisedT2IDataset(
+        data_path=t2i_path,
+        processor=processor,
+        data_args=data_args,
+        model_id=model_id,
+        anchor_model_id=anchor_model_id,
+    )
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=processor.tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)

code/umm/runtime/models/blip3o/data/t2i_covt_data.py

@@ -0,0 +1,742 @@
+
+import os
+import io
+import copy
+from dataclasses import dataclass, field
+import json
+import logging
+import pathlib
+from typing import Dict, Optional, Sequence, List, Tuple
+
+import time
+import torch, gc
+import glob
+import torch.nn as nn
+import transformers
+import tokenizers
+import random
+from torch.utils.data import Dataset
+
+import PIL
+from PIL import  ImageFile
+from datasets import load_dataset, concatenate_datasets,load_from_disk
+
+from pathlib import Path
+from datasets.utils.logging import set_verbosity_info
+from transformers import logging as tf_logging
+import torchvision.transforms as T
+from torchvision.transforms import v2
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoProcessor
+from packaging import version
+from datasets import Image
+from models.blip3o.constants import *
+from models.blip3o.params import DataArguments
+
+from .covt_data import SupervisedDataset
+# Set global configurations
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+# Image preprocessing for underspecified (contextual) images
+transform_und_images = T.Compose([
+    T.Resize(672, interpolation=InterpolationMode.BICUBIC, antialias=True), 
+    T.CenterCrop(672) # Target size for a fixed number of tokens
+]) 
+
+
+## target transform for sana
+target_transform = T.Compose(
+    [
+        v2.Resize(512),
+        v2.CenterCrop(512),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+        v2.Normalize([0.5], [0.5]),
+    ]
+    )
+set_verbosity_info()
+tf_logging.set_verbosity_info()
+local_rank = None
+
+
+# Check tokenizer version for compatibility
+IS_TOKENIZER_GREATER_THAN_0_14 = version.parse(tokenizers.__version__) >= version.parse("0.14")
+def rank0_print(*args):
+    """Prints only on the main process (rank 0)."""
+    if local_rank == 0:
+        print(*args)
+
+
+
+def get_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
+    if len(anchor_nums) == 0:
+        return response
+    
+    anchor_pads = []
+    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
+        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
+        anchor_pads.append(anchor_pad)
+    CoT_start = "<think> Because "
+    if len(anchor_names) == 1:
+        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
+    else:
+        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
+            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
+            if anchor_name == anchor_names[-2]:
+                CoT_start += ", and "
+            elif anchor_name == anchor_names[-1]:
+                CoT_start += ". "
+            else:
+                CoT_start += ", "
+    response = CoT_start + " </think>\n" + "<answer> " + response + " </answer>"
+    return response
+
+def apply_template(tokenizer, conv):
+    out = tokenizer.apply_chat_template(
+        conv,
+        tokenize=True,
+        return_tensors=None
+    )
+    
+    return out["input_ids"]
+ 
+
+def llava_to_openai(conversations, is_video=False):
+    role_mapping = {"human": "user", "gpt": "assistant"}
+    
+    # import ipdb; ipdb.set_trace()
+    transformed_data = []
+    conversations = conversations[0]
+    for conversation in conversations:
+        transformed_content = conversation["value"]
+        transformed_entry = {
+            "role": role_mapping.get(conversation["from"], conversation["from"]),
+            "content": transformed_content,
+        }
+        transformed_data.append(transformed_entry)
+
+    return transformed_data
+
+def preprocess_multimodal(sources, data_args):
+    """
+    Replaces generic image placeholder tokens with model-specific query tokens 
+    (T2I, I2I, or UND_IMAGE) in the conversation structure.
+
+    Args:
+        sources (Sequence[List[Dict]]): A sequence of conversations, where each conversation 
+                                        is a list of message dictionaries.
+        data_args (DataArguments): Data arguments containing token length (n_und_query).
+    Returns:
+        Tuple[Sequence[List[Dict]], Optional[str]]: 
+            - sources: The modified conversation structure.
+            - inst_type: The type of image task ('und' for context, 'gen' for generation), 
+                         or None if no image token is found.
+    """
+
+    # Placeholder for the underspecified (contextual) image
+    # und_placeholder = "<|vision_start|>" + UND_IMAGE_TOKEN * data_args.n_und_query + "<|vision_end|>"
+    gen_placeholder =  f"\n<begin_of_img>" + "".join([f"<img{i}>" for i in range(data_args.num_metaqueries)]) + "<end_of_img>"
+    inst_type = None
+    for source in sources:  # Iterate through batch/instance
+        for sentence in source: # Iterate through conversation turns
+            if sentence["from"] == "gpt" and LLAVA_IMAGE_TOKEN in sentence["value"]:
+                # The generative token will be replaced by T2I/I2I query tokens later in data collator
+                sentence["value"] = sentence["value"].replace(LLAVA_IMAGE_TOKEN, gen_placeholder).strip()
+                inst_type = "gen"
+    return sources, inst_type
+
+def preprocess_qwen(sources: Sequence[List[Dict]], tokenizer: transformers.PreTrainedTokenizer, system_message: str = "You are a helpful assistant.") -> Dict:
+    """
+    Tokenizes the conversation data according to the Qwen chat template, 
+    creating input_ids and attention mask/labels.
+    Args:
+        sources (Sequence[List[Dict]]): A sequence of preprocessed conversations.
+        tokenizer (transformers.PreTrainedTokenizer): The tokenizer instance.
+        has_image (bool): Flag indicating if the instance contains an image.
+        system_message (str): The initial system instruction.
+    Returns:
+        Dict: Dictionary containing 'input_ids' and 'labels' tensors.
+    """
+    roles = {"human": "user", "gpt": "assistant"}
+
+    # Use a deep copy to avoid modifying the global tokenizer chat template if parallel processing is used
+    tokenizer = copy.deepcopy(tokenizer)
+    chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+    tokenizer.chat_template = chat_template
+
+    input_ids, targets = [], []
+    for i, source in enumerate(sources):
+        if roles[source[0]["from"]] != roles["human"]:
+            source = source[1:]
+
+        input_id, target = [], []
+
+        # 1. Build system message
+        system_conv = [{"role" : "system", "content" : system_message}]
+        input_id += apply_template(tokenizer,system_conv)
+        target += [IGNORE_INDEX] * len(input_id)
+
+        # 2. Process conversation turns
+        for conv in source:
+            try:
+                role = conv["role"]
+                content = conv["content"]
+            except: # Handle older format
+                role = conv["from"]
+                content = conv["value"]
+            role = roles.get(role, role)
+
+            conv_turn = [{"role" : role, "content" : content}]
+            encode_id = apply_template(tokenizer,conv_turn)
+            input_id += encode_id
+            
+            # Mask out user/system input for loss calculation
+            if role in ["user", "system"]:
+                target += [IGNORE_INDEX] * len(encode_id)
+            else:
+                target += encode_id # Label the GPT responses
+        
+        assert len(input_id) == len(target), f"Input length mismatch: {len(input_id)} != {len(target)}"
+
+        input_ids.append(input_id)
+        targets.append(target)
+    
+    # import pdb
+    # pdb.set_trace()
+    input_ids = torch.tensor(input_ids, dtype=torch.long)
+    targets = torch.tensor(targets, dtype=torch.long)
+    
+    # valid = (targets != -100).sum()
+    # print("valid tokens:", valid)
+
+    return dict(
+        input_ids=input_ids,  # tensor(bs x seq_len)
+        labels=targets,  # tensor(bs x seq_len)
+    )
+
+def parse_txt(example: Dict) -> Dict:
+    """Extracts text prompt from the 'edited_prompt_list' field (for I2I data)."""
+    example["txt"] = example["edited_prompt_list"][-1]
+    return example
+
+def parse_json(example: Dict) -> Dict:
+    """Extracts text prompt from the 'json' field (for T2I data)."""
+    if isinstance(example["json"], str):
+        try:
+            parsed = json.loads(example["json"])
+            example["txt"] = parsed.get("prompt", "") 
+        except json.JSONDecodeError:
+            example["txt"] = "" 
+    else:
+        example["txt"] = example["json"].get("prompt", "")  
+    return example
+
+
+def get_tar_files(directory):
+    return [
+        os.path.join(directory, f)
+        for f in os.listdir(directory)
+        if f.endswith(".tar")
+    ]
+
+def get_t2i_data_files(data_path,num_files: int = 10) -> List[str]:
+    """Generate list of T2I dataset URLs."""
+    import random
+    data_files = get_tar_files(data_path)
+    num_files = min(46, num_files)
+    selected_files = random.sample(data_files, num_files)
+    return selected_files
+
+# def get_edit_data_files(num_files: int = 10) -> List[str]:
+#     """Generate list of image editing dataset URLs."""
+#     num_files = min(571, num_files)
+#     base_url = "https://huggingface.co/datasets/TIGER-Lab/OmniEdit-Filtered-1.2M/resolve/main/data/train-{:05d}-of-00571.parquet"
+#     return [base_url.format(i) for i in range(num_files)]  # 00000 to 00570
+
+# --------------------------------------------------------------------------
+# --- Dataset and Collation ---
+# --------------------------------------------------------------------------
+
+def get_magicbrush_data_files(data_path):
+    train_dataset = load_from_disk(data_path)
+    train_dataset = train_dataset.rename_column("source_img", "ref_image")
+    train_dataset = train_dataset.rename_column("target_img","image")
+    train_dataset = train_dataset.rename_column("instruction","txt")
+    train_dataset = train_dataset.remove_columns(
+        [
+            col
+            for col in train_dataset.column_names
+            if not col in (["ref_image", "image","txt"])
+        ]
+    )
+    return train_dataset
+
+def get_omni_data_files(data_path):
+    train_dataset = load_dataset(
+            "json",
+            data_files=data_path,
+            split="train"
+        )
+    train_dataset = train_dataset.rename_column("source_image", "ref_image")
+    train_dataset = train_dataset.rename_column("target_image","image")
+    train_dataset = train_dataset.rename_column("caption","txt")
+    train_dataset = train_dataset.remove_columns(
+        [
+            col
+            for col in train_dataset.column_names
+            if not col in (["ref_image", "image","txt"])
+        ]
+    )
+    
+    return train_dataset
+
+def get_t21_data_files(data_path):
+    import time
+    import fcntl
+    import errno
+
+    rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", 0)))
+
+    # Use reuse_dataset_if_exists to avoid re-downloading and lock conflicts
+    # Add retry mechanism
+    max_retries = 5
+    retry_delay = 2
+
+    for attempt in range(max_retries):
+        try:
+            if rank == 0 or attempt == max_retries - 1:
+                print(f"[Rank {rank}] Try loading the dataset (trying {attempt + 1}/{max_retries})...")
+            train_dataset = load_dataset(
+                "json",
+                data_files=data_path,
+                split="train",
+                download_mode="reuse_dataset_if_exists",  # Force the use of local caching.
+            )
+            if rank == 0:
+                print(f"[Rank {rank}] Dataset loaded successfully!")
+            break
+        except PermissionError as e:
+            if attempt < max_retries - 1:
+                print(f"[Rank {rank}] File conflict, will retry in {retry_delay} seconds... (trying {attempt + 1}/{max_retries})")
+                time.sleep(retry_delay)
+                retry_delay *= 2  # Exponential Backoff
+            else:
+                raise e
+        except Exception as e:
+            print(f"[Rank {rank}] An error occurred while loading the dataset.: {e}")
+            raise e
+            
+    train_dataset = train_dataset.rename_column("image_path","image")
+    train_dataset = train_dataset.rename_column("prompt","txt")
+    train_dataset = train_dataset.remove_columns(
+        [
+            col
+            for col in train_dataset.column_names
+            if not col in (["image","txt"])
+        ]
+    )
+    
+    return train_dataset
+
+def force_align_image_schema(ds):
+    STD_IMAGE = Image(decode=True)
+    # image
+    if "image" in ds.features:
+        ds = ds.cast_column("image", STD_IMAGE)
+
+    # ref_image
+    if "ref_image" in ds.features:
+        ds = ds.cast_column("ref_image", STD_IMAGE)
+    else:
+        ds = ds.map(lambda x: {**x, "ref_image": None})
+        ds = ds.cast_column("ref_image", STD_IMAGE)
+    return ds
+
+
+class LazySupervisedMixDataset(SupervisedDataset):
+    """
+    Dataset for supervised fine-tuning, loading and mixing I2I and T2I datasets 
+    from Parquet and WebDataset formats.
+    """
+
+    def __init__(
+        self,
+        data_path: str, # Unused but kept for signature consistency
+        processor: transformers.ProcessorMixin,
+        data_args: DataArguments,
+        model_id,
+        padding=True,
+        shuffle=True,
+        random_seed=42,
+        anchor_model_id=None,
+    ):
+        super(LazySupervisedMixDataset, self).__init__(
+            data_path=[],
+            processor=processor,
+            data_args=data_args,
+            model_id=model_id,
+            padding=padding,
+            shuffle=False,
+            random_seed=random_seed,
+            anchor_model_id=anchor_model_id,
+        )
+
+        list_data_dict = []
+
+        t2i_dataset_sa = get_t21_data_files(data_path)
+        t2i_dataset_sa = (t2i_dataset_sa.add_column("type",['T2I']*len(t2i_dataset_sa))
+        .select_columns(["image", "txt", "type"]))
+        rank0_print(f"Finished loading I2I dataset with {len(t2i_dataset_sa)} samples")
+        
+        t2i_dataset_sa = force_align_image_schema(t2i_dataset_sa)
+
+        list_data_dict.append(t2i_dataset_sa)
+        # list_data_dict.append(t2i_dataset_web1)
+        
+        # 3. Concatenate and Shuffle
+        if len(list_data_dict) > 1:
+            list_data_dict = concatenate_datasets(list_data_dict)
+        else:
+            list_data_dict = list_data_dict[0]
+        list_data_dict = list_data_dict.shuffle(seed=42)
+
+        rank0_print(f"Totoal number of training instance: {len(list_data_dict)}")
+        self.list_data_dict = list_data_dict
+
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        """
+        Retrieves a single data instance, constructs the conversation, 
+        and applies preprocessing.
+        """
+        is_video = False
+        processor = self.processor
+
+        while True: # Loop to handle data loading errors
+            sources = self.list_data_dict[i]
+            task_type = sources["type"]
+           
+            if task_type == "T2I":
+                #T2I: Human asks for generation, GPT replies with <image>
+                sources["conversations"] = [
+                    {"from": "human", "value": f"Please generate image based on the following caption: {sources['txt']}"},
+                    {"from": "gpt", "value": LLAVA_IMAGE_TOKEN}, # DEFAULT_IMAGE_TOKEN will be replaced by empty string in preprocess_multimodal
+                ]
+                image_files = sources["image"]   
+                ref_image_files = None  
+                
+            elif task_type == "I2I":
+                # I2I: Human provides image (<image>) and instruction, GPT replies with generation (<image>)
+                sources["conversations"] = [
+                    {
+                        "from": "human",
+                        "value": f"{LLAVA_IMAGE_TOKEN}\nPlease edit the given image based on the folllowing instruction: {sources['txt']}.",
+                    },
+                    {"from": "gpt", "value": LLAVA_IMAGE_TOKEN},
+                ]
+                image_files = sources["image"]
+                ref_image_files = sources["ref_image"]
+            else:
+                raise ValueError("Unknown source type. Please check the 'type' in 'sources'.")
+
+            # 2. Load images (handles bytes/PIL images from dataset)
+            if not isinstance(image_files, list):
+                image_files = [image_files]
+            if not isinstance(ref_image_files, list) and ref_image_files is not None:
+                ref_image_files = [ref_image_files]
+            images, ref_images = [], []
+
+            # Load target image(s)
+            if task_type == "T2I" or task_type == "I2I":
+                for img in image_files:
+                    try:
+                        # Ensure image is in RGB format
+                        if isinstance(img, bytes): img = Image.open(io.BytesIO(img))
+                        img = img.convert("RGB")
+                        images.append(img)
+                    except Exception as e:
+                        print(f"Error opening image: {e}. Skipping sample.")
+                        images = None; break
+            
+            # Load reference image(s) for I2I
+            if task_type == "I2I" and images is not None:
+                for img in ref_image_files:
+                    try:
+                        if isinstance(img, bytes): img = Image.open(io.BytesIO(img))
+                        img = img.convert("RGB")
+                        ref_images.append(img)
+                    except Exception as e:
+                        print(f"Error opening ref image: {e}. Skipping sample.")
+                        ref_images = None; break
+            
+            if images is None or (task_type == "I2I" and ref_images is None):
+                i = random.randint(0, len(self.list_data_dict) - 1)
+                continue # Skip to the next random index if image loading failed
+
+            # 3. Apply preprocessing to conversation and tokenize
+            sources = copy.deepcopy([sources["conversations"]])
+            sources, inst_type = preprocess_multimodal(sources, self.data_args)
+            sources = copy.deepcopy(llava_to_openai(sources))
+                        
+            all_input_ids = []
+            all_mm_token_type_ids = []
+            all_labels = []
+            all_pixel_values = []
+            all_image_grid_thw = []
+            all_second_gird = []
+
+            roles = {"human": "user", "gpt": "assistant"}
+            # all_dino_encoded_values = []
+
+            # Qwen2-VL uses a default system message so I've added this.
+            if len(SYSTEM_MESSAGE) > 0:
+                system_message = f"{DEFAULT_IM_START_TOKEN}system\n{SYSTEM_MESSAGE}\n{DEFAULT_IM_END_TOKEN}\n"
+                system_message_input_ids = processor.tokenizer(system_message, add_special_tokens=False, return_tensors='pt')['input_ids']
+                system_labels = torch.full_like(system_message_input_ids, IGNORE_INDEX) 
+                
+                all_input_ids.append(system_message_input_ids.squeeze(0))
+                all_mm_token_type_ids.append(
+                    torch.zeros(system_message_input_ids.shape[1], dtype=torch.long)
+                )
+                all_labels.append(system_labels.squeeze(0))
+
+            # import ipdb; ipdb.set_trace()
+            for _, j in enumerate(range(0, len(sources), 2)):
+                
+                if j >= 2:
+                    break
+                
+                user_input = sources[j]
+                gpt_response = sources[j + 1]
+                                
+                if self.stage_1_step <= self.cur_step < self.stage_2_step:
+                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    gpt_response = f"{gpt_response['content']}"
+                    if DEFAULT_IMAGE_TOKEN in user_input:
+                        gpt_response = get_comt_data_in_response(gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
+                    gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # print(f"\033[92m gpt_response: {gpt_response}\033[0m")
+                else:
+                    # user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                    # gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    # gpt_response = replace_pad_with_anchor_tokens(gpt_response)
+                    import random
+                    xxx = random.randint(0, 5)
+                    if xxx == 0:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
+                    else:
+                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
+                        gpt_response = f"{gpt_response['content']}"
+                        if DEFAULT_IMAGE_TOKEN in user_input:
+                            # INSERT_YOUR_CODE
+                            total = len(self.anchor_tokens)
+                            if total == 0:
+                                selected_anchor_token_nums = []
+                                selected_anchor_tokens = []
+                                selected_anchor_task_names = []
+                            else:
+                                x = random.randint(1, total)
+                                idxs = sorted(random.sample(range(total), x)) if x > 0 else []
+                                selected_anchor_token_nums = [self.anchor_token_nums[i] for i in idxs]
+                                selected_anchor_tokens = [self.anchor_tokens[i] for i in idxs]
+                                selected_anchor_task_names = [self.anchor_task_names[i] for i in idxs]
+                            gpt_response = get_comt_data_in_response(gpt_response, selected_anchor_token_nums, selected_anchor_tokens, selected_anchor_task_names)
+                        gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
+                
+                prompt_input_ids = processor.tokenizer(user_input, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+                all_mm_token_type_ids.append(
+                    torch.zeros(prompt_input_ids.shape[1], dtype=torch.long)
+                )
+
+                response_input_ids = processor.tokenizer(gpt_response, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']
+                all_mm_token_type_ids.append(
+                    torch.zeros(response_input_ids.shape[1], dtype=torch.long)
+                )
+
+                input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
+                labels = torch.cat(
+                    [
+                        torch.tensor([IGNORE_INDEX] * len(prompt_input_ids[0])),  
+                        response_input_ids.squeeze(0),
+                    ],
+                    dim=0,
+                )
+
+                all_input_ids.append(input_ids)
+                all_labels.append(labels)
+            
+            input_ids = torch.cat(all_input_ids, dim=0).to(torch.long)
+            labels = torch.cat(all_labels, dim=0).to(torch.long)
+            mm_token_type_ids = torch.cat(all_mm_token_type_ids, dim=0).to(torch.long)
+
+            if mm_token_type_ids.shape[0] != input_ids.shape[0]:
+                raise ValueError(
+                    f"mm_token_type_ids length {mm_token_type_ids.shape[0]} != input_ids {input_ids.shape[0]}"
+                )
+
+            # eos_token_id = processor.tokenizer.convert_tokens_to_ids(DEFAULT_IM_END_TOKEN)
+            # input_ids, labels = truncate_sequence(input_ids, labels, self.max_length, eos_token_id)
+
+            attention_mask = (input_ids > -1000000).to(torch.long)
+            
+            data_dict = dict(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+                mm_token_type_ids=mm_token_type_ids,
+                image_files=image_files,
+                task_types=task_type
+            )
+            # 4. Process images using VAE/Vision Tower processors
+            if inst_type == "gen": # T2I or I2I (where the LLM output involves generation)
+                # Target image for VAE (latents for diffusion)
+                images = [img.resize((self.data_args.image_size, self.data_args.image_size), 
+                                    resample=PIL.Image.Resampling.BICUBIC) for img in images]
+                data_dict["gen_image"] = self.data_args.gen_image_processor.preprocess(images, return_tensors="pt")["pixel_values"]
+
+                # I2I reference image (for diffusion conditioning)
+                if task_type == "I2I":
+                    data_dict["gen_ref_image"] = self.data_args.gen_image_processor.preprocess(ref_images, return_tensors="pt")["pixel_values"]
+                    # ref_mask is 1 for I2I
+                    b = data_dict["gen_ref_image"].shape[0]
+                    data_dict["ref_mask"] = torch.ones((b, 1, 1, 1), dtype=data_dict["gen_image"].dtype, device=data_dict["gen_image"].device)
+                # else: # T2I uses a masked/zeroed reference
+                #     data_dict["gen_ref_image"] = torch.zeros_like(data_dict["gen_image"])
+                #     b = data_dict["gen_ref_image"].shape[0]
+                #     # ref_mask is 0 for T2I (no reference image conditioning)
+                #     data_dict["ref_mask"] = torch.zeros((b, 1, 1, 1), dtype=data_dict["gen_image"].dtype, device=data_dict["gen_image"].device)
+        
+            if task_type == "I2I":      # I2I (where the LLM input involves image grounding)
+                # Underspecified image (resized for LLM context tokens)
+                resized_images = [transform_und_images(img) for img in ref_images]
+                image_inputs = self.data_args.image_processor(resized_images, return_tensors="pt")
+                data_dict["und_image"] = image_inputs.pixel_values
+                data_dict["grid_thw"] = image_inputs.image_grid_thw # Grid info for the visual encoder
+
+            self.cur_step += 1
+            return data_dict
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """
+    Collate examples for supervised fine-tuning, primarily handling padding 
+    and appending the T2I/I2I query tokens to the sequence.
+    """
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels, ids, task_types = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels", "ids", "task_types"))
+        multi_input_ids = []
+        multi_labels = []
+        i_s_pos = []
+
+        # 1. Truncate and Append Image Query Tokens
+        for input_id, label, task_type in zip(input_ids, labels, task_types):
+            # Truncate to make space for the image query tokens (n_query)
+            # The final length will be truncated to model_max_length in step 2.
+            input_id = input_id[: self.tokenizer.model_max_length]
+            label = label[: self.tokenizer.model_max_length]
+    
+            # i_s_pos records the starting index of the appended image query tokens
+            i_s_pos.append(input_id.shape[0])
+            
+            # # Create the sequence of image query tokens
+            # img_id = torch.full((self.n_query,), T2I_TOKEN_IDX, dtype=input_id.dtype, device=input_id.device) if task_type == "T2I" else torch.full((self.n_query,), I2I_TOKEN_IDX, dtype=input_id.dtype, device=input_id.device)
+            
+            # # Append tokens to input_id and label
+            # input_id = torch.cat([input_id, img_id])
+            # img_label = torch.full((self.n_query,), T2I_TOKEN_IDX, dtype=label.dtype, device=label.device) if task_type == "T2I" else torch.full((self.n_query,), I2I_TOKEN_IDX, dtype=label.dtype, device=label.device)
+            # label = torch.cat([label, img_label])
+            
+            multi_input_ids.append(input_id)
+            multi_labels.append(label)
+
+        input_ids = multi_input_ids
+        labels = multi_labels
+
+        # 2. Pad Sequences
+        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX) 
+        
+        # Final truncation to max length
+        if input_ids.shape[1] > self.tokenizer.model_max_length:
+            rank0_print(f"Warning: input with length {input_ids.shape[1]} is longer than max length {self.tokenizer.model_max_length}. Truncating.")
+        input_ids = input_ids[:, : self.tokenizer.model_max_length]
+        labels = labels[:, : self.tokenizer.model_max_length]
+
+        batch = dict(
+            input_ids=input_ids,
+            labels=labels,
+            # Attention mask ignores the padding tokens
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id), 
+        )
+        # 3. Collate Image Data (VAE latents and contextual images)
+        batch_gen_images = []
+        batch_ref_gen_images = []
+        batch_und_images = []
+        batch_grid_thw = []
+        batch_ref_mask = []
+        # Collate gen_image (target VAE latents)
+        for instance in instances:
+            if "gen_image" in instance:
+                batch_gen_images.append(instance["gen_image"])
+
+        if len(batch_gen_images) > 0:
+            # Check for consistent shape and concatenate
+            if all(x is not None and y.shape == batch_gen_images[0][0].shape for x in batch_gen_images for y in x):
+                batch["gen_image"] = torch.cat([images for images in batch_gen_images], dim=0)
+            else:
+                batch["gen_image"] = batch_gen_images # Keep as list if shapes vary (unlikely for VAE)
+
+        # Collate gen_ref_image (reference VAE latents) and ref_mask
+        for instance in instances:
+            if "gen_ref_image" in instance:
+                batch_ref_gen_images.append(instance["gen_ref_image"])
+            if "ref_mask" in instance:
+                batch_ref_mask.append(instance["ref_mask"])
+
+        if len(batch_ref_gen_images) > 0:
+            if all(x is not None and y.shape == batch_ref_gen_images[0][0].shape for x in batch_ref_gen_images for y in x):
+                batch["gen_ref_image"] = torch.cat([images for images in batch_ref_gen_images], dim=0)
+            else:
+                batch["gen_ref_image"] = batch_ref_gen_images
+        if len(batch_ref_mask) > 0:
+            batch["ref_mask"] = torch.cat([images for images in batch_ref_mask], dim=0)
+
+        # Collate und_image (contextual image pixels) and grid_thw
+        for instance in instances:
+            if "und_image" in instance:
+                batch_und_images.append(instance["und_image"].unsqueeze(0))
+                batch_grid_thw.append(instance["grid_thw"]) 
+
+        if len(batch_und_images) > 0:
+            batch["und_image"] = torch.cat([images for images in batch_und_images], dim=0)
+            batch["grid_thw"] = torch.cat([images for images in batch_grid_thw], dim=0)
+        else:
+            batch["und_image"] = None
+            batch["grid_thw"] = None
+
+        # Final metadata
+        batch["ids"] = ids
+        batch["i_s_pos"] = i_s_pos # Start positions of the image query tokens
+        return batch
+
+def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args: DataArguments) -> Dict:
+    """
+    Creates the training dataset and data collator based on data arguments.
+    Args:
+        tokenizer (transformers.PreTrainedTokenizer): The model tokenizer.
+        data_args (DataArguments): Configuration for data loading.
+    Returns:
+        Dict: Dictionary containing 'train_dataset', 'eval_dataset', and 'data_collator'.
+    """
+
+    if data_args.data_type == "mix":
+        train_dataset = LazySupervisedMixDataset(tokenizer=tokenizer, data_path=data_args.data_path, data_args=data_args)
+    else:
+        raise ValueError(f"Unknown data type: {data_args.data_type}.")
+
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)

code/umm/runtime/models/blip3o/model/__init__.py

@@ -0,0 +1,42 @@
+from .language_model.blip3o_qwen import blip3oQwenForCausalLM, blip3oQwenConfig
+from .language_model.blip3o_qwen_inference import blip3oQwenForInferenceLM, blip3oQwenConfig
+from .language_model.fm_step_prob import sde_step_with_logprob
+from .language_model.blip3o_qwen_kontext import blip3oQwenKontextForCausalLM,blip3oQwenConfig
+from .language_model.blip3o_qwen_kontext_inference import blip3oQwenKontextForInferenceLM,blip3oQwenConfig
+from .language_model.blip3o_qwen_stage1 import blip3oAlignQwenForCausalLM,blip3oAlignQwenConfig
+from .language_model.blip3o_qwen_align_infer import blip3oAlignQwenForInferenceLM,blip3oAlignQwenConfig
+from .language_model.covt_qwen_stage2_inference import CoVTForConditionalGenerationInference,blip3oCOVTQwenConfig
+from .language_model.omni_und_gen import OmniGenUndForConditionalGeneration, QwenOmniGenUndConfig
+from .language_model.omni_und_gen_infer import OmniForConditionalInference,QwenOmniConfig
+
+# Stage-2 CoVT: lazy so `import models.blip3o.model.language_model.moe_qwen_stage3_flux` does not
+# execute covt_qwen_stage2.py (Stage-3 Flux uses moe_qwen_stage3_flux instead).
+def __getattr__(name: str):
+    if name == "CoVTForConditionalGeneration":
+        from .language_model.covt_qwen_stage2 import CoVTForConditionalGeneration
+
+        return CoVTForConditionalGeneration
+    if name == "blip3oCOVTQwenConfig":
+        from .language_model.covt_qwen_stage2 import blip3oCOVTQwenConfig
+
+        return blip3oCOVTQwenConfig
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "blip3oQwenForCausalLM",
+    "blip3oQwenConfig",
+    "blip3oQwenForInferenceLM",
+    "sde_step_with_logprob",
+    "blip3oQwenKontextForCausalLM",
+    "blip3oQwenKontextForInferenceLM",
+    "blip3oAlignQwenForCausalLM",
+    "blip3oAlignQwenConfig",
+    "blip3oAlignQwenForInferenceLM",
+    "CoVTForConditionalGeneration",
+    "blip3oCOVTQwenConfig",
+    "OmniGenUndForConditionalGeneration",
+    "QwenOmniGenUndConfig",
+    "OmniForConditionalInference",
+    "QwenOmniConfig",
+]

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/DA-2K.md

@@ -0,0 +1,51 @@
+# DA-2K Evaluation Benchmark
+
+## Introduction
+
+![DA-2K](assets/DA-2K.png)
+
+DA-2K is proposed in [Depth Anything V2](https://depth-anything-v2.github.io) to evaluate the relative depth estimation capability. It encompasses eight representative scenarios of `indoor`, `outdoor`, `non_real`, `transparent_reflective`, `adverse_style`, `aerial`, `underwater`, and `object`. It consists of 1K diverse high-quality images and 2K precise pair-wise relative depth annotations.
+
+Please refer to our [paper](https://arxiv.org/abs/2406.09414) for details in constructing this benchmark.
+
+
+## Usage
+
+Please first [download the benchmark](https://huggingface.co/datasets/depth-anything/DA-2K/tree/main).
+
+All annotations are stored in `annotations.json`. The annotation file is a JSON object where each key is the path to an image file, and the value is a list of annotations associated with that image. Each annotation describes two points and identifies which point is closer to the camera. The structure is detailed below:
+
+```
+{
+  "image_path": [
+    {
+      "point1": [h1, w1], # (vertical position, horizontal position)
+      "point2": [h2, w2], # (vertical position, horizontal position)
+      "closer_point": "point1" # we always set "point1" as the closer one
+    },
+    ...
+  ],
+  ...
+}
+```
+
+To visualize the annotations:
+```bash
+python visualize.py [--scene-type <type>]
+```
+
+**Options**
+- `--scene-type <type>` (optional): Specify the scene type (`indoor`, `outdoor`, `non_real`, `transparent_reflective`, `adverse_style`, `aerial`, `underwater`, and `object`). Skip this argument or set <type> as `""` to include all scene types.
+
+## Citation
+
+If you find this benchmark useful, please consider citing:
+
+```bibtex
+@article{depth_anything_v2,
+  title={Depth Anything V2},
+  author={Yang, Lihe and Kang, Bingyi and Huang, Zilong and Zhao, Zhen and Xu, Xiaogang and Feng, Jiashi and Zhao, Hengshuang},
+  journal={arXiv:2406.09414},
+  year={2024}
+}
+```

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/README.md

@@ -0,0 +1,201 @@
+<div align="center">
+<h1>Depth Anything V2</h1>
+
+[**Lihe Yang**](https://liheyoung.github.io/)<sup>1</sup> · [**Bingyi Kang**](https://bingykang.github.io/)<sup>2&dagger;</sup> · [**Zilong Huang**](http://speedinghzl.github.io/)<sup>2</sup>
+<br>
+[**Zhen Zhao**](http://zhaozhen.me/) · [**Xiaogang Xu**](https://xiaogang00.github.io/) · [**Jiashi Feng**](https://sites.google.com/site/jshfeng/)<sup>2</sup> · [**Hengshuang Zhao**](https://hszhao.github.io/)<sup>1*</sup>
+
+<sup>1</sup>HKU&emsp;&emsp;&emsp;<sup>2</sup>TikTok
+<br>
+&dagger;project lead&emsp;*corresponding author
+
+<a href="https://arxiv.org/abs/2406.09414"><img src='https://img.shields.io/badge/arXiv-Depth Anything V2-red' alt='Paper PDF'></a>
+<a href='https://depth-anything-v2.github.io'><img src='https://img.shields.io/badge/Project_Page-Depth Anything V2-green' alt='Project Page'></a>
+<a href='https://huggingface.co/spaces/depth-anything/Depth-Anything-V2'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Demo-blue'></a>
+<a href='https://huggingface.co/datasets/depth-anything/DA-2K'><img src='https://img.shields.io/badge/Benchmark-DA--2K-yellow' alt='Benchmark'></a>
+</div>
+
+This work presents Depth Anything V2. It significantly outperforms [V1](https://github.com/LiheYoung/Depth-Anything) in fine-grained details and robustness. Compared with SD-based models, it enjoys faster inference speed, fewer parameters, and higher depth accuracy.
+
+![teaser](assets/teaser.png)
+
+
+## News
+- **2025-01-22:** [Video Depth Anything](https://videodepthanything.github.io) has been released. It generates consistent depth maps for super-long videos (e.g., over 5 minutes).
+- **2024-12-22:** [Prompt Depth Anything](https://promptda.github.io/) has been released. It supports 4K resolution metric depth estimation when low-res LiDAR is used to prompt the DA models.
+- **2024-07-06:** Depth Anything V2 is supported in [Transformers](https://github.com/huggingface/transformers/). See the [instructions](https://huggingface.co/docs/transformers/main/en/model_doc/depth_anything_v2) for convenient usage.
+- **2024-06-25:** Depth Anything is integrated into [Apple Core ML Models](https://developer.apple.com/machine-learning/models/). See the instructions ([V1](https://huggingface.co/apple/coreml-depth-anything-small), [V2](https://huggingface.co/apple/coreml-depth-anything-v2-small)) for usage.
+- **2024-06-22:** We release [smaller metric depth models](https://github.com/DepthAnything/Depth-Anything-V2/tree/main/metric_depth#pre-trained-models) based on Depth-Anything-V2-Small and Base.
+- **2024-06-20:** Our repository and project page are flagged by GitHub and removed from the public for 6 days. Sorry for the inconvenience.
+- **2024-06-14:** Paper, project page, code, models, demo, and benchmark are all released.
+
+
+## Pre-trained Models
+
+We provide **four models** of varying scales for robust relative depth estimation:
+
+| Model | Params | Checkpoint |
+|:-|-:|:-:|
+| Depth-Anything-V2-Small | 24.8M | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Small/resolve/main/depth_anything_v2_vits.pth?download=true) |
+| Depth-Anything-V2-Base | 97.5M | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Base/resolve/main/depth_anything_v2_vitb.pth?download=true) |
+| Depth-Anything-V2-Large | 335.3M | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth?download=true) |
+| Depth-Anything-V2-Giant | 1.3B | Coming soon |
+
+
+## Usage
+
+### Prepraration
+
+```bash
+git clone https://github.com/DepthAnything/Depth-Anything-V2
+cd Depth-Anything-V2
+pip install -r requirements.txt
+```
+
+Download the checkpoints listed [here](#pre-trained-models) and put them under the `checkpoints` directory.
+
+### Use our models
+```python
+import cv2
+import torch
+
+from depth_anything_v2.dpt import DepthAnythingV2
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+
+model_configs = {
+    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
+    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
+    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
+    'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
+}
+
+encoder = 'vitl' # or 'vits', 'vitb', 'vitg'
+
+model = DepthAnythingV2(**model_configs[encoder])
+model.load_state_dict(torch.load(f'checkpoints/depth_anything_v2_{encoder}.pth', map_location='cpu'))
+model = model.to(DEVICE).eval()
+
+raw_img = cv2.imread('your/image/path')
+depth = model.infer_image(raw_img) # HxW raw depth map in numpy
+```
+
+If you do not want to clone this repository, you can also load our models through [Transformers](https://github.com/huggingface/transformers/). Below is a simple code snippet. Please refer to the [official page](https://huggingface.co/docs/transformers/main/en/model_doc/depth_anything_v2) for more details.
+
+- Note 1: Make sure you can connect to Hugging Face and have installed the latest Transformers.
+- Note 2: Due to the [upsampling difference](https://github.com/huggingface/transformers/pull/31522#issuecomment-2184123463) between OpenCV (we used) and Pillow (HF used), predictions may differ slightly. So you are more recommended to use our models through the way introduced above.
+```python
+from transformers import pipeline
+from PIL import Image
+
+pipe = pipeline(task="depth-estimation", model="depth-anything/Depth-Anything-V2-Small-hf")
+image = Image.open('your/image/path')
+depth = pipe(image)["depth"]
+```
+
+### Running script on *images*
+
+```bash
+python run.py \
+  --encoder <vits | vitb | vitl | vitg> \
+  --img-path <path> --outdir <outdir> \
+  [--input-size <size>] [--pred-only] [--grayscale]
+```
+Options:
+- `--img-path`: You can either 1) point it to an image directory storing all interested images, 2) point it to a single image, or 3) point it to a text file storing all image paths.
+- `--input-size` (optional): By default, we use input size `518` for model inference. ***You can increase the size for even more fine-grained results.***
+- `--pred-only` (optional): Only save the predicted depth map, without raw image.
+- `--grayscale` (optional): Save the grayscale depth map, without applying color palette.
+
+For example:
+```bash
+python run.py --encoder vitl --img-path assets/examples --outdir depth_vis
+```
+
+### Running script on *videos*
+
+```bash
+python run_video.py \
+  --encoder <vits | vitb | vitl | vitg> \
+  --video-path assets/examples_video --outdir video_depth_vis \
+  [--input-size <size>] [--pred-only] [--grayscale]
+```
+
+***Our larger model has better temporal consistency on videos.***
+
+### Gradio demo
+
+To use our gradio demo locally:
+
+```bash
+python app.py
+```
+
+You can also try our [online demo](https://huggingface.co/spaces/Depth-Anything/Depth-Anything-V2).
+
+***Note: Compared to V1, we have made a minor modification to the DINOv2-DPT architecture (originating from this [issue](https://github.com/LiheYoung/Depth-Anything/issues/81)).*** In V1, we *unintentionally* used features from the last four layers of DINOv2 for decoding. In V2, we use [intermediate features](https://github.com/DepthAnything/Depth-Anything-V2/blob/2cbc36a8ce2cec41d38ee51153f112e87c8e42d8/depth_anything_v2/dpt.py#L164-L169) instead. Although this modification did not improve details or accuracy, we decided to follow this common practice.
+
+
+## Fine-tuned to Metric Depth Estimation
+
+Please refer to [metric depth estimation](./metric_depth).
+
+
+## DA-2K Evaluation Benchmark
+
+Please refer to [DA-2K benchmark](./DA-2K.md).
+
+
+## Community Support
+
+**We sincerely appreciate all the community support for our Depth Anything series. Thank you a lot!**
+
+- Apple Core ML:
+    - https://developer.apple.com/machine-learning/models
+    - https://huggingface.co/apple/coreml-depth-anything-v2-small
+    - https://huggingface.co/apple/coreml-depth-anything-small
+- Transformers:
+    - https://huggingface.co/docs/transformers/main/en/model_doc/depth_anything_v2
+    - https://huggingface.co/docs/transformers/main/en/model_doc/depth_anything
+- TensorRT:
+    - https://github.com/spacewalk01/depth-anything-tensorrt
+    - https://github.com/zhujiajian98/Depth-Anythingv2-TensorRT-python
+- ONNX: https://github.com/fabio-sim/Depth-Anything-ONNX
+- ComfyUI: https://github.com/kijai/ComfyUI-DepthAnythingV2
+- Transformers.js (real-time depth in web): https://huggingface.co/spaces/Xenova/webgpu-realtime-depth-estimation
+- Android:
+  - https://github.com/shubham0204/Depth-Anything-Android
+  - https://github.com/FeiGeChuanShu/ncnn-android-depth_anything
+
+
+## Acknowledgement
+
+We are sincerely grateful to the awesome Hugging Face team ([@Pedro Cuenca](https://huggingface.co/pcuenq), [@Niels Rogge](https://huggingface.co/nielsr), [@Merve Noyan](https://huggingface.co/merve), [@Amy Roberts](https://huggingface.co/amyeroberts), et al.) for their huge efforts in supporting our models in Transformers and Apple Core ML.
+
+We also thank the [DINOv2](https://github.com/facebookresearch/dinov2) team for contributing such impressive models to our community.
+
+
+## LICENSE
+
+Depth-Anything-V2-Small model is under the Apache-2.0 license. Depth-Anything-V2-Base/Large/Giant models are under the CC-BY-NC-4.0 license.
+
+
+## Citation
+
+If you find this project useful, please consider citing:
+
+```bibtex
+@article{depth_anything_v2,
+  title={Depth Anything V2},
+  author={Yang, Lihe and Kang, Bingyi and Huang, Zilong and Zhao, Zhen and Xu, Xiaogang and Feng, Jiashi and Zhao, Hengshuang},
+  journal={arXiv:2406.09414},
+  year={2024}
+}
+
+@inproceedings{depth_anything_v1,
+  title={Depth Anything: Unleashing the Power of Large-Scale Unlabeled Data}, 
+  author={Yang, Lihe and Kang, Bingyi and Huang, Zilong and Xu, Xiaogang and Feng, Jiashi and Zhao, Hengshuang},
+  booktitle={CVPR},
+  year={2024}
+}
+```

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/app.py

@@ -0,0 +1,88 @@
+import glob
+import gradio as gr
+import matplotlib
+import numpy as np
+from PIL import Image
+import torch
+import tempfile
+from gradio_imageslider import ImageSlider
+
+from depth_anything_v2.dpt import DepthAnythingV2
+
+css = """
+#img-display-container {
+    max-height: 100vh;
+}
+#img-display-input {
+    max-height: 80vh;
+}
+#img-display-output {
+    max-height: 80vh;
+}
+#download {
+    height: 62px;
+}
+"""
+DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+model_configs = {
+    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
+    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
+    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
+    'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
+}
+encoder = 'vitl'
+model = DepthAnythingV2(**model_configs[encoder])
+state_dict = torch.load(f'checkpoints/depth_anything_v2_{encoder}.pth', map_location="cpu")
+model.load_state_dict(state_dict)
+model = model.to(DEVICE).eval()
+
+title = "# Depth Anything V2"
+description = """Official demo for **Depth Anything V2**.
+Please refer to our [paper](https://arxiv.org/abs/2406.09414), [project page](https://depth-anything-v2.github.io), or [github](https://github.com/DepthAnything/Depth-Anything-V2) for more details."""
+
+def predict_depth(image):
+    return model.infer_image(image)
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+    gr.Markdown("### Depth Prediction demo")
+
+    with gr.Row():
+        input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
+        depth_image_slider = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5)
+    submit = gr.Button(value="Compute Depth")
+    gray_depth_file = gr.File(label="Grayscale depth map", elem_id="download",)
+    raw_file = gr.File(label="16-bit raw output (can be considered as disparity)", elem_id="download",)
+
+    cmap = matplotlib.colormaps.get_cmap('Spectral_r')
+
+    def on_submit(image):
+        original_image = image.copy()
+
+        h, w = image.shape[:2]
+
+        depth = predict_depth(image[:, :, ::-1])
+
+        raw_depth = Image.fromarray(depth.astype('uint16'))
+        tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+        raw_depth.save(tmp_raw_depth.name)
+
+        depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
+        depth = depth.astype(np.uint8)
+        colored_depth = (cmap(depth)[:, :, :3] * 255).astype(np.uint8)
+
+        gray_depth = Image.fromarray(depth)
+        tmp_gray_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+        gray_depth.save(tmp_gray_depth.name)
+
+        return [(original_image, colored_depth), tmp_gray_depth.name, tmp_raw_depth.name]
+
+    submit.click(on_submit, inputs=[input_image], outputs=[depth_image_slider, gray_depth_file, raw_file])
+
+    example_files = glob.glob('assets/examples/*')
+    examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=[depth_image_slider, gray_depth_file, raw_file], fn=on_submit)
+
+
+if __name__ == '__main__':
+    demo.queue().launch()

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2.py

@@ -0,0 +1,415 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+
+from .dinov2_layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
+
+
+logger = logging.getLogger("dinov2")
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    def forward(self, x):
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        ffn_bias=True,
+        proj_bias=True,
+        drop_path_rate=0.0,
+        drop_path_uniform=False,
+        init_values=None,  # for layerscale: None or 0 => no layerscale
+        embed_layer=PatchEmbed,
+        act_layer=nn.GELU,
+        block_fn=Block,
+        ffn_layer="mlp",
+        block_chunks=1,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            proj_bias (bool): enable bias for proj in attn if True
+            ffn_bias (bool): enable bias for ffn if True
+            drop_path_rate (float): stochastic depth rate
+            drop_path_uniform (bool): apply uniform drop rate across blocks
+            weight_init (str): weight init scheme
+            init_values (float): layer-scale init values
+            embed_layer (nn.Module): patch embedding layer
+            act_layer (nn.Module): MLP activation layer
+            block_fn (nn.Module): transformer block class
+            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+            num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+        """
+        super().__init__()
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.interpolate_antialias = interpolate_antialias
+        self.interpolate_offset = interpolate_offset
+
+        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        assert num_register_tokens >= 0
+        self.register_tokens = (
+            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+        )
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == "mlp":
+            logger.info("using MLP layer as FFN")
+            ffn_layer = Mlp
+        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+            logger.info("using SwiGLU layer as FFN")
+            ffn_layer = SwiGLUFFNFused
+        elif ffn_layer == "identity":
+            logger.info("using Identity layer as FFN")
+
+            def f(*args, **kwargs):
+                return nn.Identity()
+
+            ffn_layer = f
+        else:
+            raise NotImplementedError
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+        self.init_weights()
+
+    def init_weights(self):
+        trunc_normal_(self.pos_embed, std=0.02)
+        nn.init.normal_(self.cls_token, std=1e-6)
+        if self.register_tokens is not None:
+            nn.init.normal_(self.register_tokens, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        previous_dtype = x.dtype
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        pos_embed = self.pos_embed.float()
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        # DINOv2 with register modify the interpolate_offset from 0.1 to 0.0
+        w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
+        # w0, h0 = w0 + 0.1, h0 + 0.1
+        
+        sqrt_N = math.sqrt(N)
+        sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
+            scale_factor=(sx, sy),
+            # (int(w0), int(h0)), # to solve the upsampling shape issue
+            mode="bicubic",
+            antialias=self.interpolate_antialias
+        )
+        
+        assert int(w0) == patch_pos_embed.shape[-2]
+        assert int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        if self.register_tokens is not None:
+            x = torch.cat(
+                (
+                    x[:, :1],
+                    self.register_tokens.expand(x.shape[0], -1, -1),
+                    x[:, 1:],
+                ),
+                dim=1,
+            )
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        return {
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+    def _get_intermediate_layers_not_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm=True
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
+        if reshape:
+            B, _, w, h = x.shape
+            outputs = [
+                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+    def forward(self, *args, is_training=False, **kwargs):
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
+    """
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def DINOv2(model_name):
+    model_zoo = {
+        "vits": vit_small, 
+        "vitb": vit_base, 
+        "vitl": vit_large, 
+        "vitg": vit_giant2
+    }
+    
+    return model_zoo[model_name](
+        img_size=518,
+        patch_size=14,
+        init_values=1.0,
+        ffn_layer="mlp" if model_name != "vitg" else "swiglufused",
+        block_chunks=0,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1
+    )

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
+from .block import NestedTensorBlock
+from .attention import MemEffAttention

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/attention.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+
+from torch import Tensor
+from torch import nn
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+    from xformers.ops import memory_efficient_attention, unbind, fmha
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    logger.debug("xFormers not available (optional; DINOv2 falls back without it).")
+    XFORMERS_AVAILABLE = False
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            assert attn_bias is None, "xFormers is required for nested tensors usage"
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+        

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/block.py

@@ -0,0 +1,252 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+from typing import Callable, List, Any, Tuple, Dict
+
+import torch
+from torch import nn, Tensor
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+    from xformers.ops import fmha
+    from xformers.ops import scaled_index_add, index_select_cat
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    logger.debug("xFormers not available (optional; DINOv2 falls back without it).")
+    XFORMERS_AVAILABLE = False
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: Tensor) -> Tensor:
+        def attn_residual_func(x: Tensor) -> Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: Tensor) -> Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: Tensor,
+    residual_func: Callable[[Tensor], Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[Tensor],
+    residual_func: Callable[[Tensor, Any], Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+class NestedTensorBlock(Block):
+    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/drop_path.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/layer_scale.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch import nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/mlp.py

@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        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.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/patch_embed.py

@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+from typing import Callable, Optional, Tuple, Union
+
+from torch import Tensor
+import torch.nn as nn
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dinov2_layers/swiglu_ffn.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+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: Tensor) -> Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+try:
+    from xformers.ops import SwiGLU
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    SwiGLU = SwiGLUFFN
+    XFORMERS_AVAILABLE = False
+
+
+class SwiGLUFFNFused(SwiGLU):
+    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:
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/dpt.py

@@ -0,0 +1,221 @@
+import cv2
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms import Compose
+
+from .dinov2 import DINOv2
+from .util.blocks import FeatureFusionBlock, _make_scratch
+from .util.transform import Resize, NormalizeImage, PrepareForNet
+
+
+def _make_fusion_block(features, use_bn, size=None):
+    return FeatureFusionBlock(
+        features,
+        nn.ReLU(False),
+        deconv=False,
+        bn=use_bn,
+        expand=False,
+        align_corners=True,
+        size=size,
+    )
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_feature, out_feature):
+        super().__init__()
+        
+        self.conv_block = nn.Sequential(
+            nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_feature),
+            nn.ReLU(True)
+        )
+    
+    def forward(self, x):
+        return self.conv_block(x)
+
+
+class DPTHead(nn.Module):
+    def __init__(
+        self, 
+        in_channels, 
+        features=256, 
+        use_bn=False, 
+        out_channels=[256, 512, 1024, 1024], 
+        use_clstoken=False
+    ):
+        super(DPTHead, self).__init__()
+        
+        self.use_clstoken = use_clstoken
+        
+        self.projects = nn.ModuleList([
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channel,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+            ) for out_channel in out_channels
+        ])
+        
+        self.resize_layers = nn.ModuleList([
+            nn.ConvTranspose2d(
+                in_channels=out_channels[0],
+                out_channels=out_channels[0],
+                kernel_size=4,
+                stride=4,
+                padding=0),
+            nn.ConvTranspose2d(
+                in_channels=out_channels[1],
+                out_channels=out_channels[1],
+                kernel_size=2,
+                stride=2,
+                padding=0),
+            nn.Identity(),
+            nn.Conv2d(
+                in_channels=out_channels[3],
+                out_channels=out_channels[3],
+                kernel_size=3,
+                stride=2,
+                padding=1)
+        ])
+        
+        if use_clstoken:
+            self.readout_projects = nn.ModuleList()
+            for _ in range(len(self.projects)):
+                self.readout_projects.append(
+                    nn.Sequential(
+                        nn.Linear(2 * in_channels, in_channels),
+                        nn.GELU()))
+        
+        self.scratch = _make_scratch(
+            out_channels,
+            features,
+            groups=1,
+            expand=False,
+        )
+        
+        self.scratch.stem_transpose = None
+        
+        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
+        
+        head_features_1 = features
+        head_features_2 = 32
+        
+        self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
+        self.scratch.output_conv2 = nn.Sequential(
+            nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(True),
+            nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
+            nn.ReLU(True),
+            nn.Identity(),
+        )
+    
+    def forward(self, out_features, patch_h, patch_w):
+        out = []
+        for i, x in enumerate(out_features):
+            if self.use_clstoken:
+                x, cls_token = x[0], x[1]
+                readout = cls_token.unsqueeze(1).expand_as(x)
+                x = self.readout_projects[i](torch.cat((x, readout), -1))
+            else:
+                x = x[0]
+            
+            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+            
+            x = self.projects[i](x)
+            x = self.resize_layers[i](x)
+            
+            out.append(x)
+        
+        layer_1, layer_2, layer_3, layer_4 = out
+        
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+        
+        path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])        
+        path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
+        path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
+        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
+        
+        out = self.scratch.output_conv1(path_1)
+        out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
+        out = self.scratch.output_conv2(out)
+        
+        return out
+
+
+class DepthAnythingV2(nn.Module):
+    def __init__(
+        self, 
+        encoder='vitl', 
+        features=256, 
+        out_channels=[256, 512, 1024, 1024], 
+        use_bn=False, 
+        use_clstoken=False
+    ):
+        super(DepthAnythingV2, self).__init__()
+        
+        self.intermediate_layer_idx = {
+            'vits': [2, 5, 8, 11],
+            'vitb': [2, 5, 8, 11], 
+            'vitl': [4, 11, 17, 23], 
+            'vitg': [9, 19, 29, 39]
+        }
+        
+        self.encoder = encoder
+        self.pretrained = DINOv2(model_name=encoder)
+        
+        self.depth_head = DPTHead(self.pretrained.embed_dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken)
+    
+    def forward(self, x):
+        patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14
+        
+        features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], return_class_token=True)
+        
+        depth = self.depth_head(features, patch_h, patch_w)
+        depth = F.relu(depth)
+        
+        return depth.squeeze(1)
+    
+    @torch.no_grad()
+    def infer_image(self, raw_image, input_size=518):
+        image, (h, w) = self.image2tensor(raw_image, input_size)
+        
+        depth = self.forward(image)
+        
+        depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0]
+        
+        return depth.cpu().numpy()
+    
+    def image2tensor(self, raw_image, input_size=518):        
+        transform = Compose([
+            Resize(
+                width=input_size,
+                height=input_size,
+                resize_target=False,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=14,
+                resize_method='lower_bound',
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            PrepareForNet(),
+        ])
+        
+        h, w = raw_image.shape[:2]
+        
+        image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0
+        
+        image = transform({'image': image})['image']
+        image = torch.from_numpy(image).unsqueeze(0)
+        
+        DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+        image = image.to(DEVICE)
+        
+        return image, (h, w)

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/util/blocks.py

@@ -0,0 +1,148 @@
+import torch.nn as nn
+
+
+def _make_scratch(in_shape, out_shape, groups=1, expand=False):
+    scratch = nn.Module()
+
+    out_shape1 = out_shape
+    out_shape2 = out_shape
+    out_shape3 = out_shape
+    if len(in_shape) >= 4:
+        out_shape4 = out_shape
+
+    if expand:
+        out_shape1 = out_shape
+        out_shape2 = out_shape * 2
+        out_shape3 = out_shape * 4
+        if len(in_shape) >= 4:
+            out_shape4 = out_shape * 8
+
+    scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
+    scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
+    scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
+    if len(in_shape) >= 4:
+        scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
+
+    return scratch
+
+
+class ResidualConvUnit(nn.Module):
+    """Residual convolution module.
+    """
+
+    def __init__(self, features, activation, bn):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.bn = bn
+
+        self.groups=1
+
+        self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+        
+        self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+
+        if self.bn == True:
+            self.bn1 = nn.BatchNorm2d(features)
+            self.bn2 = nn.BatchNorm2d(features)
+
+        self.activation = activation
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: output
+        """
+        
+        out = self.activation(x)
+        out = self.conv1(out)
+        if self.bn == True:
+            out = self.bn1(out)
+       
+        out = self.activation(out)
+        out = self.conv2(out)
+        if self.bn == True:
+            out = self.bn2(out)
+
+        if self.groups > 1:
+            out = self.conv_merge(out)
+
+        return self.skip_add.add(out, x)
+
+
+class FeatureFusionBlock(nn.Module):
+    """Feature fusion block.
+    """
+
+    def __init__(
+        self, 
+        features, 
+        activation, 
+        deconv=False, 
+        bn=False, 
+        expand=False, 
+        align_corners=True,
+        size=None
+    ):
+        """Init.
+        
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock, self).__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+
+        self.groups=1
+
+        self.expand = expand
+        out_features = features
+        if self.expand == True:
+            out_features = features // 2
+        
+        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
+
+        self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
+        self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
+        
+        self.skip_add = nn.quantized.FloatFunctional()
+
+        self.size=size
+
+    def forward(self, *xs, size=None):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if len(xs) == 2:
+            res = self.resConfUnit1(xs[1])
+            output = self.skip_add.add(output, res)
+
+        output = self.resConfUnit2(output)
+
+        if (size is None) and (self.size is None):
+            modifier = {"scale_factor": 2}
+        elif size is None:
+            modifier = {"size": self.size}
+        else:
+            modifier = {"size": size}
+
+        output = nn.functional.interpolate(output, **modifier, mode="bilinear", align_corners=self.align_corners)
+        
+        output = self.out_conv(output)
+
+        return output

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/depth_anything_v2/util/transform.py

@@ -0,0 +1,158 @@
+import numpy as np
+import cv2
+
+
+class Resize(object):
+    """Resize sample to given size (width, height).
+    """
+
+    def __init__(
+        self,
+        width,
+        height,
+        resize_target=True,
+        keep_aspect_ratio=False,
+        ensure_multiple_of=1,
+        resize_method="lower_bound",
+        image_interpolation_method=cv2.INTER_AREA,
+    ):
+        """Init.
+
+        Args:
+            width (int): desired output width
+            height (int): desired output height
+            resize_target (bool, optional):
+                True: Resize the full sample (image, mask, target).
+                False: Resize image only.
+                Defaults to True.
+            keep_aspect_ratio (bool, optional):
+                True: Keep the aspect ratio of the input sample.
+                Output sample might not have the given width and height, and
+                resize behaviour depends on the parameter 'resize_method'.
+                Defaults to False.
+            ensure_multiple_of (int, optional):
+                Output width and height is constrained to be multiple of this parameter.
+                Defaults to 1.
+            resize_method (str, optional):
+                "lower_bound": Output will be at least as large as the given size.
+                "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
+                "minimal": Scale as least as possible.  (Output size might be smaller than given size.)
+                Defaults to "lower_bound".
+        """
+        self.__width = width
+        self.__height = height
+
+        self.__resize_target = resize_target
+        self.__keep_aspect_ratio = keep_aspect_ratio
+        self.__multiple_of = ensure_multiple_of
+        self.__resize_method = resize_method
+        self.__image_interpolation_method = image_interpolation_method
+
+    def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
+        y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if max_val is not None and y > max_val:
+            y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if y < min_val:
+            y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        return y
+
+    def get_size(self, width, height):
+        # determine new height and width
+        scale_height = self.__height / height
+        scale_width = self.__width / width
+
+        if self.__keep_aspect_ratio:
+            if self.__resize_method == "lower_bound":
+                # scale such that output size is lower bound
+                if scale_width > scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "upper_bound":
+                # scale such that output size is upper bound
+                if scale_width < scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "minimal":
+                # scale as least as possbile
+                if abs(1 - scale_width) < abs(1 - scale_height):
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            else:
+                raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        if self.__resize_method == "lower_bound":
+            new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
+            new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
+        elif self.__resize_method == "upper_bound":
+            new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
+            new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
+        elif self.__resize_method == "minimal":
+            new_height = self.constrain_to_multiple_of(scale_height * height)
+            new_width = self.constrain_to_multiple_of(scale_width * width)
+        else:
+            raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        return (new_width, new_height)
+
+    def __call__(self, sample):
+        width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
+        
+        # resize sample
+        sample["image"] = cv2.resize(sample["image"], (width, height), interpolation=self.__image_interpolation_method)
+
+        if self.__resize_target:
+            if "depth" in sample:
+                sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
+                
+            if "mask" in sample:
+                sample["mask"] = cv2.resize(sample["mask"].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST)
+        
+        return sample
+
+
+class NormalizeImage(object):
+    """Normlize image by given mean and std.
+    """
+
+    def __init__(self, mean, std):
+        self.__mean = mean
+        self.__std = std
+
+    def __call__(self, sample):
+        sample["image"] = (sample["image"] - self.__mean) / self.__std
+
+        return sample
+
+
+class PrepareForNet(object):
+    """Prepare sample for usage as network input.
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, sample):
+        image = np.transpose(sample["image"], (2, 0, 1))
+        sample["image"] = np.ascontiguousarray(image).astype(np.float32)
+
+        if "depth" in sample:
+            depth = sample["depth"].astype(np.float32)
+            sample["depth"] = np.ascontiguousarray(depth)
+        
+        if "mask" in sample:
+            sample["mask"] = sample["mask"].astype(np.float32)
+            sample["mask"] = np.ascontiguousarray(sample["mask"])
+        
+        return sample

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/README.md

@@ -0,0 +1,114 @@
+# Depth Anything V2 for Metric Depth Estimation
+
+![teaser](./assets/compare_zoedepth.png)
+
+We here provide a simple codebase to fine-tune our Depth Anything V2 pre-trained encoder for metric depth estimation. Built on our powerful encoder, we use a simple DPT head to regress the depth. We fine-tune our pre-trained encoder on synthetic Hypersim / Virtual KITTI datasets for indoor / outdoor metric depth estimation, respectively.
+
+
+# Pre-trained Models
+
+We provide **six metric depth models** of three scales for indoor and outdoor scenes, respectively.
+
+| Base Model | Params | Indoor (Hypersim) | Outdoor (Virtual KITTI 2) |
+|:-|-:|:-:|:-:|
+| Depth-Anything-V2-Small | 24.8M | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-Hypersim-Small/resolve/main/depth_anything_v2_metric_hypersim_vits.pth?download=true) | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-VKITTI-Small/resolve/main/depth_anything_v2_metric_vkitti_vits.pth?download=true) |
+| Depth-Anything-V2-Base | 97.5M | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-Hypersim-Base/resolve/main/depth_anything_v2_metric_hypersim_vitb.pth?download=true) | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-VKITTI-Base/resolve/main/depth_anything_v2_metric_vkitti_vitb.pth?download=true) |
+| Depth-Anything-V2-Large | 335.3M | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-Hypersim-Large/resolve/main/depth_anything_v2_metric_hypersim_vitl.pth?download=true) | [Download](https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-VKITTI-Large/resolve/main/depth_anything_v2_metric_vkitti_vitl.pth?download=true) |
+
+*We recommend to first try our larger models (if computational cost is affordable) and the indoor version.*
+
+## Usage
+
+### Prepraration
+
+```bash
+git clone https://github.com/DepthAnything/Depth-Anything-V2
+cd Depth-Anything-V2/metric_depth
+pip install -r requirements.txt
+```
+
+Download the checkpoints listed [here](#pre-trained-models) and put them under the `checkpoints` directory.
+
+### Use our models
+```python
+import cv2
+import torch
+
+from depth_anything_v2.dpt import DepthAnythingV2
+
+model_configs = {
+    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
+    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
+    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}
+}
+
+encoder = 'vitl' # or 'vits', 'vitb'
+dataset = 'hypersim' # 'hypersim' for indoor model, 'vkitti' for outdoor model
+max_depth = 20 # 20 for indoor model, 80 for outdoor model
+
+model = DepthAnythingV2(**{**model_configs[encoder], 'max_depth': max_depth})
+model.load_state_dict(torch.load(f'checkpoints/depth_anything_v2_metric_{dataset}_{encoder}.pth', map_location='cpu'))
+model.eval()
+
+raw_img = cv2.imread('your/image/path')
+depth = model.infer_image(raw_img) # HxW depth map in meters in numpy
+```
+
+### Running script on images
+
+Here, we take the `vitl` encoder as an example. You can also use `vitb` or `vits` encoders.
+
+```bash
+# indoor scenes
+python run.py \
+  --encoder vitl \
+  --load-from checkpoints/depth_anything_v2_metric_hypersim_vitl.pth \
+  --max-depth 20 \
+  --img-path <path> --outdir <outdir> [--input-size <size>] [--save-numpy]
+
+# outdoor scenes
+python run.py \
+  --encoder vitl \
+  --load-from checkpoints/depth_anything_v2_metric_vkitti_vitl.pth \
+  --max-depth 80 \
+  --img-path <path> --outdir <outdir> [--input-size <size>] [--save-numpy]
+```
+
+### Project 2D images to point clouds:
+
+```bash
+python depth_to_pointcloud.py \
+  --encoder vitl \
+  --load-from checkpoints/depth_anything_v2_metric_hypersim_vitl.pth \
+  --max-depth 20 \
+  --img-path <path> --outdir <outdir>
+```
+
+### Reproduce training
+
+Please first prepare the [Hypersim](https://github.com/apple/ml-hypersim) and [Virtual KITTI 2](https://europe.naverlabs.com/research/computer-vision/proxy-virtual-worlds-vkitti-2/) datasets. Then:
+
+```bash
+bash dist_train.sh
+```
+
+
+## Citation
+
+If you find this project useful, please consider citing:
+
+```bibtex
+@article{depth_anything_v2,
+  title={Depth Anything V2},
+  author={Yang, Lihe and Kang, Bingyi and Huang, Zilong and Zhao, Zhen and Xu, Xiaogang and Feng, Jiashi and Zhao, Hengshuang},
+  journal={arXiv:2406.09414},
+  year={2024}
+}
+
+@inproceedings{depth_anything_v1,
+  title={Depth Anything: Unleashing the Power of Large-Scale Unlabeled Data}, 
+  author={Yang, Lihe and Kang, Bingyi and Huang, Zilong and Xu, Xiaogang and Feng, Jiashi and Zhao, Hengshuang},
+  booktitle={CVPR},
+  year={2024}
+}
+```

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/dataset/hypersim.py

@@ -0,0 +1,74 @@
+import cv2
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision.transforms import Compose
+
+from dataset.transform import Resize, NormalizeImage, PrepareForNet, Crop
+
+
+def hypersim_distance_to_depth(npyDistance):
+    intWidth, intHeight, fltFocal = 1024, 768, 886.81
+
+    npyImageplaneX = np.linspace((-0.5 * intWidth) + 0.5, (0.5 * intWidth) - 0.5, intWidth).reshape(
+        1, intWidth).repeat(intHeight, 0).astype(np.float32)[:, :, None]
+    npyImageplaneY = np.linspace((-0.5 * intHeight) + 0.5, (0.5 * intHeight) - 0.5,
+                                 intHeight).reshape(intHeight, 1).repeat(intWidth, 1).astype(np.float32)[:, :, None]
+    npyImageplaneZ = np.full([intHeight, intWidth, 1], fltFocal, np.float32)
+    npyImageplane = np.concatenate(
+        [npyImageplaneX, npyImageplaneY, npyImageplaneZ], 2)
+
+    npyDepth = npyDistance / np.linalg.norm(npyImageplane, 2, 2) * fltFocal
+    return npyDepth
+
+
+class Hypersim(Dataset):
+    def __init__(self, filelist_path, mode, size=(518, 518)):
+        
+        self.mode = mode
+        self.size = size
+        
+        with open(filelist_path, 'r') as f:
+            self.filelist = f.read().splitlines()
+        
+        net_w, net_h = size
+        self.transform = Compose([
+            Resize(
+                width=net_w,
+                height=net_h,
+                resize_target=True if mode == 'train' else False,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=14,
+                resize_method='lower_bound',
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            PrepareForNet(),
+        ] + ([Crop(size[0])] if self.mode == 'train' else []))
+        
+    def __getitem__(self, item):
+        img_path = self.filelist[item].split(' ')[0]
+        depth_path = self.filelist[item].split(' ')[1]
+        
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
+        
+        depth_fd = h5py.File(depth_path, "r")
+        distance_meters = np.array(depth_fd['dataset'])
+        depth = hypersim_distance_to_depth(distance_meters)
+        
+        sample = self.transform({'image': image, 'depth': depth})
+
+        sample['image'] = torch.from_numpy(sample['image'])
+        sample['depth'] = torch.from_numpy(sample['depth'])
+        
+        sample['valid_mask'] = (torch.isnan(sample['depth']) == 0)
+        sample['depth'][sample['valid_mask'] == 0] = 0
+        
+        sample['image_path'] = self.filelist[item].split(' ')[0]
+        
+        return sample
+
+    def __len__(self):
+        return len(self.filelist)

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/dataset/kitti.py

@@ -0,0 +1,57 @@
+import cv2
+import torch
+from torch.utils.data import Dataset
+from torchvision.transforms import Compose
+
+from dataset.transform import Resize, NormalizeImage, PrepareForNet
+
+
+class KITTI(Dataset):
+    def __init__(self, filelist_path, mode, size=(518, 518)):
+        if mode != 'val':
+            raise NotImplementedError
+        
+        self.mode = mode
+        self.size = size
+        
+        with open(filelist_path, 'r') as f:
+            self.filelist = f.read().splitlines()
+        
+        net_w, net_h = size
+        self.transform = Compose([
+            Resize(
+                width=net_w,
+                height=net_h,
+                resize_target=True if mode == 'train' else False,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=14,
+                resize_method='lower_bound',
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            PrepareForNet(),
+        ])
+    
+    def __getitem__(self, item):
+        img_path = self.filelist[item].split(' ')[0]
+        depth_path = self.filelist[item].split(' ')[1]
+        
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
+        
+        depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED).astype('float32')
+        
+        sample = self.transform({'image': image, 'depth': depth})
+        
+        sample['image'] = torch.from_numpy(sample['image'])
+        sample['depth'] = torch.from_numpy(sample['depth'])
+        sample['depth'] = sample['depth'] / 256.0  # convert in meters
+        
+        sample['valid_mask'] = sample['depth'] > 0
+        
+        sample['image_path'] = self.filelist[item].split(' ')[0]
+        
+        return sample
+
+    def __len__(self):
+        return len(self.filelist)

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/dataset/transform.py

@@ -0,0 +1,277 @@
+import cv2
+import math
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
+    """Rezise the sample to ensure the given size. Keeps aspect ratio.
+
+    Args:
+        sample (dict): sample
+        size (tuple): image size
+
+    Returns:
+        tuple: new size
+    """
+    shape = list(sample["disparity"].shape)
+
+    if shape[0] >= size[0] and shape[1] >= size[1]:
+        return sample
+
+    scale = [0, 0]
+    scale[0] = size[0] / shape[0]
+    scale[1] = size[1] / shape[1]
+
+    scale = max(scale)
+
+    shape[0] = math.ceil(scale * shape[0])
+    shape[1] = math.ceil(scale * shape[1])
+
+    # resize
+    sample["image"] = cv2.resize(
+        sample["image"], tuple(shape[::-1]), interpolation=image_interpolation_method
+    )
+
+    sample["disparity"] = cv2.resize(
+        sample["disparity"], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST
+    )
+    sample["mask"] = cv2.resize(
+        sample["mask"].astype(np.float32),
+        tuple(shape[::-1]),
+        interpolation=cv2.INTER_NEAREST,
+    )
+    sample["mask"] = sample["mask"].astype(bool)
+
+    return tuple(shape)
+
+
+class Resize(object):
+    """Resize sample to given size (width, height).
+    """
+
+    def __init__(
+        self,
+        width,
+        height,
+        resize_target=True,
+        keep_aspect_ratio=False,
+        ensure_multiple_of=1,
+        resize_method="lower_bound",
+        image_interpolation_method=cv2.INTER_AREA,
+    ):
+        """Init.
+
+        Args:
+            width (int): desired output width
+            height (int): desired output height
+            resize_target (bool, optional):
+                True: Resize the full sample (image, mask, target).
+                False: Resize image only.
+                Defaults to True.
+            keep_aspect_ratio (bool, optional):
+                True: Keep the aspect ratio of the input sample.
+                Output sample might not have the given width and height, and
+                resize behaviour depends on the parameter 'resize_method'.
+                Defaults to False.
+            ensure_multiple_of (int, optional):
+                Output width and height is constrained to be multiple of this parameter.
+                Defaults to 1.
+            resize_method (str, optional):
+                "lower_bound": Output will be at least as large as the given size.
+                "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
+                "minimal": Scale as least as possible.  (Output size might be smaller than given size.)
+                Defaults to "lower_bound".
+        """
+        self.__width = width
+        self.__height = height
+
+        self.__resize_target = resize_target
+        self.__keep_aspect_ratio = keep_aspect_ratio
+        self.__multiple_of = ensure_multiple_of
+        self.__resize_method = resize_method
+        self.__image_interpolation_method = image_interpolation_method
+
+    def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
+        y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if max_val is not None and y > max_val:
+            y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if y < min_val:
+            y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        return y
+
+    def get_size(self, width, height):
+        # determine new height and width
+        scale_height = self.__height / height
+        scale_width = self.__width / width
+
+        if self.__keep_aspect_ratio:
+            if self.__resize_method == "lower_bound":
+                # scale such that output size is lower bound
+                if scale_width > scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "upper_bound":
+                # scale such that output size is upper bound
+                if scale_width < scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "minimal":
+                # scale as least as possbile
+                if abs(1 - scale_width) < abs(1 - scale_height):
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            else:
+                raise ValueError(
+                    f"resize_method {self.__resize_method} not implemented"
+                )
+
+        if self.__resize_method == "lower_bound":
+            new_height = self.constrain_to_multiple_of(
+                scale_height * height, min_val=self.__height
+            )
+            new_width = self.constrain_to_multiple_of(
+                scale_width * width, min_val=self.__width
+            )
+        elif self.__resize_method == "upper_bound":
+            new_height = self.constrain_to_multiple_of(
+                scale_height * height, max_val=self.__height
+            )
+            new_width = self.constrain_to_multiple_of(
+                scale_width * width, max_val=self.__width
+            )
+        elif self.__resize_method == "minimal":
+            new_height = self.constrain_to_multiple_of(scale_height * height)
+            new_width = self.constrain_to_multiple_of(scale_width * width)
+        else:
+            raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        return (new_width, new_height)
+
+    def __call__(self, sample):
+        width, height = self.get_size(
+            sample["image"].shape[1], sample["image"].shape[0]
+        )
+
+        # resize sample
+        sample["image"] = cv2.resize(
+            sample["image"],
+            (width, height),
+            interpolation=self.__image_interpolation_method,
+        )
+
+        if self.__resize_target:
+            if "disparity" in sample:
+                sample["disparity"] = cv2.resize(
+                    sample["disparity"],
+                    (width, height),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+
+            if "depth" in sample:
+                sample["depth"] = cv2.resize(
+                    sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST
+                )
+
+            if "semseg_mask" in sample:
+                # sample["semseg_mask"] = cv2.resize(
+                #     sample["semseg_mask"], (width, height), interpolation=cv2.INTER_NEAREST
+                # )
+                sample["semseg_mask"] = F.interpolate(torch.from_numpy(sample["semseg_mask"]).float()[None, None, ...], (height, width), mode='nearest').numpy()[0, 0]
+                
+            if "mask" in sample:
+                sample["mask"] = cv2.resize(
+                    sample["mask"].astype(np.float32),
+                    (width, height),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+                # sample["mask"] = sample["mask"].astype(bool)
+
+        # print(sample['image'].shape, sample['depth'].shape)
+        return sample
+
+
+class NormalizeImage(object):
+    """Normlize image by given mean and std.
+    """
+
+    def __init__(self, mean, std):
+        self.__mean = mean
+        self.__std = std
+
+    def __call__(self, sample):
+        sample["image"] = (sample["image"] - self.__mean) / self.__std
+
+        return sample
+
+
+class PrepareForNet(object):
+    """Prepare sample for usage as network input.
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, sample):
+        image = np.transpose(sample["image"], (2, 0, 1))
+        sample["image"] = np.ascontiguousarray(image).astype(np.float32)
+
+        if "mask" in sample:
+            sample["mask"] = sample["mask"].astype(np.float32)
+            sample["mask"] = np.ascontiguousarray(sample["mask"])
+        
+        if "depth" in sample:
+            depth = sample["depth"].astype(np.float32)
+            sample["depth"] = np.ascontiguousarray(depth)
+            
+        if "semseg_mask" in sample:
+            sample["semseg_mask"] = sample["semseg_mask"].astype(np.float32)
+            sample["semseg_mask"] = np.ascontiguousarray(sample["semseg_mask"])
+
+        return sample
+
+
+class Crop(object):
+    """Crop sample for batch-wise training. Image is of shape CxHxW
+    """
+
+    def __init__(self, size):
+        if isinstance(size, int):
+            self.size = (size, size)
+        else:
+            self.size = size
+
+    def __call__(self, sample):
+        h, w = sample['image'].shape[-2:]
+        assert h >= self.size[0] and w >= self.size[1], 'Wrong size'
+        
+        h_start = np.random.randint(0, h - self.size[0] + 1)
+        w_start = np.random.randint(0, w - self.size[1] + 1)
+        h_end = h_start + self.size[0]
+        w_end = w_start + self.size[1]
+        
+        sample['image'] = sample['image'][:, h_start: h_end, w_start: w_end]
+        
+        if "depth" in sample:
+            sample["depth"] = sample["depth"][h_start: h_end, w_start: w_end]
+        
+        if "mask" in sample:
+            sample["mask"] = sample["mask"][h_start: h_end, w_start: w_end]
+            
+        if "semseg_mask" in sample:
+            sample["semseg_mask"] = sample["semseg_mask"][h_start: h_end, w_start: w_end]
+            
+        return sample

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/dataset/vkitti2.py

@@ -0,0 +1,54 @@
+import cv2
+import torch
+from torch.utils.data import Dataset
+from torchvision.transforms import Compose
+
+from dataset.transform import Resize, NormalizeImage, PrepareForNet, Crop
+
+
+class VKITTI2(Dataset):
+    def __init__(self, filelist_path, mode, size=(518, 518)):
+        
+        self.mode = mode
+        self.size = size
+        
+        with open(filelist_path, 'r') as f:
+            self.filelist = f.read().splitlines()
+        
+        net_w, net_h = size
+        self.transform = Compose([
+            Resize(
+                width=net_w,
+                height=net_h,
+                resize_target=True if mode == 'train' else False,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=14,
+                resize_method='lower_bound',
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            PrepareForNet(),
+        ] + ([Crop(size[0])] if self.mode == 'train' else []))
+    
+    def __getitem__(self, item):
+        img_path = self.filelist[item].split(' ')[0]
+        depth_path = self.filelist[item].split(' ')[1]
+        
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
+        
+        depth = cv2.imread(depth_path, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) / 100.0  # cm to m
+        
+        sample = self.transform({'image': image, 'depth': depth})
+
+        sample['image'] = torch.from_numpy(sample['image'])
+        sample['depth'] = torch.from_numpy(sample['depth'])
+        
+        sample['valid_mask'] = (sample['depth'] <= 80)
+        
+        sample['image_path'] = self.filelist[item].split(' ')[0]
+        
+        return sample
+
+    def __len__(self):
+        return len(self.filelist)

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2.py

@@ -0,0 +1,415 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+
+from .dinov2_layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
+
+
+logger = logging.getLogger("dinov2")
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    def forward(self, x):
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        ffn_bias=True,
+        proj_bias=True,
+        drop_path_rate=0.0,
+        drop_path_uniform=False,
+        init_values=None,  # for layerscale: None or 0 => no layerscale
+        embed_layer=PatchEmbed,
+        act_layer=nn.GELU,
+        block_fn=Block,
+        ffn_layer="mlp",
+        block_chunks=1,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            proj_bias (bool): enable bias for proj in attn if True
+            ffn_bias (bool): enable bias for ffn if True
+            drop_path_rate (float): stochastic depth rate
+            drop_path_uniform (bool): apply uniform drop rate across blocks
+            weight_init (str): weight init scheme
+            init_values (float): layer-scale init values
+            embed_layer (nn.Module): patch embedding layer
+            act_layer (nn.Module): MLP activation layer
+            block_fn (nn.Module): transformer block class
+            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+            num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+        """
+        super().__init__()
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.interpolate_antialias = interpolate_antialias
+        self.interpolate_offset = interpolate_offset
+
+        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        assert num_register_tokens >= 0
+        self.register_tokens = (
+            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+        )
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == "mlp":
+            logger.info("using MLP layer as FFN")
+            ffn_layer = Mlp
+        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+            logger.info("using SwiGLU layer as FFN")
+            ffn_layer = SwiGLUFFNFused
+        elif ffn_layer == "identity":
+            logger.info("using Identity layer as FFN")
+
+            def f(*args, **kwargs):
+                return nn.Identity()
+
+            ffn_layer = f
+        else:
+            raise NotImplementedError
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+        self.init_weights()
+
+    def init_weights(self):
+        trunc_normal_(self.pos_embed, std=0.02)
+        nn.init.normal_(self.cls_token, std=1e-6)
+        if self.register_tokens is not None:
+            nn.init.normal_(self.register_tokens, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        previous_dtype = x.dtype
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        pos_embed = self.pos_embed.float()
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        # DINOv2 with register modify the interpolate_offset from 0.1 to 0.0
+        w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
+        # w0, h0 = w0 + 0.1, h0 + 0.1
+        
+        sqrt_N = math.sqrt(N)
+        sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
+            scale_factor=(sx, sy),
+            # (int(w0), int(h0)), # to solve the upsampling shape issue
+            mode="bicubic",
+            antialias=self.interpolate_antialias
+        )
+        
+        assert int(w0) == patch_pos_embed.shape[-2]
+        assert int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        if self.register_tokens is not None:
+            x = torch.cat(
+                (
+                    x[:, :1],
+                    self.register_tokens.expand(x.shape[0], -1, -1),
+                    x[:, 1:],
+                ),
+                dim=1,
+            )
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        return {
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+    def _get_intermediate_layers_not_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm=True
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
+        if reshape:
+            B, _, w, h = x.shape
+            outputs = [
+                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+    def forward(self, *args, is_training=False, **kwargs):
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
+    """
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def DINOv2(model_name):
+    model_zoo = {
+        "vits": vit_small, 
+        "vitb": vit_base, 
+        "vitl": vit_large, 
+        "vitg": vit_giant2
+    }
+    
+    return model_zoo[model_name](
+        img_size=518,
+        patch_size=14,
+        init_values=1.0,
+        ffn_layer="mlp" if model_name != "vitg" else "swiglufused",
+        block_chunks=0,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1
+    )

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
+from .block import NestedTensorBlock
+from .attention import MemEffAttention

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/attention.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+
+from torch import Tensor
+from torch import nn
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+    from xformers.ops import memory_efficient_attention, unbind, fmha
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    logger.debug("xFormers not available (optional; DINOv2 falls back without it).")
+    XFORMERS_AVAILABLE = False
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            assert attn_bias is None, "xFormers is required for nested tensors usage"
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+        

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/block.py

@@ -0,0 +1,252 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+from typing import Callable, List, Any, Tuple, Dict
+
+import torch
+from torch import nn, Tensor
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+    from xformers.ops import fmha
+    from xformers.ops import scaled_index_add, index_select_cat
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    logger.debug("xFormers not available (optional; DINOv2 falls back without it).")
+    XFORMERS_AVAILABLE = False
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: Tensor) -> Tensor:
+        def attn_residual_func(x: Tensor) -> Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: Tensor) -> Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: Tensor,
+    residual_func: Callable[[Tensor], Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[Tensor],
+    residual_func: Callable[[Tensor, Any], Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+class NestedTensorBlock(Block):
+    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/drop_path.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/layer_scale.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch import nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/mlp.py

@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        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.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/patch_embed.py

@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+from typing import Callable, Optional, Tuple, Union
+
+from torch import Tensor
+import torch.nn as nn
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops

code/umm/runtime/models/blip3o/model/anchors/DepthAnything/metric_depth/depth_anything_v2/dinov2_layers/swiglu_ffn.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+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: Tensor) -> Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+try:
+    from xformers.ops import SwiGLU
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    SwiGLU = SwiGLUFFN
+    XFORMERS_AVAILABLE = False
+
+
+class SwiGLUFFNFused(SwiGLU):
+    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:
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )