Deploy MODUS 3-tab any-to-any demo (ZeroGPU)

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,33 @@
 ---
 title: MODUS
-emoji: 💻
-colorFrom: purple
+emoji: 🎨
+colorFrom: indigo
 colorTo: purple
 sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.13'
+sdk_version: 4.44.1
 app_file: app.py
 pinned: false
+short_description: MODUS any-to-any multimodal demo (16 aligned modalities)
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MODUS — any-to-any multimodal demo
+
+Three-tab Gradio demo for the MODUS 16-modality any-to-any model:
+
+1. **Any-to-Any** — one condition modality (image or caption) → any set of target
+   modalities, shown in a gallery (4M-style).
+2. **Chained** — condition → intermediate (bridge) → target, showing both.
+3. **Representation Analysis** — RGB→Depth/Normal conditioned on
+   {ViT only, VAE only, ViT+VAE}, side by side.
+
+## Setup (Space Settings → Variables and secrets)
+- **Hardware:** ZeroGPU.
+- **`HF_TOKEN`** (secret): a *read* token with access to the gated weights repo
+  `mqye/modus-16mod-stage3`.
+- Optional **`MODUS_WEIGHTS_REPO`**: defaults to `mqye/modus-16mod-stage3`.
+
+The weights (bf16, ~30GB) + VAE + config + tokenizer are pulled once at startup
+via `snapshot_download`; the model is built on CPU (~5min) and moved to the
+ZeroGPU GPU per request. Inference config is baked in: modality config
+`conf/modalities/instruction_16mod_stage2.yaml`, `use_instruction` off for the
+generic modalities (seg/det/grounding keep their own instructions).

app.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+"""MODUS any-to-any demo — HuggingFace Space entrypoint (ZeroGPU).
+
+Thin ZeroGPU wrapper over the existing 3-tab demo backend (``demo_modus.py`` +
+``demo_my/*``). The model is loaded ONCE at startup on CPU (no GPU, no time
+limit); each inference call moves it to the ZeroGPU-provided GPU via the
+``@spaces.GPU`` decorator (moving ~30GB bf16 over PCIe is seconds, whereas the
+build+weight-load is ~5min and must NOT happen inside the GPU-time window).
+
+Space setup (Settings -> Variables and secrets):
+    HF_TOKEN   read token for the gated weights repo (below)
+Optional env (have sane defaults):
+    MODUS_WEIGHTS_REPO   gated HF model repo with the bf16 weights + config + VAE
+"""
+import os
+import sys
+
+try:
+    import torch
+    print(f"[app] torch={torch.__version__} python={sys.version.split()[0]}", flush=True)
+except Exception as _e:
+    print(f"[app] torch import failed: {_e}", flush=True)
+
+# ── Env MUST be set before importing the demo backend (it reads these at import) ─
+os.environ.setdefault("MODUS_NO_MEAN_RESIZING", "1")  # avoid gradio-BLAS deadlock
+os.environ.setdefault("MODUS_FORCE_SDPA_ATTN", "1")   # no flash-attn on the Space
+os.environ.setdefault("MODUS_DEMO_MODALITY_CONFIG",
+                      "conf/modalities/instruction_16mod_stage2.yaml")
+os.environ.setdefault("MODUS_DEMO_MODEL_NAME", "bagel_from_json")
+os.environ.setdefault("MODUS_DEMO_USE_EMA", "0")
+
+WEIGHTS_REPO = os.environ.get("MODUS_WEIGHTS_REPO", "mqye/modus-16mod-stage3")
+
+# ── Pull the gated weights once (model.safetensors + ae + config + tokenizer) ────
+from huggingface_hub import snapshot_download  # noqa: E402
+
+_weights_dir = snapshot_download(
+    repo_id=WEIGHTS_REPO,
+    repo_type="model",
+    token=os.environ.get("HF_TOKEN"),
+)
+# The snapshot dir holds BOTH the checkpoint (model.safetensors) and the base
+# config/tokenizer/VAE, so it serves as CHECKPOINT_PATH and MODEL_PATH at once.
+os.environ["MODUS_DEMO_CHECKPOINT"] = _weights_dir
+os.environ["BAGEL_MODEL_PATH"] = _weights_dir
+print(f"[app] weights ready at {_weights_dir}", flush=True)
+
+import spaces  # noqa: E402  (ZeroGPU)
+
+# diffusers 0.20 (imported by the fourm VQVAE feature tokenizers) does
+# `from huggingface_hub import cached_download`, which was removed in hub>=0.26
+# (the Space has 0.36). Shim it to hf_hub_download so `import diffusers` succeeds
+# and the dino/clip/imagebind tokenizers can load.
+import huggingface_hub as _hh  # noqa: E402
+if not hasattr(_hh, "cached_download"):
+    _hh.cached_download = _hh.hf_hub_download
+
+# Importing demo_modus builds the UI + backend and reads the env set above.
+import demo_modus  # noqa: E402
+
+
+# ── ZeroGPU: wrap the two inference entry points so each call gets a GPU slice ───
+# tab1/tab2 call demo_modus.run_task; tab3 calls demo_modus.run_representation_task.
+# The tab handlers resolve these by module-global name at call time, so replacing
+# the module attribute makes them use the GPU-wrapped versions.
+# duration = max GPU-seconds ZeroGPU reserves per call. The Chained tab runs TWO
+# generations inside a single run_task call (~60s+), so 60 is too tight ("GPU task
+# aborted"); 120 covers chained + cold-start model materialisation. (With PRO the
+# reservation size is a non-issue quota-wise.)
+_GPU_DURATION = int(os.environ.get("MODUS_GPU_DURATION", "120"))
+demo_modus.run_task = spaces.GPU(duration=_GPU_DURATION)(demo_modus.run_task)
+# Tab3 runs THREE generations (vit/vae/both) per click. Wrapping run_representation_task
+# would take three separate GPU acquisitions in one handler → the 3rd hits "Expired
+# ZeroGPU proxy token". Instead wrap the whole tab3 HANDLER so all three run inside a
+# single GPU session (one token, one reservation big enough for 3 gens).
+demo_modus.tab3_generate = spaces.GPU(duration=180)(demo_modus.tab3_generate)
+
+# Load the model ONCE at startup. The heavy CPU work (build arch + read 30GB
+# weights, ~5min) runs here, OUTSIDE any @spaces.GPU window; the model's .cuda()
+# moves are deferred by `spaces` and materialise on the first GPU call.
+try:
+    demo_modus.HOLDER.ensure_loaded()
+    print("[app] model loaded (CPU) at startup", flush=True)
+except Exception as e:  # surface in UI, retry lazily on first request
+    demo_modus.HOLDER.load_error = str(e)
+    print(f"[app] startup model load failed: {e}", flush=True)
+
+# Work around a gradio 4.44.1 bug: get_api_info() (called when rendering the main
+# "/" route) crashes with `TypeError: argument of type 'bool' is not iterable` when
+# a component's JSON schema contains a bool (additionalProperties: true/false).
+# Patch gradio_client's schema helpers to tolerate bool schemas.
+try:
+    import gradio_client.utils as _gcu
+
+    _orig_j2p = _gcu._json_schema_to_python_type
+
+    def _safe_j2p(schema, defs=None):
+        if isinstance(schema, bool):
+            return "Any"
+        return _orig_j2p(schema, defs)
+
+    _gcu._json_schema_to_python_type = _safe_j2p
+
+    _orig_get_type = _gcu.get_type
+
+    def _safe_get_type(schema):
+        if isinstance(schema, bool):
+            return "Any"
+        return _orig_get_type(schema)
+
+    _gcu.get_type = _safe_get_type
+except Exception as _e:
+    print(f"[app] gradio_client schema patch skipped: {_e}", flush=True)
+
+demo = demo_modus.build_ui()
+demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_api=False)

conf/modalities/bagel_stage3_exact3cond.yaml

@@ -0,0 +1,191 @@
+# Stage 3 modality config for the standalone BAGEL exact-3-condition project.
+# Intentionally duplicated from instruction_stage2.yaml so this project can
+# evolve independently without affecting MODUS experiments.
+
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    codebook_spatial_shape: [32, 32]
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    codebook_spatial_shape: [28, 28]
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    codebook_spatial_shape: [32, 32]
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0

conf/modalities/hunyuan_image_3.yaml

@@ -0,0 +1,31 @@
+# Modality configuration for HunyuanImage-3.0 training.
+#
+# Only the RGB image modality is needed for text-to-image generation.
+# Token keys (start_of_image / end_of_image) are populated by
+# data_utils.add_special_tokens(), which adds <|vision_start|> and
+# <|vision_end|> to the tokenizer when they are absent.
+# These serve as sequence delimiters; actual image patches are placed at
+# packed_vae_token_indexes and overwritten by the model's patch_embed.
+
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vae: true
+    represent_vit: true

conf/modalities/instruction.yaml

@@ -0,0 +1,201 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding   # checkpoint compat: attr is "grounding_pos_embed"
+    conditions: [rgb]
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0

conf/modalities/instruction_10modality_stage2.yaml

@@ -0,0 +1,142 @@
+# Stage 2 modality config — unconstrained conditions.
+#
+# Same modality definitions as instruction.yaml, but with ALL condition
+# restrictions removed so every modality can be conditioned on every other.
+
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    # conditions: unconstrained (any modality can condition caption)
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    # conditions: unconstrained
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding   # checkpoint compat: attr is "grounding_pos_embed"
+    # conditions: unconstrained
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    # conditions: unconstrained
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    # conditions: unconstrained
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  

conf/modalities/instruction_16mod_instr.yaml

@@ -0,0 +1,259 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_depth     # REPLACE: per-modality start (token already exists in base vocab)
+    start_token: "<|depth_start|>"
+    end_token_key: end_of_image          # shared end (image end is a structural terminator)
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_normal    # REPLACE: per-modality start (token already exists in base vocab)
+    start_token: "<|normal_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    # NOTE: det's learnable pos_embed (grounding_pos_embed, size 4) is REMOVED here.
+    # The forward pos_embed loop is keyed by token-id range, and cocodet reuses
+    # det's x1/y1/x2/y2 tokens; with the pos_embed on, a multi-box cocodet sample
+    # (>4 coord tokens) would index the size-4 embedding out of bounds. det is not
+    # trained in modus (only present for coord-token-id alignment), so dropping its
+    # pos_embed is safe. reweight also removed (no reweight for det/cocodet).
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_seg       # REPLACE: per-modality start (NEW token)
+    start_token: "<|seg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_canny     # REPLACE: per-modality start (NEW token)
+    start_token: "<|canny_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  # ── NEW: cocodet (Pix2seq detection). stage1 = rgb→cocodet only.
+  # Coords reuse det's x1/y1/x2/y2 tokens (dedup, 0 new); class = COCO-id-aligned
+  # tokens (~91 new). dispersed_code_tokens → CE is gathered over the
+  # non-contiguous {coords ∪ class ∪ end} set, not a contiguous range.
+  # No reweight, no learnable pos_embed (RoPE + token-id carry order/role).
+  - name: cocodet
+    id: 13
+    kind: codebook
+    start_token_key: start_of_cocodet
+    end_token_key: end_of_cocodet
+    start_token: "<|cocodet_start|>"
+    end_token: "<|cocodet_end|>"
+    represent_vae: true
+    dispersed_code_tokens: true
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|coco_cls_{i:02d}|>"
+        start: 0
+        end: 90
+    # Inference pipeline config.
+    # decode_method is 'cocodet' (NOT 'detection') — cocodet's output is
+    # coords+class per box ending on cocodet_end, a DIFFERENT format from the
+    # grounding/det coords-only decode (generate_detection_coordonly). Handled by
+    # the dedicated model.generate_cocodet + inferencer.gen_cocodet (which do NOT
+    # touch the grounding decode).
+    inference_decode_method: cocodet
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false
+
+  # ── NEW image modalities (stage2): SAM segmentation / edge. Mirror seg/canny
+  # (direct-PNG image, VIT+VAE, MSE). REPLACE: per-modality start token, shared
+  # end_of_image. These start tokens are NEW (created via start_token field).
+  - name: samseg
+    id: 14
+    kind: image
+    start_token_key: start_of_samseg
+    start_token: "<|samseg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: samedge
+    id: 15
+    kind: image
+    start_token_key: start_of_samedge
+    start_token: "<|samedge_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true

conf/modalities/instruction_16mod_stage1.yaml

@@ -0,0 +1,243 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    # NOTE: det's learnable pos_embed (grounding_pos_embed, size 4) is REMOVED here.
+    # The forward pos_embed loop is keyed by token-id range, and cocodet reuses
+    # det's x1/y1/x2/y2 tokens; with the pos_embed on, a multi-box cocodet sample
+    # (>4 coord tokens) would index the size-4 embedding out of bounds. det is not
+    # trained in modus (only present for coord-token-id alignment), so dropping its
+    # pos_embed is safe. reweight also removed (no reweight for det/cocodet).
+    conditions: [rgb]
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  # ── NEW: cocodet (Pix2seq detection). stage1 = rgb→cocodet only.
+  # Coords reuse det's x1/y1/x2/y2 tokens (dedup, 0 new); class = COCO-id-aligned
+  # tokens (~91 new). dispersed_code_tokens → CE is gathered over the
+  # non-contiguous {coords ∪ class ∪ end} set, not a contiguous range.
+  # No reweight, no learnable pos_embed (RoPE + token-id carry order/role).
+  - name: cocodet
+    id: 13
+    kind: codebook
+    start_token_key: start_of_cocodet
+    end_token_key: end_of_cocodet
+    start_token: "<|cocodet_start|>"
+    end_token: "<|cocodet_end|>"
+    represent_vae: true
+    conditions: [rgb]
+    dispersed_code_tokens: true
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|coco_cls_{i:02d}|>"
+        start: 0
+        end: 90
+    # Inference pipeline config.
+    # decode_method is 'cocodet' (NOT 'detection') — cocodet's output is
+    # coords+class per box ending on cocodet_end, a DIFFERENT format from the
+    # grounding/det coords-only decode (generate_detection_coordonly). Handled by
+    # the dedicated model.generate_cocodet + inferencer.gen_cocodet (which do NOT
+    # touch the grounding decode).
+    inference_decode_method: cocodet
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false

conf/modalities/instruction_16mod_stage1_rgbcond.yaml

@@ -0,0 +1,266 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_depth     # REPLACE: per-modality start (token already exists in base vocab)
+    start_token: "<|depth_start|>"
+    end_token_key: end_of_image          # shared end (image end is a structural terminator)
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_normal    # REPLACE: per-modality start (token already exists in base vocab)
+    start_token: "<|normal_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    conditions: ["rgb"]   # stage1 rgbcond: condition on rgb only (rgb->X warmup)
+    end_token_key: end_of_det
+    represent_vae: true
+    # NOTE: det's learnable pos_embed (grounding_pos_embed, size 4) is REMOVED here.
+    # The forward pos_embed loop is keyed by token-id range, and cocodet reuses
+    # det's x1/y1/x2/y2 tokens; with the pos_embed on, a multi-box cocodet sample
+    # (>4 coord tokens) would index the size-4 embedding out of bounds. det is not
+    # trained in modus (only present for coord-token-id alignment), so dropping its
+    # pos_embed is safe. reweight also removed (no reweight for det/cocodet).
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_seg       # REPLACE: per-modality start (NEW token)
+    start_token: "<|seg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_canny     # REPLACE: per-modality start (NEW token)
+    start_token: "<|canny_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    conditions: ["rgb"]   # stage1 rgbcond: condition on rgb only (rgb->X warmup)
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    conditions: ["rgb"]   # stage1 rgbcond: condition on rgb only (rgb->X warmup)
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    conditions: ["rgb"]   # stage1 rgbcond: condition on rgb only (rgb->X warmup)
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    conditions: ["rgb"]   # stage1 rgbcond: condition on rgb only (rgb->X warmup)
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    conditions: ["rgb"]   # stage1 rgbcond: condition on rgb only (rgb->X warmup)
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  # ── NEW: cocodet (Pix2seq detection). stage1 = rgb→cocodet only.
+  # Coords reuse det's x1/y1/x2/y2 tokens (dedup, 0 new); class = COCO-id-aligned
+  # tokens (~91 new). dispersed_code_tokens → CE is gathered over the
+  # non-contiguous {coords ∪ class ∪ end} set, not a contiguous range.
+  # No reweight, no learnable pos_embed (RoPE + token-id carry order/role).
+  - name: cocodet
+    id: 13
+    kind: codebook
+    start_token_key: start_of_cocodet
+    conditions: ["rgb"]   # stage1 rgbcond: condition on rgb only (rgb->X warmup)
+    end_token_key: end_of_cocodet
+    start_token: "<|cocodet_start|>"
+    end_token: "<|cocodet_end|>"
+    represent_vae: true
+    dispersed_code_tokens: true
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|coco_cls_{i:02d}|>"
+        start: 0
+        end: 90
+    # Inference pipeline config.
+    # decode_method is 'cocodet' (NOT 'detection') — cocodet's output is
+    # coords+class per box ending on cocodet_end, a DIFFERENT format from the
+    # grounding/det coords-only decode (generate_detection_coordonly). Handled by
+    # the dedicated model.generate_cocodet + inferencer.gen_cocodet (which do NOT
+    # touch the grounding decode).
+    inference_decode_method: cocodet
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false
+
+  # ── NEW image modalities (stage2): SAM segmentation / edge. Mirror seg/canny
+  # (direct-PNG image, VIT+VAE, MSE). REPLACE: per-modality start token, shared
+  # end_of_image. These start tokens are NEW (created via start_token field).
+  - name: samseg
+    id: 14
+    kind: image
+    start_token_key: start_of_samseg
+    start_token: "<|samseg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: samedge
+    id: 15
+    kind: image
+    start_token_key: start_of_samedge
+    start_token: "<|samedge_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true

conf/modalities/instruction_16mod_stage2.yaml

@@ -0,0 +1,272 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    inference_add_instruction: false   # stage2/3 trained with use_target_instruction=false (no generic instr)
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    inference_add_instruction: false   # any2caption (stage2) trained without a target instruction
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    inference_add_instruction: false   # generic [start X] not trained in stage2/3 (use_target_instruction=false)
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_depth     # REPLACE: per-modality start (token already exists in base vocab)
+    start_token: "<|depth_start|>"
+    end_token_key: end_of_image          # shared end (image end is a structural terminator)
+    represent_vit: true
+    represent_vae: true
+    inference_add_instruction: false
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_normal    # REPLACE: per-modality start (token already exists in base vocab)
+    start_token: "<|normal_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    inference_add_instruction: false
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    # NOTE: det's learnable pos_embed (grounding_pos_embed, size 4) is REMOVED here.
+    # The forward pos_embed loop is keyed by token-id range, and cocodet reuses
+    # det's x1/y1/x2/y2 tokens; with the pos_embed on, a multi-box cocodet sample
+    # (>4 coord tokens) would index the size-4 embedding out of bounds. det is not
+    # trained in modus (only present for coord-token-id alignment), so dropping its
+    # pos_embed is safe. reweight also removed (no reweight for det/cocodet).
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_seg       # REPLACE: per-modality start (NEW token)
+    start_token: "<|seg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_canny     # REPLACE: per-modality start (NEW token)
+    start_token: "<|canny_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    inference_add_instruction: false
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+    inference_add_instruction: false   # start_of_dino token triggers it; generic instr not in stage2/3
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+    inference_add_instruction: false
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+    inference_add_instruction: false
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+    inference_add_instruction: false
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+    inference_add_instruction: false
+
+  # ── NEW: cocodet (Pix2seq detection). stage1 = rgb→cocodet only.
+  # Coords reuse det's x1/y1/x2/y2 tokens (dedup, 0 new); class = COCO-id-aligned
+  # tokens (~91 new). dispersed_code_tokens → CE is gathered over the
+  # non-contiguous {coords ∪ class ∪ end} set, not a contiguous range.
+  # No reweight, no learnable pos_embed (RoPE + token-id carry order/role).
+  - name: cocodet
+    id: 13
+    kind: codebook
+    start_token_key: start_of_cocodet
+    end_token_key: end_of_cocodet
+    start_token: "<|cocodet_start|>"
+    end_token: "<|cocodet_end|>"
+    represent_vae: true
+    dispersed_code_tokens: true
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|coco_cls_{i:02d}|>"
+        start: 0
+        end: 90
+    # Inference pipeline config.
+    # decode_method is 'cocodet' (NOT 'detection') — cocodet's output is
+    # coords+class per box ending on cocodet_end, a DIFFERENT format from the
+    # grounding/det coords-only decode (generate_detection_coordonly). Handled by
+    # the dedicated model.generate_cocodet + inferencer.gen_cocodet (which do NOT
+    # touch the grounding decode).
+    inference_decode_method: cocodet
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false
+
+  # ── NEW image modalities (stage2): SAM segmentation / edge. Mirror seg/canny
+  # (direct-PNG image, VIT+VAE, MSE). REPLACE: per-modality start token, shared
+  # end_of_image. These start tokens are NEW (created via start_token field).
+  - name: samseg
+    id: 14
+    kind: image
+    start_token_key: start_of_samseg
+    start_token: "<|samseg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    inference_add_instruction: false
+
+  - name: samedge
+    id: 15
+    kind: image
+    start_token_key: start_of_samedge
+    start_token: "<|samedge_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    inference_add_instruction: false

conf/modalities/instruction_2d_only_stage1.yaml

@@ -0,0 +1,65 @@
+# 2D-only stage1 modality registry.
+#
+# Keeps text/caption and image-like modalities only. Drops DET/DINO/CLIP/
+# ImageBind codebook modalities so this run can isolate whether 1D/codebook CE
+# is the main source of Hunyuan T2I prior drift.
+
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, depth, normal, seg, canny]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, depth, normal, seg, canny]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, rgb, normal, seg, canny]
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, rgb, depth, seg, canny]
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, rgb, depth, normal, canny]
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, rgb, depth, normal, seg]

conf/modalities/instruction_9modality_stage1.yaml

@@ -0,0 +1,114 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, grounding, dino]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, grounding, dino]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding   # checkpoint compat: attr is "grounding_pos_embed"
+    conditions: [rgb]
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0

conf/modalities/instruction_9modality_stage3.yaml

@@ -0,0 +1,115 @@
+# Stage 3 modality config — 9 modalities (text + 8 task), unconstrained conditions.
+#
+# Mirrors instruction_9modality_stage1.yaml but drops every `conditions: [...]`
+# restriction so any modality can be conditioned on any other (matches the old
+# stage3 training run that was originally configured via CLI flags only).
+# det reweight + det/dino learnable pos-embeds are preserved.
+
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding   # checkpoint compat: attr is "grounding_pos_embed"
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0

conf/modalities/instruction_hunyuan_16mod_stage1.yaml

@@ -0,0 +1,257 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    # det's learnable pos_embed (size 4) + reweight are REMOVED so cocodet — which
+    # reuses det's x1/y1/x2/y2 tokens (coco_cls is cocodet's own) — can carry >4
+    # coord tokens per sample without indexing the size-4 pos_embed out of bounds.
+    # det is not a target here (only present for coord-token-id alignment);
+    # grounding reuses these tokens and is now RoPE-only for its bbox coords.
+    conditions: [rgb]
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  # ── cocodet (Pix2seq detection). Coords reuse det's x1/y1/x2/y2 tokens (dedup,
+  # 0 new); class = COCO-id-aligned tokens (~91 new). dispersed_code_tokens → CE
+  # is gathered over the non-contiguous {coords ∪ class ∪ end} set, not a
+  # contiguous range. No reweight, no learnable pos_embed (RoPE + token-id).
+  - name: cocodet
+    id: 13
+    kind: codebook
+    start_token_key: start_of_cocodet
+    end_token_key: end_of_cocodet
+    start_token: "<|cocodet_start|>"
+    end_token: "<|cocodet_end|>"
+    represent_vae: true
+    conditions: [rgb]   # rgb→cocodet, consistent with det/dino/clip in this config
+    dispersed_code_tokens: true
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|coco_cls_{i:02d}|>"
+        start: 0
+        end: 90
+    # Inference: dedicated model.generate_cocodet + inferencer.gen_cocodet
+    # (coords+class per box ending on cocodet_end; NOT the grounding/det decode).
+    inference_decode_method: cocodet
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false
+
+  # ── SAM segmentation / edge. Mirror seg/canny (direct-PNG image, ViT+VAE, MSE).
+  # Per-modality start token, shared end_of_image. Start tokens are NEW.
+  - name: samseg
+    id: 14
+    kind: image
+    start_token_key: start_of_samseg
+    start_token: "<|samseg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: samedge
+    id: 15
+    kind: image
+    start_token_key: start_of_samedge
+    start_token: "<|samedge_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true

conf/modalities/instruction_hunyuan_16mod_stage2.yaml

@@ -0,0 +1,250 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, grounding, dino, dinolocal, clip, imagebind, imagebindlocal]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    # det's learnable pos_embed (size 4) + reweight are REMOVED so cocodet — which
+    # reuses det's x1/y1/x2/y2 tokens (coco_cls is cocodet's own) — can carry >4
+    # coord tokens per sample without indexing the size-4 pos_embed out of bounds.
+    # det is not a target here (only present for coord-token-id alignment);
+    # grounding reuses these tokens and is now RoPE-only for its bbox coords.
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  # ── cocodet (Pix2seq detection). Coords reuse det's x1/y1/x2/y2 tokens (dedup,
+  # 0 new); class = COCO-id-aligned tokens (~91 new). dispersed_code_tokens → CE
+  # is gathered over the non-contiguous {coords ∪ class ∪ end} set, not a
+  # contiguous range. No reweight, no learnable pos_embed (RoPE + token-id).
+  - name: cocodet
+    id: 13
+    kind: codebook
+    start_token_key: start_of_cocodet
+    end_token_key: end_of_cocodet
+    start_token: "<|cocodet_start|>"
+    end_token: "<|cocodet_end|>"
+    represent_vae: true
+    dispersed_code_tokens: true
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|coco_cls_{i:02d}|>"
+        start: 0
+        end: 90
+    # Inference: dedicated model.generate_cocodet + inferencer.gen_cocodet
+    # (coords+class per box ending on cocodet_end; NOT the grounding/det decode).
+    inference_decode_method: cocodet
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false
+
+  # ── SAM segmentation / edge. Mirror seg/canny (direct-PNG image, ViT+VAE, MSE).
+  # Per-modality start token, shared end_of_image. Start tokens are NEW.
+  - name: samseg
+    id: 14
+    kind: image
+    start_token_key: start_of_samseg
+    start_token: "<|samseg_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: samedge
+    id: 15
+    kind: image
+    start_token_key: start_of_samedge
+    start_token: "<|samedge_start|>"
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true

conf/modalities/instruction_stage2.yaml

@@ -0,0 +1,204 @@
+# Stage 2 modality config — unconstrained conditions.
+#
+# Same modality definitions as instruction.yaml, but with ALL condition
+# restrictions removed so every modality can be conditioned on every other.
+
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    # conditions: unconstrained (any modality can condition caption)
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    # conditions: unconstrained
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding   # checkpoint compat: attr is "grounding_pos_embed"
+    # conditions: unconstrained
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    # conditions: unconstrained
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    # conditions: unconstrained
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    # conditions: unconstrained
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    # conditions: unconstrained
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    # conditions: unconstrained
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+

conf/modalities/instruction_t2i_only_stage1.yaml

@@ -0,0 +1,27 @@
+# v65 = T2I-only stage1 modality registry.
+# Keeps text + caption + rgb only. Drops every other modality.
+# Pair with `hunyuan_image_3_v65_t2i_only.yaml`.
+
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption]

conf/modalities/legacy.yaml

@@ -0,0 +1,38 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: start_of_caption
+    end_token_key: end_of_caption
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_depth
+    end_token_key: end_of_depth
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_normal
+    end_token_key: end_of_normal
+    represent_vit: true
+    represent_vae: true
+
+

conf/modalities/rebuttal_rgb2target.yaml

@@ -0,0 +1,116 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, grounding, dino]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, grounding, dino]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [rgb]
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [rgb]
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding   # checkpoint compat: attr is "grounding_pos_embed"
+    conditions: [rgb]
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0

conf/modalities/stage1_oversample_bbox2rgb_dinolocal2rgb.yaml

@@ -0,0 +1,203 @@
+modalities:
+  - name: text
+    id: 0
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+
+  - name: caption
+    id: 1
+    kind: text
+    start_token_key: bos_token_id
+    end_token_key: eos_token_id
+    represent_vae: true
+    conditions: [rgb, grounding, dinolocal]
+    condition_probs: [0.1, 0.45, 0.45]
+
+  - name: rgb
+    id: 2
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+    conditions: [caption, grounding, dinolocal]
+    condition_probs: [0.1, 0.45, 0.45]
+
+  - name: depth
+    id: 3
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: normal
+    id: 4
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: det
+    id: 5
+    kind: codebook
+    start_token_key: start_of_det
+    end_token_key: end_of_det
+    represent_vae: true
+    pos_embed_size: 4
+    apply_pos_embed_in_forward: true
+    pos_embed_name: grounding   # checkpoint compat: attr is "grounding_pos_embed"
+    conditions: [rgb]
+    loss:
+      reweight: true
+      reweight_min_w: 0.005
+    start_token: "<|det_start|>"
+    end_token: "<|det_end|>"
+    extra_tokens: ["<|box_start|>", "<|box_end|>"]
+    code_token_groups:
+      - token_format: "<|x1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y1_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|x2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|y2_{i:03d}|>"
+        start: 0
+        end: 999
+      - token_format: "<|score_{i:02d}|>"
+        start: 0
+        end: 99
+    # Inference pipeline config
+    inference_decode_method: detection
+    inference_max_tokens: 1000
+    inference_cfg_uncond: text
+    inference_add_instruction: false   # det uses start_of_det token, not text instruction
+
+  - name: seg
+    id: 6
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: canny
+    id: 7
+    kind: image
+    start_token_key: start_of_image
+    end_token_key: end_of_image
+    represent_vit: true
+    represent_vae: true
+
+  - name: dino
+    id: 8
+    kind: codebook
+    start_token_key: start_of_dino
+    end_token_key: end_of_dino
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|dino_start|>"
+    end_token: "<|dino_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dino_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: dino
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: dinolocal
+    id: 9
+    kind: codebook
+    start_token_key: start_of_dinolocal
+    end_token_key: end_of_dinolocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|dinolocal_start|>"
+    end_token: "<|dinolocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|dinolocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: dinolocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: clip
+    id: 10
+    kind: codebook
+    start_token_key: start_of_clip
+    end_token_key: end_of_clip
+    pos_embed_size: 784
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [28, 28]   # 784 tokens, 28×28 grid → 2D RoPE
+    start_token: "<|clip_start|>"
+    end_token: "<|clip_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|clip_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: clip
+    inference_max_tokens: 785
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebind
+    id: 11
+    kind: codebook
+    start_token_key: start_of_imagebind
+    end_token_key: end_of_imagebind
+    pos_embed_size: 16
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    start_token: "<|imagebind_start|>"
+    end_token: "<|imagebind_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebind_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224"
+    # Inference pipeline config
+    inference_decode_method: imagebind
+    inference_max_tokens: 17
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0
+
+  - name: imagebindlocal
+    id: 12
+    kind: codebook
+    start_token_key: start_of_imagebindlocal
+    end_token_key: end_of_imagebindlocal
+    pos_embed_size: 1024
+    apply_pos_embed_in_forward: true
+    represent_vae: true
+    conditions: [rgb]
+    codebook_spatial_shape: [32, 32]   # 1024 tokens, 32×32 grid → 2D RoPE
+    start_token: "<|imagebindlocal_start|>"
+    end_token: "<|imagebindlocal_end|>"
+    code_vocab_size: 8192
+    code_token_format: "<|imagebindlocal_{i:04d}|>"
+    external_tokenizer_kind: vqvae
+    external_tokenizer_repo: "EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448"
+    # Inference pipeline config
+    inference_decode_method: imagebindlocal
+    inference_max_tokens: 1025
+    inference_cfg_uncond: img
+    inference_cfg_img_scale: 1.0

core/__init__.py

@@ -0,0 +1,11 @@
+"""
+Core utilities for MODUS.
+
+This package is intentionally small and dependency-light so it can be used by:
+- training scripts
+- inference/demo scripts
+- dataset packing
+- model code
+"""
+
+

core/modality.py

@@ -0,0 +1,394 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field as dc_field, replace as dc_replace
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple
+
+
+TokenRange = Tuple[int, int]  # (base_id, length)
+
+
+@dataclass(frozen=True)
+class LossConfig:
+    """Per-modality loss configuration (lives in the modality YAML)."""
+    reweight: bool = False
+    reweight_min_w: float = 0.005
+    reweight_det_vocab_dim: int = 4000   # only used when reweight=True for coordinate codebooks
+
+
+@dataclass(frozen=True)
+class ModalitySpec:
+    """
+    Config-driven definition of a single modality.
+
+    Every tunable per-modality knob lives here so that the training script,
+    model forward, and dataset packing can all be driven from one YAML.
+    """
+
+    name: str
+    id: int
+    start_token_key: str
+    end_token_key: str
+    kind: str = "text"  # text | image | codebook
+
+    # Token range ``(base, length)`` for CE-loss vocab slicing.
+    # Codebook modalities get their own range from the tokenizer; text
+    # modalities are auto-assigned ``(0, text_vocab_end)`` in
+    # ``ModalityRegistry.from_config``.  Image modalities (ViT/VAE)
+    # never produce CE tokens, so they keep ``None``.
+    code_token_range: Optional[TokenRange] = None
+
+    # Explicit, possibly **non-contiguous** set of CE-loss token ids. Used for
+    # modalities whose alphabet is dispersed across the vocab and therefore
+    # cannot be expressed as a single ``(base, length)`` range (e.g. cocodet,
+    # which reuses det's low-vocab coordinate tokens + new high-vocab class/end
+    # tokens). When set, the model does gather-CE over exactly these ids instead
+    # of the contiguous-slice path. ``None`` for every existing modality.
+    code_token_ids: Optional[Tuple[int, ...]] = None
+
+    # Learnable positional embedding for code tokens.
+    pos_embed_size: Optional[int] = None
+    apply_pos_embed_in_forward: bool = False
+
+    # Image representation flags (used by dataset packing and inferencer).
+    represent_vit: bool = True
+    represent_vae: bool = False
+
+    # Which other modalities may condition this one during training.
+    # ``None`` means "use all available"; an explicit list restricts.
+    conditions: Optional[Tuple[str, ...]] = None
+    # Optional sampling probabilities aligned with ``conditions`` order.
+    # If omitted, the dataset uses a uniform distribution.
+    condition_probs: Optional[Tuple[float, ...]] = None
+
+    # Per-modality loss parameters (CE smoothing / reweighting).
+    loss: LossConfig = LossConfig()
+
+    # Override the model attribute name for the pos-embed (for checkpoint compat).
+    # Defaults to ``name`` (e.g. modality "dino" → attr ``dino_pos_embed``).
+    # Set to e.g. "grounding" for det → ``grounding_pos_embed``.
+    pos_embed_name: Optional[str] = None
+
+    # 2D spatial shape for codebook modalities (h, w).
+    # None → 1D sequential (uses text region of build_2d_rope, gets sequential positions).
+    # Set e.g. (32, 32) for dinolocal, (28, 28) for clip → gets 2D RoPE same as VAE latents.
+    codebook_spatial_shape: Optional[Tuple[int, int]] = None
+
+    # External tokenizer for inference decoding/encoding (e.g. DINO VQVAE).
+    external_tokenizer_repo: Optional[str] = None
+    external_tokenizer_kind: str = "vqvae"
+
+    # ── Inference-time configuration ─────────────────────────────────────
+    # These fields drive the unified inference pipeline so that the
+    # inferencer does not need per-modality if/elif chains.
+
+    # Decode method: "auto" resolves from ``kind`` (image→image, text→text).
+    # Explicit values: "image", "text", "detection", "dino".
+    inference_decode_method: str = "auto"
+
+    # Maximum AR tokens for text / codebook decoding (None → inferencer default).
+    inference_max_tokens: Optional[int] = None
+
+    # Which CFG unconditional context to use when this modality is a *target*:
+    #   "text"  → cfg_text_context (drop text, keep image)   e.g. detection
+    #   "img"   → cfg_img_context  (drop image, keep text)   e.g. dino
+    #   "both"  → dual CFG with both contexts                e.g. image generation
+    #   "none"  → no CFG                                     e.g. plain text
+    inference_cfg_uncond: str = "auto"
+
+    # Whether to prepend "[start {name} x10]" instruction before generation.
+    inference_add_instruction: bool = True
+
+    # Default CFG image scale for this modality when it is the *target*.
+    # ``None`` → use the global default from the base config.
+    inference_cfg_img_scale: Optional[float] = None
+
+
+class ModalityRegistry:
+    """Runtime registry consumed by dataset packing, model forward, and inferencer."""
+
+    def __init__(self, specs: Iterable[ModalitySpec]):
+        specs = list(specs)
+        if len(specs) == 0:
+            raise ValueError("ModalityRegistry requires at least one ModalitySpec.")
+
+        self._by_name: Dict[str, ModalitySpec] = {}
+        self._by_id: Dict[int, ModalitySpec] = {}
+        for s in specs:
+            if s.name in self._by_name:
+                raise ValueError(f"Duplicate modality name: {s.name}")
+            if s.id in self._by_id:
+                raise ValueError(f"Duplicate modality id: {s.id}")
+            self._by_name[s.name] = s
+            self._by_id[s.id] = s
+
+    # ── lookups ──────────────────────────────────────────────────────────
+
+    @property
+    def specs(self) -> List[ModalitySpec]:
+        return list(self._by_name.values())
+
+    def get(self, name: str) -> ModalitySpec:
+        return self._by_name[name]
+
+    def get_by_id(self, modality_id: int) -> ModalitySpec:
+        return self._by_id[modality_id]
+
+    def name_to_id(self) -> Dict[str, int]:
+        return {k: v.id for k, v in self._by_name.items()}
+
+    def id_to_name(self) -> Dict[int, str]:
+        return {k: v.name for k, v in self._by_id.items()}
+
+    def modality_name(self, modality_id: int) -> str:
+        spec = self._by_id.get(modality_id)
+        return spec.name if spec is not None else f"unknown_{modality_id}"
+
+    def start_token_key(self, name: str) -> str:
+        return self._by_name[name].start_token_key
+
+    def end_token_key(self, name: str) -> str:
+        return self._by_name[name].end_token_key
+
+    def start_token_id(self, new_token_ids: Mapping[str, Any], name: str) -> int:
+        return int(new_token_ids[self.start_token_key(name)])
+
+    def end_token_id(self, new_token_ids: Mapping[str, Any], name: str) -> int:
+        return int(new_token_ids[self.end_token_key(name)])
+
+    def code_token_range(self, name: str) -> Optional[TokenRange]:
+        return self._by_name[name].code_token_range
+
+    def has_codebook_modalities(self) -> bool:
+        """True if any registered modality is ``kind == 'codebook'``."""
+        return any(s.kind == "codebook" for s in self._by_name.values())
+
+    def conditions_for(self, target_modality: str) -> Optional[List[str]]:
+        """
+        Return the allowed conditioning modalities for *target_modality*.
+
+        ``None`` means "no restriction — use all available".
+        """
+        spec = self._by_name.get(target_modality)
+        if spec is None or spec.conditions is None:
+            return None
+        return list(spec.conditions)
+
+    def condition_probs_for(self, target_modality: str) -> Optional[List[float]]:
+        """
+        Return optional condition sampling probabilities for *target_modality*.
+
+        The returned probabilities align with ``conditions_for(target_modality)``.
+        ``None`` means "use uniform condition sampling".
+        """
+        spec = self._by_name.get(target_modality)
+        if spec is None or spec.condition_probs is None:
+            return None
+        return list(spec.condition_probs)
+
+    def resolve_decode_method(self, name: str) -> str:
+        """Return the concrete decode method for *name* (resolves ``"auto"``)."""
+        spec = self._by_name[name]
+        if spec.inference_decode_method != "auto":
+            return spec.inference_decode_method
+        if spec.kind == "image":
+            return "image"
+        if spec.kind == "text":
+            return "text"
+        # codebook — guess from name
+        if "det" in spec.name:
+            return "detection"
+        if spec.name == "dinolocal":
+            return "dinolocal"
+        if "dino" in spec.name:
+            return "dino"
+        return "text"
+
+    def resolve_cfg_uncond(self, name: str) -> str:
+        """Return the concrete CFG-uncond context type (resolves ``"auto"``)."""
+        spec = self._by_name[name]
+        if spec.inference_cfg_uncond != "auto":
+            return spec.inference_cfg_uncond
+        dm = self.resolve_decode_method(name)
+        if dm == "image":
+            return "both"
+        if dm == "detection":
+            return "text"
+        if dm in ("dino", "dinolocal"):
+            return "img"
+        return "none"
+
+    def resolve_cfg_img_scale(self, name: str, default: float = 2.0) -> float:
+        """Return the per-modality CFG image scale, or *default* if not specified."""
+        spec = self._by_name[name]
+        if spec.inference_cfg_img_scale is not None:
+            return spec.inference_cfg_img_scale
+        return default
+
+    def needs_external_tokenizer(self, name: str) -> bool:
+        """Return True if the modality has an external tokenizer (e.g. DINO VQVAE)."""
+        spec = self._by_name.get(name)
+        return spec is not None and spec.external_tokenizer_repo is not None
+
+    def modalities_with_forward_pos_embed(self) -> List[ModalitySpec]:
+        return [
+            s
+            for s in self._by_name.values()
+            if s.pos_embed_size is not None and s.apply_pos_embed_in_forward
+        ]
+
+    # ── construction from YAML dict ─────────────────��───────────────────
+
+    @staticmethod
+    def from_config(
+        cfg: Any,
+        *,
+        token_ranges: Optional[Mapping[str, TokenRange]] = None,
+        code_token_ids: Optional[Mapping[str, List[int]]] = None,
+    ) -> "ModalityRegistry":
+        """
+        Build a registry from the parsed YAML config dict.
+
+        Expected structure::
+
+            cfg["modalities"]: list[{name, id, start_token_key, end_token_key, kind, ...}]
+
+        ``token_ranges`` maps modality name → ``(base, length)`` and is typically
+        computed after tokenizer token additions.
+        """
+        if cfg is None:
+            raise ValueError("cfg is required")
+
+        # Accept either ``{"modalities": [...]}`` or a bare list.
+        modalities = cfg.get("modalities") if isinstance(cfg, dict) else cfg
+        if modalities is None:
+            raise ValueError("No 'modalities' key found in cfg")
+
+        specs: List[ModalitySpec] = []
+        token_ranges = dict(token_ranges or {})
+        code_token_ids = dict(code_token_ids or {})
+
+        for m in modalities:
+            name = str(m["name"])
+
+            # Resolve code_token_range: prefer runtime-computed, fall back to cfg.
+            _range = token_ranges.get(name)
+            if _range is None and m.get("code_token_range") is not None:
+                _range = tuple(m["code_token_range"])
+
+            # Resolve explicit (possibly dispersed) code_token_ids, if any.
+            _ids = code_token_ids.get(name)
+            if _ids is None and m.get("code_token_ids") is not None:
+                _ids = m["code_token_ids"]
+            _ids = tuple(int(i) for i in _ids) if _ids is not None else None
+
+            # pos_embed_size may be None.
+            _pos = m.get("pos_embed_size")
+            pos_embed_size = int(_pos) if _pos is not None else None
+
+            # Parse conditions (list of strings or None).
+            _conds = m.get("conditions")
+            conditions = tuple(_conds) if _conds is not None else None
+            _cond_probs = m.get("condition_probs")
+            condition_probs = tuple(float(x) for x in _cond_probs) if _cond_probs is not None else None
+            if condition_probs is not None:
+                if conditions is None:
+                    raise ValueError(
+                        f"Modality '{name}' sets 'condition_probs' without 'conditions'."
+                    )
+                if len(condition_probs) != len(conditions):
+                    raise ValueError(
+                        f"Modality '{name}' has {len(conditions)} conditions but "
+                        f"{len(condition_probs)} condition_probs."
+                    )
+                if any(p < 0.0 for p in condition_probs):
+                    raise ValueError(f"Modality '{name}' has negative values in condition_probs.")
+                probs_sum = float(sum(condition_probs))
+                if abs(probs_sum - 1.0) > 1e-6:
+                    raise ValueError(
+                        f"Modality '{name}' condition_probs must sum to 1.0, got {probs_sum}."
+                    )
+
+            # Parse per-modality loss config.
+            _loss_dict = m.get("loss") or {}
+            loss_cfg = LossConfig(
+                reweight=bool(_loss_dict.get("reweight", False)),
+                reweight_min_w=float(_loss_dict.get("reweight_min_w", 0.005)),
+                reweight_det_vocab_dim=int(_loss_dict.get("reweight_det_vocab_dim", 4000)),
+            )
+
+            _pe_name = m.get("pos_embed_name")
+            pe_name = str(_pe_name) if _pe_name is not None else None
+
+            # 2D spatial shape for codebook modalities.
+            _shape = m.get("codebook_spatial_shape")
+            codebook_spatial_shape = tuple(int(x) for x in _shape) if _shape is not None else None
+
+            # Inference-time fields (optional in YAML; sensible defaults).
+            _infer = m.get("inference") or {}
+            _infer_decode = str(_infer.get("decode_method", m.get("inference_decode_method", "auto")))
+            _infer_max = _infer.get("max_tokens", m.get("inference_max_tokens"))
+            _infer_cfg = str(_infer.get("cfg_uncond", m.get("inference_cfg_uncond", "auto")))
+            _infer_add_instr = bool(_infer.get("add_instruction", m.get("inference_add_instruction", True)))
+            _infer_cfg_img = _infer.get("cfg_img_scale", m.get("inference_cfg_img_scale"))
+            _infer_cfg_img = float(_infer_cfg_img) if _infer_cfg_img is not None else None
+
+            spec = ModalitySpec(
+                name=name,
+                id=int(m["id"]),
+                start_token_key=str(m["start_token_key"]),
+                end_token_key=str(m["end_token_key"]),
+                kind=str(m.get("kind", "text")),
+                code_token_range=_range,
+                code_token_ids=_ids,
+                pos_embed_size=pos_embed_size,
+                apply_pos_embed_in_forward=bool(m.get("apply_pos_embed_in_forward", False)),
+                represent_vit=bool(m.get("represent_vit", True)),
+                represent_vae=bool(m.get("represent_vae", False)),
+                conditions=conditions,
+                condition_probs=condition_probs,
+                loss=loss_cfg,
+                pos_embed_name=pe_name,
+                codebook_spatial_shape=codebook_spatial_shape,
+                external_tokenizer_repo=m.get("external_tokenizer_repo"),
+                external_tokenizer_kind=str(m.get("external_tokenizer_kind", "vqvae")),
+                inference_decode_method=_infer_decode,
+                inference_max_tokens=int(_infer_max) if _infer_max is not None else None,
+                inference_cfg_uncond=_infer_cfg,
+                inference_add_instruction=_infer_add_instr,
+                inference_cfg_img_scale=_infer_cfg_img,
+            )
+            specs.append(spec)
+
+        # ── Auto-assign text token range to text-kind modalities ─────────
+        # Find text_vocab_end = min base across all codebook ranges.
+        text_vocab_end: Optional[int] = None
+        for s in specs:
+            if s.code_token_range is not None:
+                cb_base, _ = s.code_token_range
+                text_vocab_end = int(cb_base) if text_vocab_end is None else min(text_vocab_end, int(cb_base))
+
+        if text_vocab_end is not None:
+            text_range: TokenRange = (0, text_vocab_end)
+            specs = [
+                dc_replace(s, code_token_range=text_range)
+                if s.kind == "text" and s.code_token_range is None else s
+                for s in specs
+            ]
+
+        return ModalityRegistry(specs)
+
+
+def infer_contiguous_token_range(token_ids: List[int]) -> TokenRange:
+    """Given a list of token IDs, infer ``(base, length)`` and validate contiguity."""
+    if len(token_ids) == 0:
+        raise ValueError("token_ids cannot be empty")
+    token_ids_sorted = sorted(int(x) for x in token_ids)
+    base = token_ids_sorted[0]
+    length = len(token_ids_sorted)
+    expected = list(range(base, base + length))
+    if token_ids_sorted != expected:
+        raise ValueError(
+            "Token IDs are not contiguous; cannot represent as a range efficiently. "
+            f"base={base}, length={length}, first10={token_ids_sorted[:10]}"
+        )
+    return base, length

core/model_registry.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional, Tuple
+
+
+ModelBuilder = Callable[..., Any]
+
+
+class ModelRegistry:
+    """
+    Minimal model registry to decouple scripts from concrete model classes.
+
+    Goal:
+    - training/inference scripts choose `model.name` in config
+    - swapping architectures doesn't require editing scripts, only adding a builder registration
+    """
+
+    def __init__(self):
+        self._builders: Dict[str, ModelBuilder] = {}
+
+    def register(self, name: str, builder: ModelBuilder) -> None:
+        if name in self._builders:
+            raise ValueError(f"Model '{name}' is already registered.")
+        self._builders[name] = builder
+
+    def build(self, name: str, **kwargs) -> Any:
+        if name not in self._builders:
+            known = ", ".join(sorted(self._builders.keys()))
+            raise KeyError(f"Unknown model '{name}'. Known: {known}")
+        return self._builders[name](**kwargs)
+
+
+GLOBAL_MODEL_REGISTRY = ModelRegistry()
+
+
+def register_model(name: str) -> Callable[[ModelBuilder], ModelBuilder]:
+    def _decorator(fn: ModelBuilder) -> ModelBuilder:
+        GLOBAL_MODEL_REGISTRY.register(name, fn)
+        return fn
+    return _decorator
+
+
+def build_model(name: str, **kwargs) -> Any:
+    return GLOBAL_MODEL_REGISTRY.build(name, **kwargs)
+
+

core/tokenizer_utils.py

@@ -0,0 +1,257 @@
+"""
+Config-driven tokenizer setup.
+
+All modality tokens (delimiters, codebook entries) are specified in the modality
+YAML and added here.  No hard-coded ``use_det`` / ``use_dino`` flags.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from core.modality import TokenRange, infer_contiguous_token_range
+from data.data_utils import add_special_tokens
+
+
+@dataclass
+class TokenizerArtifacts:
+    """Return value of :func:`build_tokenizer_and_special_tokens`."""
+    tokenizer: Any
+    new_token_ids: Dict[str, Any]
+    token_ranges: Dict[str, TokenRange]
+    # Explicit dispersed CE-token id sets, for modalities that set
+    # ``dispersed_code_tokens: true`` (e.g. cocodet). modality name -> id list.
+    code_token_ids: Dict[str, List[int]]
+
+
+def load_base_tokenizer(*, model_args: Any, training_args: Any):
+    """Load the base tokenizer (before any token additions).
+
+    Dispatches to HunyuanImage3TokenizerFast when model_name == 'hunyuan_image_3',
+    otherwise falls back to the standard Qwen2Tokenizer.
+    """
+    if getattr(model_args, "model_name", "bagel") == "hunyuan_image_3":
+        from modeling.hunyuan_image_3.src.tokenization_hunyuan_image_3 import (
+            HunyuanImage3TokenizerFast,
+        )
+        return HunyuanImage3TokenizerFast.from_pretrained(
+            model_args.model_path,
+            model_version="HunyuanImage-3.0-Instruct",
+        )
+
+    from modeling.qwen2.tokenization_qwen2 import Qwen2Tokenizer
+
+    pretrained_path = (
+        model_args.model_path
+        if getattr(training_args, "finetune_from_hf", False)
+        else model_args.llm_path
+    )
+    return Qwen2Tokenizer.from_pretrained(pretrained_path)
+
+
+# ── internal helpers ─────────────────────────────────────────────────────────
+
+def _add_tokens(tokenizer, tokens: List[str]) -> None:
+    """Add tokens in order, deduplicating while preserving first occurrence."""
+    seen: set = set()
+    ordered: List[str] = []
+    for t in tokens:
+        if t not in seen:
+            seen.add(t)
+            ordered.append(t)
+    tokenizer.add_tokens(ordered)
+
+
+def _ensure_token_key(
+    tokenizer,
+    new_token_ids: Dict[str, Any],
+    *,
+    token_key: str,
+    token_str: Optional[str],
+) -> None:
+    """
+    Make sure ``new_token_ids[token_key]`` exists, creating the token if needed.
+    """
+    if token_key in new_token_ids:
+        return
+    if hasattr(tokenizer, token_key):
+        new_token_ids[token_key] = int(getattr(tokenizer, token_key))
+        return
+    if token_str is None:
+        raise KeyError(
+            f"Token key '{token_key}' not in tokenizer or new_token_ids and no token_str provided."
+        )
+    try:
+        tokenizer.add_special_tokens({"additional_special_tokens": [token_str]})
+    except Exception:
+        tokenizer.add_tokens([token_str])
+    new_token_ids[token_key] = int(tokenizer.convert_tokens_to_ids(token_str))
+
+
+def _code_tokens_from_cfg(m: Dict[str, Any]) -> List[str]:
+    """Build the list of code-token strings from a modality config dict."""
+
+    # Preferred: explicit groups (e.g. DET coordinate tokens).
+    groups = m.get("code_token_groups")
+    if groups is not None:
+        out: List[str] = []
+        for g in groups:
+            if not isinstance(g, dict):
+                g = dict(g)
+            token_format = str(g["token_format"])
+            start = int(g.get("start", 0))
+            end = int(g["end"])
+            for i in range(start, end + 1):
+                out.append(token_format.format(
+                    i=i,
+                    prefix=m.get("code_token_prefix", m.get("name")),
+                ))
+        return out
+
+    # Fallback: contiguous vocab with format / prefix.
+    vocab = m.get("code_vocab_size")
+    if vocab is None:
+        raise ValueError(
+            f"Codebook modality '{m.get('name')}' needs 'code_vocab_size' or 'code_token_groups'."
+        )
+    vocab = int(vocab)
+    token_format = str(m.get("code_token_format", "<|{prefix}_{i:04d}|>"))
+    prefix = str(m.get("code_token_prefix", m.get("name")))
+    return [token_format.format(prefix=prefix, i=i) for i in range(vocab)]
+
+
+def _add_modality_tokens(
+    tokenizer,
+    *,
+    m: Dict[str, Any],
+    new_token_ids: Dict[str, Any],
+    token_ranges: Dict[str, TokenRange],
+    code_token_ids_out: Dict[str, List[int]],
+    deferred_tokens: List[Any],
+) -> None:
+    """
+    Add all tokens for one modality entry (text/image/codebook) from its YAML dict.
+
+    For **codebook** modalities the addition order is *code tokens first, then
+    delimiter / extra tokens*.  This matches the ordering used when the
+    original checkpoint was trained so that token IDs remain identical.
+    """
+    kind = str(m.get("kind", "text"))
+    start_key = str(m["start_token_key"])
+    end_key = str(m["end_token_key"])
+
+    if kind != "codebook":
+        # Image/text modalities: resolve already-existing start/end tokens now
+        # (shared start_of_image, or start_of_depth/normal added by
+        # add_special_tokens). Any genuinely NEW token (e.g. a per-modality
+        # start_of_seg/canny/samseg/samedge under REPLACE) is DEFERRED and added
+        # at the very tail in build_tokenizer_and_special_tokens — so introducing
+        # an image modality with its own start token does NOT shift the ids of
+        # codebook tokens added later in the loop, preserving cross-stage ckpt
+        # alignment.
+        vocab = tokenizer.get_vocab()
+        for token_key, token_str in ((start_key, m.get("start_token")), (end_key, m.get("end_token"))):
+            if token_key in new_token_ids:
+                continue
+            if token_str is not None and str(token_str) in vocab:
+                new_token_ids[token_key] = int(vocab[str(token_str)])
+            elif token_str is not None:
+                deferred_tokens.append((token_key, str(token_str)))
+            else:
+                _ensure_token_key(tokenizer, new_token_ids, token_key=token_key, token_str=None)
+        return
+
+    # ── codebook modality ────────────────────────────────────────────────
+    code_tokens = _code_tokens_from_cfg(m)
+
+    start_tok = m.get("start_token")
+    end_tok = m.get("end_token")
+    delim_tokens: List[str] = []
+    if start_tok is not None:
+        delim_tokens.append(str(start_tok))
+    if end_tok is not None:
+        delim_tokens.append(str(end_tok))
+    delim_tokens.extend(str(t) for t in (m.get("extra_tokens") or []))
+
+    # One batch: code tokens first, delimiters after → matches checkpoint ordering.
+    _add_tokens(tokenizer, code_tokens + delim_tokens)
+
+    # Populate delimiter keys.
+    if start_tok is not None and start_key not in new_token_ids:
+        new_token_ids[start_key] = int(tokenizer.convert_tokens_to_ids(str(start_tok)))
+    elif start_key not in new_token_ids:
+        _ensure_token_key(tokenizer, new_token_ids, token_key=start_key, token_str=None)
+    if end_tok is not None and end_key not in new_token_ids:
+        new_token_ids[end_key] = int(tokenizer.convert_tokens_to_ids(str(end_tok)))
+    elif end_key not in new_token_ids:
+        _ensure_token_key(tokenizer, new_token_ids, token_key=end_key, token_str=None)
+
+    name = str(m["name"])
+    token_ids = [int(tokenizer.convert_tokens_to_ids(t)) for t in code_tokens]
+    if bool(m.get("dispersed_code_tokens", False)):
+        # Dispersed alphabet (e.g. cocodet: reused det coords at low vocab ids +
+        # new class tokens at the tail) — cannot be a single contiguous range.
+        # Record the explicit CE-token id set = all code tokens + the end
+        # delimiter (a CE target / stop token). The start delimiter is
+        # input-only, so it is excluded.
+        ids = list(token_ids)
+        end_id = new_token_ids.get(end_key)
+        if end_id is not None:
+            ids.append(int(end_id))
+        code_token_ids_out[name] = ids
+    else:
+        # Record contiguous code-token range.
+        token_ranges[name] = infer_contiguous_token_range(token_ids)
+
+
+# ── public API ───────────────────────────────────────────────────────────────
+
+def build_tokenizer_and_special_tokens(
+    tokenizer,
+    *,
+    modalities_cfg: Dict[str, Any],
+) -> TokenizerArtifacts:
+    """
+    Fully config-driven tokenizer setup.
+
+    1. Adds universal special tokens (``<|im_start|>``, ``<|vision_start|>``, …)
+       via the legacy ``add_special_tokens`` helper (for checkpoint compat).
+    2. Iterates over ``modalities_cfg["modalities"]`` and adds delimiter / code
+       tokens for each entry.
+
+    Returns a :class:`TokenizerArtifacts` with the final tokenizer, a dict of
+    all special-token IDs, and a dict of code-token ranges for codebook modalities.
+    """
+    token_ranges: Dict[str, TokenRange] = {}
+    code_token_ids: Dict[str, List[int]] = {}
+    deferred_tokens: List[Any] = []
+
+    # Step 1: universal special tokens (checkpoint-compatible ordering).
+    tokenizer, new_token_ids, _num = add_special_tokens(tokenizer)
+
+    # Step 2: per-modality tokens from YAML.
+    modalities = modalities_cfg.get("modalities", modalities_cfg)
+    for m in modalities:
+        if not isinstance(m, dict):
+            m = dict(m)
+        _add_modality_tokens(
+            tokenizer,
+            m=m,
+            new_token_ids=new_token_ids,
+            token_ranges=token_ranges,
+            code_token_ids_out=code_token_ids,
+            deferred_tokens=deferred_tokens,
+        )
+
+    # Step 3: append deferred NEW image-modality start/end tokens at the tail
+    # (after all codebook tokens) so they never shift existing ids across stages.
+    for token_key, token_str in deferred_tokens:
+        if token_key not in new_token_ids:
+            _ensure_token_key(tokenizer, new_token_ids, token_key=token_key, token_str=token_str)
+
+    return TokenizerArtifacts(
+        tokenizer=tokenizer,
+        new_token_ids=new_token_ids,
+        token_ranges=token_ranges,
+        code_token_ids=code_token_ids,
+    )

data/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0

data/any2any_preprocess/_build_preview_and_montage.py

@@ -0,0 +1,52 @@
+import io, os
+import pyarrow.parquet as pq
+from PIL import Image, ImageDraw
+
+BASE_SRC = ('datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg_grounding'
+            '_canny_dino_global_clip448_imagebind_samseg_samedge_cocodet')
+Q95 = 'datasets/blip3o/parquet_16mod_normalq95_sample/sa_000000.parquet'
+PREVIEW_DIR = 'datasets/blip3o/modus_preview'
+MONTAGE = 'eval_outputs/modus_hero_montage.png'
+N_ROWS = 500
+
+# ---- 1. print sample uids per source to confirm the alignment key ----
+print('==== uid samples (alignment key check) ====', flush=True)
+for f in ['sa_000000.parquet', 'webdataset_shard_000.parquet', 'webdataset_JDB_2.parquet']:
+    p = f'{BASE_SRC}/{f}'
+    t = pq.ParquetFile(p).read_row_group(0, columns=['uid'])
+    uids = t.column('uid').to_pylist()[:4]
+    print(f'  {f}: {uids}', flush=True)
+
+# ---- 2. build small preview parquet (subset of q95-compressed sa_000000) ----
+os.makedirs(PREVIEW_DIR, exist_ok=True)
+pf = pq.ParquetFile(Q95)
+table = pf.read_row_group(0).slice(0, N_ROWS)
+dst = f'{PREVIEW_DIR}/modus_preview_sa_500.parquet'
+pq.write_table(table, dst, compression='snappy')
+print(f'\npreview: {dst}  rows={table.num_rows}  size={os.path.getsize(dst)/1e6:.1f} MB', flush=True)
+
+# ---- 3. hero montage: rows=samples, cols=modalities ----
+MODS = ['rgb', 'depth', 'normal', 'canny', 'sam_seg', 'sam_edge']
+N_SHOW = 5
+cell, pad, hdr = 300, 6, 26
+cols = {m: table.column(m).to_pylist()[:N_SHOW] for m in MODS}
+
+def dec(b):
+    if b is None: return None
+    if isinstance(b, memoryview): b = b.tobytes()
+    return Image.open(io.BytesIO(b)).convert('RGB')
+
+W = len(MODS)*(cell+pad)+pad
+H = hdr + N_SHOW*(cell+pad)+pad
+grid = Image.new('RGB', (W, H), 'white')
+d = ImageDraw.Draw(grid)
+for c, m in enumerate(MODS):
+    d.text((pad+c*(cell+pad)+4, 6), m, fill='black')
+for r in range(N_SHOW):
+    for c, m in enumerate(MODS):
+        im = dec(cols[m][r])
+        if im is None: continue
+        grid.paste(im.resize((cell, cell)), (pad+c*(cell+pad), hdr+pad+r*(cell+pad)))
+os.makedirs(os.path.dirname(MONTAGE), exist_ok=True)
+grid.save(MONTAGE)
+print(f'montage: {MONTAGE}  ({W}x{H})', flush=True)

data/any2any_preprocess/_cast_preview_images.py

@@ -0,0 +1,34 @@
+"""Cast the preview parquet's raw-bytes image columns to the HF `Image` feature
+so the dataset viewer renders thumbnails."""
+import json
+import os
+from datasets import load_dataset, Image, Value
+
+SRC = 'datasets/blip3o/modus_preview/modus_preview_sa_500.parquet'
+OUT_DIR = 'datasets/blip3o/modus_preview_upload/data'
+OUT = f'{OUT_DIR}/modus_preview_sa_500.parquet'
+IMG_COLS = ['rgb', 'depth', 'normal', 'canny', 'sam_seg', 'sam_edge']
+
+# Display-name order requested by the user (uid=key first; det_seg + grounding
+# raw blobs moved to the very end).
+ORDER = [
+    'uid', 'rgb', 'depth', 'normal', 'sam_seg', 'canny', 'sam_edge',
+    'coco_det', 'caption', 'dino_global', 'dino', 'clip448',
+    'imagebind_global', 'imagebind', 'det_seg', 'grounding',
+]
+
+ds = load_dataset('parquet', data_files=SRC, split='train')
+print('loaded:', ds.num_rows, 'rows; features before:', list(ds.features.keys()))
+for c in IMG_COLS:
+    ds = ds.cast_column(c, Image())
+# Flatten deep-nested blobs to JSON strings so the HF viewer worker doesn't OOM.
+for c in ['det_seg', 'grounding']:
+    ds = ds.map(lambda ex, col=c: {col: json.dumps(ex[col], default=str, ensure_ascii=False)})
+    ds = ds.cast_column(c, Value('string'))
+assert set(ORDER) == set(ds.column_names), set(ORDER) ^ set(ds.column_names)
+ds = ds.select_columns(ORDER)
+print('column order:', ds.column_names)
+
+os.makedirs(OUT_DIR, exist_ok=True)
+ds.to_parquet(OUT)
+print(f'wrote {OUT}  size={os.path.getsize(OUT)/1e6:.1f} MB', flush=True)

data/any2any_preprocess/_sb_full_rebuild.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+#SBATCH --job-name=modus_full_rebuild
+#SBATCH --partition=debug
+#SBATCH --time=01:30:00
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=700GB
+#SBATCH --account=a143
+#SBATCH --output=logs/full_rebuild_%j.out
+#SBATCH --error=logs/full_rebuild_%j.err
+#SBATCH --environment=/iopsstor/scratch/cscs/mye/workspace/BAGEL/environment/bagel.toml
+
+M=/iopsstor/scratch/cscs/mye/workspace/MODUS_RELEASE/MODUS
+srun --environment=/iopsstor/scratch/cscs/mye/workspace/BAGEL/environment/bagel.toml \
+  bash -c "cd $M && export PYTHONPATH='$M/data/any2any_preprocess:\$PYTHONPATH' && \
+    python3 data/any2any_preprocess/run_full_rebuild.py \
+      --num-tasks 1 --task-id 0 --concurrency 64"

data/any2any_preprocess/_sb_recompress_normal_sample.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+#SBATCH --job-name=norm_q95_smp
+#SBATCH --partition=debug
+#SBATCH --time=01:30:00
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=64GB
+#SBATCH --account=a143
+#SBATCH --output=logs/norm_q95_sample_%j.out
+#SBATCH --error=logs/norm_q95_sample_%j.err
+#SBATCH --environment=/iopsstor/scratch/cscs/mye/workspace/BAGEL/environment/bagel.toml
+
+M=/iopsstor/scratch/cscs/mye/workspace/MODUS_RELEASE/MODUS
+SRC="$M/datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global_clip448_imagebind_samseg_samedge_cocodet"
+OUT="$M/datasets/blip3o/parquet_16mod_normalq95_sample"
+
+srun --environment=/iopsstor/scratch/cscs/mye/workspace/BAGEL/environment/bagel.toml \
+  bash -c "cd $M && python3 data/any2any_preprocess/recompress_normal_jpeg.py \
+    --out_dir '$OUT' \
+    '$SRC/sa_000000.parquet' \
+    '$SRC/sa_000001.parquet' \
+    '$SRC/sa_000002.parquet' \
+    '$SRC/sa_000003.parquet'"

data/any2any_preprocess/_sb_upload.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+#SBATCH --job-name=modus_upload
+#SBATCH --partition=debug
+#SBATCH --time=01:30:00
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=32
+#SBATCH --mem=200GB
+#SBATCH --account=a143
+#SBATCH --output=logs/upload_%j.out
+#SBATCH --error=logs/upload_%j.err
+#SBATCH --environment=/iopsstor/scratch/cscs/mye/workspace/BAGEL/environment/bagel.toml
+
+# Resumable upload of the whole modus_full folder (all configs) + README.
+M=/iopsstor/scratch/cscs/mye/workspace/MODUS_RELEASE/MODUS
+srun --environment=/iopsstor/scratch/cscs/mye/workspace/BAGEL/environment/bagel.toml \
+  bash -c "cd $M && export HF_TOKEN_FILE=/users/mye/.hf_token && \
+    python3 -c \"from huggingface_hub import HfApi; t=open('/users/mye/.hf_token').read().strip(); \
+      HfApi(token=t).upload_file(path_or_fileobj='datasets/blip3o/modus_preview_upload/README.md', path_in_repo='README.md', repo_id='epfl-vilab-modus/MODUS-16Modality', repo_type='dataset', commit_message='full configs card'); print('README ok')\" ; \
+    python3 data/any2any_preprocess/upload_full.py \
+      epfl-vilab-modus/MODUS-16Modality datasets/blip3o/modus_full '*/*.parquet'"

data/any2any_preprocess/_verify_normal_q95.py

@@ -0,0 +1,34 @@
+import io, os
+import pyarrow.parquet as pq
+from PIL import Image
+
+SRC = ('datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg_grounding'
+       '_canny_dino_global_clip448_imagebind_samseg_samedge_cocodet/sa_000000.parquet')
+NEW = 'datasets/blip3o/parquet_16mod_normalq95_sample/sa_000000.parquet'
+OUT = 'eval_outputs/normal_q95_check'
+os.makedirs(OUT, exist_ok=True)
+
+def load(path, col, n):
+    t = pq.ParquetFile(path).read_row_group(0, columns=[col])
+    vals = t.column(col).to_pylist()[:n]
+    out = []
+    for b in vals:
+        if isinstance(b, memoryview): b = b.tobytes()
+        out.append(Image.open(io.BytesIO(b)).convert('RGB'))
+    return out
+
+N = 3
+rgb = load(SRC, 'rgb', N)
+nsrc = load(SRC, 'normal', N)
+nnew = load(NEW, 'normal', N)
+
+cell = 384
+grid = Image.new('RGB', (cell*3, cell*N), 'white')
+for r in range(N):
+    for c, im in enumerate([rgb[r], nsrc[r], nnew[r]]):
+        grid.paste(im.resize((cell, cell)), (c*cell, r*cell))
+path = f'{OUT}/rgb_normalPNG_normalJPEGq95.png'
+grid.save(path)
+print(f'saved {path}  (cols: rgb | normal-PNG(src) | normal-JPEGq95(new))', flush=True)
+for r in range(N):
+    print(f'  row{r}: normal src size={nsrc[r].size}  new size={nnew[r].size}', flush=True)

data/any2any_preprocess/build_full_release.py

@@ -0,0 +1,108 @@
+"""Build the MODUS full-release parquet for one source config.
+
+Per config:
+  sa1b / journeydb : annotations-only -> DROP ['rgb','caption'] ; recompress normal->JPEG q95
+  cc12m            : self-contained    -> keep all columns       ; recompress normal->JPEG q95
+
+Streams by row group (bounded memory). Resumable: skips outputs that already
+exist and are complete; writes to a .tmp then atomically renames.
+
+Usage:
+  python build_full_release.py --config sa1b --out_dir <dir> --workers 8 file1.parquet [...]
+"""
+import argparse
+import io
+import os
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from PIL import Image
+
+DROP = {'sa1b': ['rgb', 'caption'], 'journeydb': ['rgb', 'caption'], 'cc12m': []}
+JPEG_QUALITY = 95
+
+
+def _recompress_normal(b):
+    if b is None:
+        return None
+    if isinstance(b, memoryview):
+        b = b.tobytes()
+    with Image.open(io.BytesIO(b)) as im:
+        img = im.convert('RGB')
+    out = io.BytesIO()
+    img.save(out, format='JPEG', quality=JPEG_QUALITY)
+    img.close()
+    return out.getvalue()
+
+
+def process_one(src, dst, drop_cols):
+    if os.path.exists(dst) and os.path.getsize(dst) > 0:
+        return src, 'skip', os.path.getsize(dst)
+    tmp = dst + '.tmp'
+    pf = pq.ParquetFile(src)
+    names = pf.schema_arrow.names
+    keep = [n for n in names if n not in drop_cols]
+    out_schema = pa.schema([pf.schema_arrow.field(n) for n in keep])
+    has_normal = 'normal' in keep
+    norm_field = out_schema.field('normal') if has_normal else None
+    ni = keep.index('normal') if has_normal else -1
+
+    # Stream in small batches so peak memory is bounded regardless of row-group size.
+    writer = pq.ParquetWriter(tmp, out_schema, compression='snappy')
+    try:
+        for batch in pf.iter_batches(batch_size=128, columns=keep):
+            t = pa.Table.from_batches([batch], schema=out_schema)
+            if has_normal:
+                col = [_recompress_normal(b) for b in t.column('normal').to_pylist()]
+                t = t.set_column(ni, norm_field, pa.array(col, type=norm_field.type))
+            writer.write_table(t)
+            del t
+    finally:
+        writer.close()
+    os.replace(tmp, dst)
+    return src, 'done', os.path.getsize(dst)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--config', required=True, choices=list(DROP))
+    ap.add_argument('--out_dir', required=True)
+    ap.add_argument('--workers', type=int, default=8)
+    ap.add_argument('files', nargs='+')
+    args = ap.parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+    drop_cols = DROP[args.config]
+
+    tasks = [(f, os.path.join(args.out_dir, os.path.basename(f)), drop_cols)
+             for f in args.files]
+    done = skip = 0
+
+    if args.workers == 1:
+        # Sequential, in-process (used when spawned as a per-chunk subprocess so
+        # the OS reclaims all memory on exit — no leak accumulation).
+        for src, dst, dc in tasks:
+            _, status, _ = process_one(src, dst, dc)
+            done += status == 'done'
+            skip += status == 'skip'
+            print(f'{status} {os.path.basename(src)}', flush=True)
+        print(f'CHUNK {args.config}: {done} built, {skip} skipped', flush=True)
+        return
+
+    tot_out = 0
+    with ProcessPoolExecutor(max_workers=args.workers) as ex:
+        futs = [ex.submit(process_one, *t) for t in tasks]
+        for i, fu in enumerate(as_completed(futs), 1):
+            src, status, sz = fu.result()
+            tot_out += sz
+            done += status == 'done'
+            skip += status == 'skip'
+            if i % 10 == 0 or status == 'done':
+                print(f'[{i}/{len(tasks)}] {status} {os.path.basename(src)} '
+                      f'done={done} skip={skip} out={tot_out/1e9:.1f}GB', flush=True)
+    print(f'CONFIG {args.config}: {done} built, {skip} skipped, '
+          f'total_out={tot_out/1e9:.1f}GB', flush=True)
+
+
+if __name__ == '__main__':
+    main()

data/any2any_preprocess/check_vqa.py

@@ -0,0 +1,188 @@
+import os
+import json
+import argparse
+from typing import Dict, Any, Tuple, List
+
+
+def is_valid_name(name: Any) -> bool:
+    return isinstance(name, str) and len(name.strip()) > 0
+
+
+def process_labels_jsonl(images_dir: str, labels_path: str, max_lines: int = -1) -> Tuple[Dict[str, Any], List[str]]:
+    stats: Dict[str, Any] = dict(
+        num_records=0,
+        num_invalid_json=0,
+        num_image_refs=0,
+        num_images_missing=0,
+        num_unique_images_missing=0,
+        num_video_refs=0,
+        num_videos_missing=0,
+    )
+    missing_images_unique = set()
+
+    with open(labels_path, "r") as f:
+        for line_idx, line in enumerate(f):
+            if max_lines > 0 and line_idx >= max_lines:
+                break
+            try:
+                record = json.loads(line)
+            except Exception:
+                stats["num_invalid_json"] += 1
+                continue
+
+            stats["num_records"] += 1
+
+            # Check images (string or list)
+            if "image" in record:
+                if isinstance(record["image"], list):
+                    image_names = record["image"]
+                else:
+                    image_names = [record["image"]]
+
+                for name in image_names:
+                    stats["num_image_refs"] += 1
+                    if not is_valid_name(name):
+                        stats["num_images_missing"] += 1
+                        missing_images_unique.add(str(name))
+                        continue
+                    image_path = os.path.join(images_dir, name)
+                    if not os.path.isfile(image_path):
+                        stats["num_images_missing"] += 1
+                        missing_images_unique.add(name)
+
+            # Check video if present (optional)
+            if "video" in record:
+                stats["num_video_refs"] += 1
+                video_name = record["video"]
+                if not is_valid_name(video_name):
+                    stats["num_videos_missing"] += 1
+                else:
+                    video_path = os.path.join(images_dir, video_name)
+                    if not os.path.isfile(video_path):
+                        stats["num_videos_missing"] += 1
+
+    stats["num_unique_images_missing"] = len(missing_images_unique)
+    # Return up to some examples of missing files for debugging
+    sample_missing = list(missing_images_unique)[:50]
+    return stats, sample_missing
+
+
+def scan_base_dir(base_dir: str, max_lines: int = -1) -> Dict[str, Any]:
+    summary: Dict[str, Any] = dict(per_dataset={}, totals={})
+    totals = dict(
+        datasets_scanned=0,
+        num_records=0,
+        num_invalid_json=0,
+        num_image_refs=0,
+        num_images_missing=0,
+        num_unique_images_missing=0,  # accumulated across datasets (sum, not set-union)
+        num_video_refs=0,
+        num_videos_missing=0,
+    )
+
+    if not os.path.isdir(base_dir):
+        raise FileNotFoundError(f"Base dir does not exist: {base_dir}")
+
+    # Pre-collect datasets that have labels.jsonl for better progress logging
+    ds_entries = []
+    for entry in sorted(os.listdir(base_dir)):
+        ds_dir = os.path.join(base_dir, entry)
+        if not os.path.isdir(ds_dir):
+            continue
+        labels_path = os.path.join(ds_dir, "label", "labels.jsonl")
+        if os.path.isfile(labels_path):
+            ds_entries.append(entry)
+
+    total_datasets = len(ds_entries)
+    for idx, entry in enumerate(ds_entries, 1):
+        ds_dir = os.path.join(base_dir, entry)
+        images_dir = os.path.join(ds_dir, "images")
+        labels_path = os.path.join(ds_dir, "label", "labels.jsonl")
+
+        print(f"[{idx}/{total_datasets}] Scanning {entry} ...", flush=True)
+
+        totals["datasets_scanned"] += 1
+        ds_stats, sample_missing = process_labels_jsonl(images_dir, labels_path, max_lines=max_lines)
+
+        summary["per_dataset"][entry] = dict(stats=ds_stats, sample_missing_images=sample_missing)
+
+        # Accumulate totals (note: unique missing is summed per dataset)
+        for k in [
+            "num_records",
+            "num_invalid_json",
+            "num_image_refs",
+            "num_images_missing",
+            "num_unique_images_missing",
+            "num_video_refs",
+            "num_videos_missing",
+        ]:
+            totals[k] += ds_stats[k]
+
+        # Immediate per-dataset summary line
+        pct = 100.0 * ds_stats["num_images_missing"] / max(1, ds_stats["num_image_refs"])
+        print(
+            f"    -> records={ds_stats['num_records']}, image_refs={ds_stats['num_image_refs']}, "
+            f"missing_images={ds_stats['num_images_missing']} ({pct:.3f}%), videos_missing={ds_stats['num_videos_missing']}",
+            flush=True,
+        )
+
+    summary["totals"] = totals
+    return summary
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Check VQA JSONL image/video references against images directories.")
+    parser.add_argument(
+        "--base_dir",
+        type=str,
+        default="./datasets/llava_onevision_vqa",
+        help="Root directory containing per-dataset subfolders with images/ and label/labels.jsonl",
+    )
+    parser.add_argument(
+        "--max_lines",
+        type=int,
+        default=-1,
+        help="If >0, limit number of lines per JSONL file for a faster sample-based check.",
+    )
+    parser.add_argument(
+        "--out_json",
+        type=str,
+        default="",
+        help="Optional path to write a JSON summary report.",
+    )
+
+    args = parser.parse_args()
+
+    summary = scan_base_dir(args.base_dir, max_lines=args.max_lines)
+
+    totals = summary["totals"]
+    print("==== VQA Consistency Check ====")
+    print(
+        f"Datasets: {totals['datasets_scanned']} | Records: {totals['num_records']} | "
+        f"Image refs: {totals['num_image_refs']} | Missing images: {totals['num_images_missing']} | "
+        f"Video refs: {totals['num_video_refs']} | Missing videos: {totals['num_videos_missing']}"
+    )
+    if totals["num_image_refs"] > 0:
+        miss_pct = 100.0 * totals["num_images_missing"] / max(1, totals["num_image_refs"])
+        print(f"Image missing rate: {miss_pct:.4f}%")
+
+    # Print a few per-dataset highlights
+    for ds_name, data in list(summary["per_dataset"].items())[:20]:
+        s = data["stats"]
+        pct = 100.0 * s["num_images_missing"] / max(1, s["num_image_refs"])
+        print(
+            f"- {ds_name}: records={s['num_records']}, image_refs={s['num_image_refs']}, "
+            f"missing_images={s['num_images_missing']} ({pct:.3f}%)"
+        )
+
+    if args.out_json:
+        with open(args.out_json, "w") as f:
+            json.dump(summary, f)
+        print(f"Summary written to {args.out_json}")
+
+
+if __name__ == "__main__":
+    main()
+
+
+# python data/any2any_preprocess/check_vqa.py --base_dir ./datasets/llava_onevision_vqa

data/any2any_preprocess/generate_parquest_grounding_canny_dino_global.py

@@ -0,0 +1,205 @@
+import os
+import io
+import glob
+import tarfile
+import logging
+from typing import Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+class GlobalDinoReplacer:
+    def __init__(self, base_dir: str,
+                 input_parquet_dir: Optional[str] = None,
+                 global_dino_dir: Optional[str] = None,
+                 output_parquet_dir: Optional[str] = None) -> None:
+        self.base_dir = base_dir
+
+        # Defaults aligned with generate_parquet_grounding_canny_dino.py
+        self.input_parquet_dir = input_parquet_dir or os.path.join(
+            base_dir, 'parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino'
+        )
+        self.global_dino_dir = global_dino_dir or os.path.join(base_dir, 'dinov2_global')
+        self.output_parquet_dir = output_parquet_dir or os.path.join(
+            base_dir, 'parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global'
+        )
+
+        os.makedirs(self.output_parquet_dir, exist_ok=True)
+
+    def _list_input_parquets(self) -> List[str]:
+        files = glob.glob(os.path.join(self.input_parquet_dir, '*.parquet'))
+        return sorted(files)
+
+    def _load_global_dino_from_tar(self, tar_path: str) -> Dict[str, np.ndarray]:
+        """Load all uid -> flattened tokens from a dinov2_global tar."""
+        results: Dict[str, np.ndarray] = {}
+
+        if not os.path.exists(tar_path):
+            logger.warning(f"Global DINO tar not found: {tar_path}")
+            return results
+
+        try:
+            with tarfile.open(tar_path, 'r') as tar:
+                for member in tar.getmembers():
+                    if not member.isfile():
+                        continue
+                    name = member.name
+                    if not name.endswith('_tokens_global.npy'):
+                        continue
+
+                    file_obj = tar.extractfile(member)
+                    if file_obj is None:
+                        continue
+
+                    content = file_obj.read()
+                    npy = np.load(io.BytesIO(content))
+
+                    # Derive uid by stripping extension and the trailing `_tokens`
+                    uid = os.path.splitext(name)[0]
+                    if uid.endswith('_tokens_global'):
+                        uid = uid[:-len('_tokens_global')]
+
+                    results[uid] = npy.flatten()
+        except Exception as e:
+            logger.error(f"Failed to read {tar_path}: {e}")
+
+        return results
+
+    def _replace_dino_column(self, df: pd.DataFrame, uid_to_tokens: Dict[str, np.ndarray],
+                              strict: bool = False) -> pd.DataFrame:
+        if 'uid' not in df.columns:
+            raise ValueError("Input parquet missing required 'uid' column")
+
+        num_rows = len(df)
+        replaced = 0
+        missing = 0
+
+        # Prepare a new column without mutating during iteration to avoid pandas warnings
+        new_dino_column: List[object] = []
+        for _, row in df.iterrows():
+            uid = row['uid']
+            tokens = uid_to_tokens.get(uid)
+            if tokens is None:
+                missing += 1
+                if strict:
+                    new_dino_column.append(None)
+                else:
+                    # Keep original if present
+                    new_dino_column.append(row.get('dino_global', None))
+            else:
+                new_dino_column.append(tokens)
+                replaced += 1
+
+        logger.info(f"Replaced DINO tokens for {replaced}/{num_rows} rows; missing {missing}")
+
+        if strict and missing > 0:
+            logger.error("Strict mode: some UIDs are missing global DINO tokens")
+
+        df = df.copy()
+        df['dino_global'] = new_dino_column
+        return df
+
+    def process_one_parquet(self, parquet_path: str, overwrite: bool = False, strict: bool = False) -> Optional[str]:
+        base_name = os.path.splitext(os.path.basename(parquet_path))[0]
+        out_path = os.path.join(self.output_parquet_dir, f"{base_name}.parquet")
+
+        if os.path.exists(out_path) and not overwrite:
+            logger.info(f"Output exists, skipping: {out_path}")
+            return out_path
+
+        # Load the input parquet
+        df = pd.read_parquet(parquet_path, engine='pyarrow')
+
+        # Load corresponding global dino tar
+        global_dino_tar = os.path.join(self.global_dino_dir, f"{base_name}.tar")
+        logger.info(f"Loading global DINO tokens from: {global_dino_tar}")
+        uid_to_tokens = self._load_global_dino_from_tar(global_dino_tar)
+
+
+        if not uid_to_tokens:
+            logger.warning(f"No tokens loaded from {global_dino_tar}; leaving original 'dino' values")
+
+        # Replace column
+        df = self._replace_dino_column(df, uid_to_tokens, strict=strict)
+
+        # Persist
+        df.to_parquet(
+            out_path,
+            index=False,
+            engine='pyarrow',
+            compression='snappy',
+            row_group_size=1000,
+        )
+        logger.info(f"Wrote updated parquet with global DINO to: {out_path}")
+        return out_path
+
+    def process_all(self, overwrite: bool = False, strict: bool = False,
+                    start_from: int = 0) -> List[str]:
+        inputs = self._list_input_parquets()
+        if start_from > 0:
+            inputs = inputs[start_from:]
+
+        if not inputs:
+            logger.error("No input parquet files found")
+            return []
+
+        logger.info(f"Found {len(inputs)} parquet files to update")
+        outputs: List[str] = []
+        for parquet_path in inputs:
+            try:
+                result = self.process_one_parquet(parquet_path, overwrite=overwrite, strict=strict)
+                if result:
+                    outputs.append(result)
+            except Exception as e:
+                logger.error(f"Error processing {parquet_path}: {e}")
+
+        logger.info(f"Successfully wrote {len(outputs)} updated parquet files")
+        return outputs
+
+
+def main() -> None:
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Replace DINO tokens in existing parquet files with dinov2_global tokens')
+    parser.add_argument('--base_dir', type=str, default='./datasets/blip3o',
+                        help='Base directory (default: ./datasets/blip3o)')
+    parser.add_argument('--input_parquet_dir', type=str, default=None,
+                        help='Input parquet directory; defaults to parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino')
+    parser.add_argument('--global_dino_dir', type=str, default=None,
+                        help='Directory containing dinov2_global tars; defaults to base_dir/dinov2_global')
+    parser.add_argument('--output_parquet_dir', type=str, default=None,
+                        help='Output parquet directory; defaults to parquet_*_dino_global')
+    parser.add_argument('--overwrite', action='store_true', help='Overwrite existing output parquet if present')
+    parser.add_argument('--strict', action='store_true',
+                        help='If set, missing UIDs in global DINO will set dino=None and log error')
+    parser.add_argument('--start_from', type=int, default=0,
+                        help='Index to start processing from (for resuming)')
+
+    args = parser.parse_args()
+
+    replacer = GlobalDinoReplacer(
+        base_dir=args.base_dir,
+        input_parquet_dir=args.input_parquet_dir,
+        global_dino_dir=args.global_dino_dir,
+        output_parquet_dir=args.output_parquet_dir,
+    )
+
+    outputs = replacer.process_all(overwrite=args.overwrite, strict=args.strict, start_from=args.start_from)
+    if outputs:
+        logger.info("Outputs written:")
+        for p in outputs:
+            logger.info(f"  - {p}")
+    else:
+        logger.error("No outputs were created")
+
+
+if __name__ == '__main__':
+    main()
+
+

data/any2any_preprocess/generate_parquet.py

@@ -0,0 +1,229 @@
+import os
+import tarfile
+import pandas as pd
+import numpy as np
+from PIL import Image
+import io
+import glob
+import json
+from tqdm import tqdm
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class TarProcessor:
+    def __init__(self, base_dir):
+        self.base_dir = base_dir
+        self.rgb_caption_dir = os.path.join(base_dir, 'datasets')
+        self.depth_dir = os.path.join(base_dir, 'datasets_depth')
+        self.normal_dir = os.path.join(base_dir, 'datasets_normal')
+        self.det_seg_dir = os.path.join(base_dir, 'datasets_seg_swinb')
+        self.parquet_dir = os.path.join(base_dir, 'parquet_rgb_caption_depth_normal_det_seg')
+        
+        # Create parquet directory if it doesn't exist
+        os.makedirs(self.parquet_dir, exist_ok=True)
+        
+    def get_tar_files(self):
+        """Get all tar files from rgb_caption directory"""
+        tar_files = glob.glob(os.path.join(self.rgb_caption_dir, '*.tar'))
+        return sorted(tar_files)
+    
+    def extract_from_tar(self, tar_path, file_type):
+        """
+        Extract images and captions from tar file
+        
+        Args:
+            tar_path: Path to tar file
+            file_type: 'rgb_caption', 'depth', 'normal', 'det_seg'
+        """
+        data = {}
+        
+        try:
+            with tarfile.open(tar_path, 'r') as tar:
+                for member in tar.getmembers():
+                    if member.isfile():
+                        filename = member.name
+                        uid = os.path.splitext(filename)[0]  # Remove extension
+                        
+                        # Extract file content
+                        f = tar.extractfile(member)
+                        if f is None:
+                            continue
+                        
+                        content = f.read()
+                        
+                        if file_type == 'rgb_caption':
+                            # Handle RGB + caption files
+                            if filename.endswith('.jpg') or filename.endswith('.jpeg') or filename.endswith('.png'):
+                                data[uid] = {'image': content, 'caption': None}
+                            elif filename.endswith('.txt'):
+                                # Find corresponding image
+                                base_name = os.path.splitext(filename)[0]
+                                if base_name in data:
+                                    data[base_name]['caption'] = content.decode('utf-8').strip()
+                                else:
+                                    data[base_name] = {'image': None, 'caption': content.decode('utf-8').strip()}
+                        
+                        elif file_type == 'depth':
+                            # Handle depth maps (PNG format)
+                            if filename.endswith('.png'):
+                                data[uid] = content
+                        
+                        elif file_type == 'normal':
+                            # Handle normal maps (RGB format)
+                            if filename.endswith('.jpg') or filename.endswith('.jpeg') or filename.endswith('.png'):
+                                data[uid] = content
+                        elif file_type == 'det_seg':
+                            # Handle detection maps (JSON format)
+                            if filename.endswith('.json'):
+                                # Decode bytes to string and parse JSON
+                                json_str = content.decode('utf-8')
+                                json_data = json.loads(json_str)
+                                data[uid] = json_data
+        except Exception as e:
+            logger.error(f"Error processing {tar_path}: {e}")
+            return {}
+            
+        return data
+    
+    def process_tar_file(self, rgb_caption_tar):
+        """Process a single tar file and create corresponding parquet"""
+        
+        # Extract base name (e.g., 'sa_000000' from 'sa_000000.tar')
+        base_name = os.path.splitext(os.path.basename(rgb_caption_tar))[0]
+        
+        # Construct paths for corresponding depth and normal tar files
+        depth_tar = os.path.join(self.depth_dir, f"{base_name}.tar")
+        normal_tar = os.path.join(self.normal_dir, f"{base_name}.tar")
+        det_seg_tar = os.path.join(self.det_seg_dir, f"{base_name}.tar")
+        
+        logger.info(f"Processing {base_name}")
+        
+        # Check if corresponding files exist
+        if not os.path.exists(depth_tar):
+            logger.warning(f"Depth tar file not found: {depth_tar}")
+            return None
+        if not os.path.exists(normal_tar):
+            logger.warning(f"Normal tar file not found: {normal_tar}")
+            return None
+        if not os.path.exists(det_seg_tar):
+            logger.warning(f"Det seg tar file not found: {det_seg_tar}")
+            return None
+        
+        # Extract data from all three tar files
+        logger.info(f"Extracting from rgb_caption: {rgb_caption_tar}")
+        rgb_caption_data = self.extract_from_tar(rgb_caption_tar, 'rgb_caption')
+        
+        logger.info(f"Extracting from depth: {depth_tar}")
+        depth_data = self.extract_from_tar(depth_tar, 'depth')
+        
+        logger.info(f"Extracting from normal: {normal_tar}")
+        normal_data = self.extract_from_tar(normal_tar, 'normal')
+        
+        logger.info(f"Extracting from det_seg: {det_seg_tar}")
+        det_seg_data = self.extract_from_tar(det_seg_tar, 'det_seg')
+        
+        # Find common UIDs
+        logger.info(f"rgb_caption_data.keys() - det_seg_data.keys(): {len(set(rgb_caption_data.keys())) - len(set(det_seg_data.keys()))}")
+        common_uids = set(rgb_caption_data.keys()) & set(depth_data.keys()) & set(normal_data.keys()) & set(det_seg_data.keys())
+
+        if not common_uids:
+            logger.warning(f"No common UIDs found for {base_name}")
+            return None
+        
+        logger.info(f"Found {len(common_uids)} common UIDs")
+        
+        # Sort UIDs for consistent ordering
+        sorted_uids = sorted(common_uids)
+        logger.info(f"Sorted {len(sorted_uids)} UIDs")
+        
+        # Create DataFrame
+        rows = []
+        for uid in tqdm(sorted_uids, desc=f"Processing {base_name}"):
+            rgb_info = rgb_caption_data[uid]
+            
+            # Skip if missing image or caption
+            if rgb_info['image'] is None or rgb_info['caption'] is None:
+                continue
+                
+            row = {
+                'uid': uid,
+                'rgb': rgb_info['image'],
+                'caption': rgb_info['caption'],
+                'depth': depth_data[uid],
+                'normal': normal_data[uid],
+                'det_seg': det_seg_data[uid]
+            }
+            rows.append(row)
+        
+        if not rows:
+            logger.warning(f"No valid rows for {base_name}")
+            return None
+        
+        # Create DataFrame
+        df = pd.DataFrame(rows)
+        
+        # Debug information
+        logger.info(f"DataFrame shape: {df.shape}")
+        if not df.empty:
+            logger.info(f"Sample UIDs: {df['uid'].head().tolist()}")
+        logger.info(f"Data types: {df.dtypes.to_dict()}")
+        
+        # Save to parquet with better settings for dataset compatibility
+        output_path = os.path.join(self.parquet_dir, f"{base_name}.parquet")
+        
+        # Use better parquet settings for dataset compatibility
+        df.to_parquet(
+            output_path, 
+            index=False,
+            engine='pyarrow',
+            compression='snappy',
+            row_group_size=1000  # Smaller row group size for more groups
+        )
+        logger.info(f"Saved {len(df)} rows to {output_path}")
+        
+        return output_path
+    
+    def process_all_tars(self):
+        """Process all tar files and create corresponding parquet files"""
+        tar_files = self.get_tar_files()
+        
+        if not tar_files:
+            logger.error("No tar files found in rgb_caption directory")
+            return
+        
+        logger.info(f"Found {len(tar_files)} tar files to process")
+        
+        processed_files = []
+        for tar_file in tar_files:
+            try:
+                output_path = self.process_tar_file(tar_file)
+                if output_path:
+                    processed_files.append(output_path)
+            except Exception as e:
+                logger.error(f"Error processing {tar_file}: {e}")
+        
+        logger.info(f"Successfully processed {len(processed_files)} files")
+        return processed_files
+
+def main():
+    # Set the base directory
+    base_dir = './datasets/blip3o'
+    
+    # Create processor
+    processor = TarProcessor(base_dir)
+    
+    # Process all tar files
+    processed_files = processor.process_all_tars()
+    
+    if processed_files:
+        logger.info(f"Successfully created {len(processed_files)} parquet files:")
+        for file in processed_files:
+            logger.info(f"  - {file}")
+    else:
+        logger.error("No parquet files were created")
+
+if __name__ == "__main__":
+    main()

data/any2any_preprocess/generate_parquet_clip448_imagebind.py

@@ -0,0 +1,318 @@
+import os
+import io
+import glob
+import tarfile
+import logging
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+class MultiModalEnricher:
+    """Read existing parquet files and add clip448 and imagebind columns from tar archives."""
+
+    def __init__(self, base_dir: str,
+                 input_parquet_dir: Optional[str] = None,
+                 clip448_dir: Optional[str] = None,
+                 imagebind_dir: Optional[str] = None,
+                 output_parquet_dir: Optional[str] = None) -> None:
+        self.base_dir = base_dir
+
+        # Defaults
+        self.input_parquet_dir = input_parquet_dir or os.path.join(
+            base_dir, 'parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global'
+        )
+        self.clip448_dir = clip448_dir or os.path.join(base_dir, 'datasets_clip448')
+        self.imagebind_dir = imagebind_dir or os.path.join(base_dir, 'datasets_imagebind')
+        self.output_parquet_dir = output_parquet_dir or os.path.join(
+            base_dir, 'parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global_clip448_imagebind'
+        )
+
+        os.makedirs(self.output_parquet_dir, exist_ok=True)
+
+    # ------------------------------------------------------------------
+    # Listing helpers
+    # ------------------------------------------------------------------
+
+    def _list_input_parquets(self) -> List[str]:
+        files = glob.glob(os.path.join(self.input_parquet_dir, '*.parquet'))
+        return sorted(files)
+
+    # ------------------------------------------------------------------
+    # Tar loaders
+    # ------------------------------------------------------------------
+
+    def _load_clip448_from_tar(self, tar_path: str) -> Dict[str, np.ndarray]:
+        """Load uid -> flattened local tokens from a clip448 tar.
+
+        Files inside: {uid}_tokens.npy  (shape (28,28) int64)
+        """
+        results: Dict[str, np.ndarray] = {}
+
+        if not os.path.exists(tar_path):
+            logger.warning(f"CLIP448 tar not found: {tar_path}")
+            return results
+
+        try:
+            with tarfile.open(tar_path, 'r') as tar:
+                for member in tar.getmembers():
+                    if not member.isfile():
+                        continue
+                    name = member.name
+                    if not name.endswith('_tokens.npy'):
+                        continue
+
+                    file_obj = tar.extractfile(member)
+                    if file_obj is None:
+                        continue
+
+                    content = file_obj.read()
+                    npy = np.load(io.BytesIO(content))
+
+                    # Derive uid: strip .npy then _tokens
+                    uid = os.path.splitext(name)[0]  # remove .npy
+                    if uid.endswith('_tokens'):
+                        uid = uid[:-len('_tokens')]
+
+                    results[uid] = npy.flatten()
+        except Exception as e:
+            logger.error(f"Failed to read CLIP448 tar {tar_path}: {e}")
+
+        return results
+
+    def _load_imagebind_from_tar(self, tar_path: str) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
+        """Load uid -> flattened tokens from an imagebind tar.
+
+        Files inside:
+          {uid}_tokens.npy          (local, shape (32,32) int64)
+          {uid}_tokens_global.npy   (global, shape (16,1,1) int64)
+
+        Returns (local_map, global_map).
+        """
+        local_results: Dict[str, np.ndarray] = {}
+        global_results: Dict[str, np.ndarray] = {}
+
+        if not os.path.exists(tar_path):
+            logger.warning(f"ImageBind tar not found: {tar_path}")
+            return local_results, global_results
+
+        try:
+            with tarfile.open(tar_path, 'r') as tar:
+                for member in tar.getmembers():
+                    if not member.isfile():
+                        continue
+                    name = member.name
+                    if not name.endswith('.npy'):
+                        continue
+
+                    file_obj = tar.extractfile(member)
+                    if file_obj is None:
+                        continue
+
+                    content = file_obj.read()
+                    npy = np.load(io.BytesIO(content))
+
+                    base = os.path.splitext(name)[0]  # remove .npy
+
+                    if base.endswith('_tokens_global'):
+                        uid = base[:-len('_tokens_global')]
+                        global_results[uid] = npy.flatten()
+                    elif base.endswith('_tokens'):
+                        uid = base[:-len('_tokens')]
+                        local_results[uid] = npy.flatten()
+        except Exception as e:
+            logger.error(f"Failed to read ImageBind tar {tar_path}: {e}")
+
+        return local_results, global_results
+
+    # ------------------------------------------------------------------
+    # Column addition
+    # ------------------------------------------------------------------
+
+    def _add_columns(self, df: pd.DataFrame,
+                     clip448_map: Dict[str, np.ndarray],
+                     ib_local_map: Dict[str, np.ndarray],
+                     ib_global_map: Dict[str, np.ndarray],
+                     strict: bool = False) -> pd.DataFrame:
+        """Add clip448, imagebind, and imagebind_global columns to *df*."""
+
+        if 'uid' not in df.columns:
+            raise ValueError("Input parquet missing required 'uid' column")
+
+        num_rows = len(df)
+        stats = {'clip448': 0, 'imagebind': 0, 'imagebind_global': 0,
+                 'clip448_miss': 0, 'imagebind_miss': 0, 'imagebind_global_miss': 0}
+
+        clip448_col: List[object] = []
+        ib_local_col: List[object] = []
+        ib_global_col: List[object] = []
+
+        for _, row in df.iterrows():
+            uid = row['uid']
+
+            # clip448
+            tokens = clip448_map.get(uid)
+            if tokens is None:
+                stats['clip448_miss'] += 1
+                clip448_col.append(row.get('clip448', None) if not strict else None)
+            else:
+                clip448_col.append(tokens)
+                stats['clip448'] += 1
+
+            # imagebind local
+            tokens = ib_local_map.get(uid)
+            if tokens is None:
+                stats['imagebind_miss'] += 1
+                ib_local_col.append(row.get('imagebind', None) if not strict else None)
+            else:
+                ib_local_col.append(tokens)
+                stats['imagebind'] += 1
+
+            # imagebind global
+            tokens = ib_global_map.get(uid)
+            if tokens is None:
+                stats['imagebind_global_miss'] += 1
+                ib_global_col.append(row.get('imagebind_global', None) if not strict else None)
+            else:
+                ib_global_col.append(tokens)
+                stats['imagebind_global'] += 1
+
+        logger.info(
+            f"Rows: {num_rows} | "
+            f"clip448: {stats['clip448']} ok / {stats['clip448_miss']} miss | "
+            f"imagebind: {stats['imagebind']} ok / {stats['imagebind_miss']} miss | "
+            f"imagebind_global: {stats['imagebind_global']} ok / {stats['imagebind_global_miss']} miss"
+        )
+
+        if strict and (stats['clip448_miss'] or stats['imagebind_miss'] or stats['imagebind_global_miss']):
+            logger.error("Strict mode: some UIDs are missing tokens from one or more modalities")
+
+        df = df.copy()
+        df['clip448'] = clip448_col
+        df['imagebind'] = ib_local_col
+        df['imagebind_global'] = ib_global_col
+        return df
+
+    # ------------------------------------------------------------------
+    # Processing
+    # ------------------------------------------------------------------
+
+    def process_one_parquet(self, parquet_path: str, overwrite: bool = False,
+                           strict: bool = False) -> Optional[str]:
+        base_name = os.path.splitext(os.path.basename(parquet_path))[0]
+        out_path = os.path.join(self.output_parquet_dir, f"{base_name}.parquet")
+
+        if os.path.exists(out_path) and not overwrite:
+            logger.info(f"Output exists, skipping: {out_path}")
+            return out_path
+
+        # Load existing parquet
+        df = pd.read_parquet(parquet_path, engine='pyarrow')
+
+        # Load clip448 tar
+        clip448_tar = os.path.join(self.clip448_dir, f"{base_name}.tar")
+        logger.info(f"Loading CLIP448 tokens from: {clip448_tar}")
+        clip448_map = self._load_clip448_from_tar(clip448_tar)
+
+        # Load imagebind tar
+        imagebind_tar = os.path.join(self.imagebind_dir, f"{base_name}.tar")
+        logger.info(f"Loading ImageBind tokens from: {imagebind_tar}")
+        ib_local_map, ib_global_map = self._load_imagebind_from_tar(imagebind_tar)
+
+        if not clip448_map:
+            logger.warning(f"No CLIP448 tokens loaded from {clip448_tar}")
+        if not ib_local_map:
+            logger.warning(f"No ImageBind local tokens loaded from {imagebind_tar}")
+        if not ib_global_map:
+            logger.warning(f"No ImageBind global tokens loaded from {imagebind_tar}")
+
+        # Add new columns
+        df = self._add_columns(df, clip448_map, ib_local_map, ib_global_map, strict=strict)
+
+        # Persist
+        df.to_parquet(
+            out_path,
+            index=False,
+            engine='pyarrow',
+            compression='snappy',
+            row_group_size=1000,
+        )
+        logger.info(f"Wrote enriched parquet to: {out_path}")
+        return out_path
+
+    def process_all(self, overwrite: bool = False, strict: bool = False,
+                    start_from: int = 0) -> List[str]:
+        inputs = self._list_input_parquets()
+        if start_from > 0:
+            inputs = inputs[start_from:]
+
+        if not inputs:
+            logger.error("No input parquet files found")
+            return []
+
+        logger.info(f"Found {len(inputs)} parquet files to enrich")
+        outputs: List[str] = []
+        for parquet_path in tqdm(inputs, desc="Enriching parquets", unit="file"):
+            try:
+                result = self.process_one_parquet(parquet_path, overwrite=overwrite, strict=strict)
+                if result:
+                    outputs.append(result)
+            except Exception as e:
+                logger.error(f"Error processing {parquet_path}: {e}")
+
+        logger.info(f"Successfully wrote {len(outputs)} enriched parquet files")
+        return outputs
+
+
+def main() -> None:
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description='Add CLIP448 and ImageBind tokens to existing parquet files'
+    )
+    parser.add_argument('--base_dir', type=str, default='./datasets/blip3o',
+                        help='Base directory (default: ./datasets/blip3o)')
+    parser.add_argument('--input_parquet_dir', type=str, default=None,
+                        help='Input parquet directory; defaults to parquet_*_dino_global')
+    parser.add_argument('--clip448_dir', type=str, default=None,
+                        help='Directory containing clip448 tars; defaults to base_dir/datasets_clip448')
+    parser.add_argument('--imagebind_dir', type=str, default=None,
+                        help='Directory containing imagebind tars; defaults to base_dir/datasets_imagebind')
+    parser.add_argument('--output_parquet_dir', type=str, default=None,
+                        help='Output parquet directory; defaults to parquet_*_clip448_imagebind')
+    parser.add_argument('--overwrite', action='store_true',
+                        help='Overwrite existing output parquet if present')
+    parser.add_argument('--strict', action='store_true',
+                        help='If set, missing UIDs will set column=None and log error')
+    parser.add_argument('--start_from', type=int, default=0,
+                        help='Index to start processing from (for resuming)')
+
+    args = parser.parse_args()
+
+    enricher = MultiModalEnricher(
+        base_dir=args.base_dir,
+        input_parquet_dir=args.input_parquet_dir,
+        clip448_dir=args.clip448_dir,
+        imagebind_dir=args.imagebind_dir,
+        output_parquet_dir=args.output_parquet_dir,
+    )
+
+    outputs = enricher.process_all(overwrite=args.overwrite, strict=args.strict, start_from=args.start_from)
+    if outputs:
+        logger.info("Outputs written:")
+        for p in outputs:
+            logger.info(f"  - {p}")
+    else:
+        logger.error("No outputs were created")
+
+
+if __name__ == '__main__':
+    main()
+

data/any2any_preprocess/generate_parquet_grounding.py

@@ -0,0 +1,583 @@
+import os
+import tarfile
+import pandas as pd
+import numpy as np
+from PIL import Image
+import io
+import glob
+import json
+from tqdm import tqdm
+import logging
+import pycocotools.mask as mask_util
+import matplotlib.pyplot as plt
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class GroundingProcessor:
+    def __init__(self, base_dir):
+        self.base_dir = base_dir
+        self.rgb_caption_dir = os.path.join(base_dir, 'datasets')
+        self.depth_dir = os.path.join(base_dir, 'datasets_depth')
+        self.normal_dir = os.path.join(base_dir, 'datasets_normal')
+        self.det_seg_dir = os.path.join(base_dir, 'datasets_seg_swinb')
+        self.grounding_dir = os.path.join(base_dir, 'datasets_grounding')  # Pre-processed grounding tars
+        self.parquet_dir = os.path.join(base_dir, 'parquet_rgb_caption_depth_normal_det_seg_grounding')
+        
+        # Create parquet directory if it doesn't exist
+        os.makedirs(self.parquet_dir, exist_ok=True)
+        
+    def get_tar_files(self):
+        """Get all tar files from rgb_caption directory"""
+        tar_files = glob.glob(os.path.join(self.rgb_caption_dir, '*.tar'))
+        return sorted(tar_files)
+    
+    
+    def extract_from_tar(self, tar_path, file_type):
+        """
+        Extract data from tar file
+        
+        Args:
+            tar_path: Path to tar file
+            file_type: 'rgb_caption', 'depth', 'normal', 'det_seg', 'grounding'
+        """
+        data = {}
+        
+        try:
+            with tarfile.open(tar_path, 'r') as tar:
+                for member in tar.getmembers():
+                    if member.isfile():
+                        filename = member.name
+                        uid = os.path.splitext(filename)[0]  # Remove extension
+                        
+                        # Extract file content
+                        f = tar.extractfile(member)
+                        if f is None:
+                            continue
+                        
+                        content = f.read()
+                        
+                        if file_type == 'rgb_caption':
+                            # Handle RGB + caption files
+                            if filename.endswith('.jpg') or filename.endswith('.jpeg') or filename.endswith('.png'):
+                                data[uid] = {'image': content, 'caption': None}
+                            elif filename.endswith('.txt'):
+                                # Find corresponding image
+                                base_name = os.path.splitext(filename)[0]
+                                if base_name in data:
+                                    data[base_name]['caption'] = content.decode('utf-8').strip()
+                                else:
+                                    data[base_name] = {'image': None, 'caption': content.decode('utf-8').strip()}
+                        
+                        elif file_type == 'depth':
+                            # Handle depth maps (PNG format)
+                            if filename.endswith('.png'):
+                                data[uid] = content
+                        
+                        elif file_type == 'normal':
+                            # Handle normal maps (RGB format)
+                            if filename.endswith('.jpg') or filename.endswith('.jpeg') or filename.endswith('.png'):
+                                data[uid] = content
+                        
+                        elif file_type == 'det_seg':
+                            # Handle detection maps (JSON format)
+                            if filename.endswith('.json'):
+                                # Decode bytes to string and parse JSON
+                                json_str = content.decode('utf-8')
+                                json_data = json.loads(json_str)
+                                data[uid] = json_data
+                        
+                        elif file_type == 'grounding':
+                            # Handle grounding JSON files
+                            if filename.endswith('.json'):
+                                # Decode bytes to string and parse JSON
+                                json_str = content.decode('utf-8')
+                                json_data = json.loads(json_str)
+                                data[uid] = json_data
+        except Exception as e:
+            logger.error(f"Error processing {tar_path}: {e}")
+            return {}
+            
+        return data
+    
+    
+    def extract_grounding_data(self, grounding_data):
+        """Extract grounding information from JSON data"""
+        if not grounding_data:
+            return None
+            
+        # Get the first (and usually only) key which is the image name
+        image_key = list(grounding_data.keys())[0]
+        image_data = grounding_data[image_key]
+        
+        # Extract objects and their bounding boxes
+        objects = image_data.get('objects', [])
+        floating_objects = image_data.get('floating_objects', [])
+        
+        # Extract grounding relationships
+        relationships = image_data.get('relationships', {})
+        grounding_phrases = relationships.get('grounding', [])
+        
+        # Extract captions with grounding details
+        short_captions = image_data.get('short_captions', [])
+        dense_caption = image_data.get('dense_caption', {})
+        
+        return {
+            'objects': objects,
+            'floating_objects': floating_objects,
+            'grounding_phrases': grounding_phrases,
+            'short_captions': short_captions,
+            'dense_caption': dense_caption
+        }
+    
+    def normalize_bbox(self, bbox, image_size):
+        """Normalize bbox to 0-999 based on image dimensions"""
+        if not bbox or len(bbox) != 4:
+            return []
+        
+        # RLE format: size = [width, height]
+        img_height, img_width = image_size
+        
+        x1, y1, x2, y2 = bbox
+        
+        # Normalize to 0-999 and clamp values
+        norm_x1 = max(0, min(999, int((x1 / img_width) * 999)))
+        norm_y1 = max(0, min(999, int((y1 / img_height) * 999)))
+        norm_x2 = max(0, min(999, int((x2 / img_width) * 999)))
+        norm_y2 = max(0, min(999, int((y2 / img_height) * 999)))
+        
+        return [norm_x1, norm_y1, norm_x2, norm_y2]
+
+
+    def resize_and_crop_bbox(self, bbox, orig_size, target_size=512):
+        """
+        Resize and center-crop a bounding box using CLIP/BLIP-style transform.
+
+        Args:
+            bbox: list or tuple [x1, y1, x2, y2] in pixel coords.
+            orig_size: [W, H] original image size.
+            target_size: int, final square crop size (default=512).
+
+        Returns:
+            [x1', y1', x2', y2'] after resize+crop.
+            Also returns (scale, left, top) for debugging.
+        """
+        if not bbox or len(bbox) != 4:
+            return []
+
+        H, W = orig_size
+        x1, y1, x2, y2 = bbox
+
+        # --- 1. Resize (keep aspect ratio, short side -> target_size) ---
+        scale = target_size / min(W, H)
+        new_W, new_H = W * scale, H * scale
+
+        # --- 2. Center crop to target_size×target_size ---
+        left = max(0, (new_W - target_size) / 2)
+        top = max(0, (new_H - target_size) / 2)
+
+        # --- 3. Apply transform to bbox ---
+        x1_new = x1 * scale - left
+        y1_new = y1 * scale - top
+        x2_new = x2 * scale - left
+        y2_new = y2 * scale - top
+
+        # Clip to [0, target_size]
+        x1_new, y1_new, x2_new, y2_new = np.clip(
+            [x1_new, y1_new, x2_new, y2_new],
+            0, target_size
+        )
+
+        return [x1_new, y1_new, x2_new, y2_new]
+
+
+    def visualize_rle_mask(self, rle_data, save_path=None):
+        """Visualize RLE mask from segmentation data"""
+        try:
+            # Decode RLE to binary mask
+            mask = mask_util.decode(rle_data)
+            
+            # Create visualization
+            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
+            
+            # Show mask
+            ax1.imshow(mask, cmap='gray')
+            ax1.set_title('RLE Mask')
+            ax1.axis('off')
+            
+            # Show mask with bbox overlay
+            ax2.imshow(mask, cmap='gray', alpha=0.7)
+            
+            # Add bbox if available
+            if 'bbox' in rle_data:
+                bbox = rle_data['bbox']
+                x1, y1, x2, y2 = bbox
+                width = x2 - x1
+                height = y2 - y1
+                rect = plt.Rectangle((x1, y1), width, height, 
+                                  fill=False, edgecolor='red', linewidth=2)
+                ax2.add_patch(rect)
+                ax2.set_title('Mask + Bbox')
+            else:
+                ax2.set_title('Mask Only')
+            ax2.axis('off')
+            
+            plt.tight_layout()
+            
+            if save_path:
+                plt.savefig(save_path, dpi=150, bbox_inches='tight')
+                logger.info(f"Mask visualization saved to {save_path}")
+            else:
+                plt.show()
+            
+            plt.close()
+            
+        except Exception as e:
+            logger.error(f"Error visualizing RLE mask: {e}")
+    
+    def visualize_object_mask(self, grounding_data, uid, object_id=0, save_path=None):
+        """Visualize a specific object's mask from grounding data"""
+        try:
+            if not grounding_data or 'objects' not in grounding_data:
+                logger.warning(f"No objects found for {uid}")
+                return
+            
+            objects = grounding_data['objects']
+            if object_id >= len(objects):
+                logger.warning(f"Object ID {object_id} not found for {uid}")
+                return
+            
+            obj = objects[object_id]
+            if 'segmentation' not in obj:
+                logger.warning(f"No segmentation found for object {object_id} in {uid}")
+                return
+            
+            segmentation = obj['segmentation']
+            bbox = obj.get('bbox', [])
+            
+            # Create RLE data for visualization
+            rle_data = {
+                'size': segmentation['size'],
+                'counts': segmentation['counts'],
+                'bbox': bbox
+            }
+            
+            # Visualize
+            if save_path:
+                save_path_with_id = save_path.replace('.png', f'_obj_{object_id}.png')
+            else:
+                save_path_with_id = None
+                
+            self.visualize_rle_mask(rle_data, save_path_with_id)
+            
+        except Exception as e:
+            logger.error(f"Error visualizing object mask for {uid}: {e}")
+    
+    def create_grounding_samples(self, grounding_data, uid):
+        """Create training samples from grounding data"""
+        if not grounding_data:
+            return []
+        
+        samples = []
+        
+        # Get image size from first object's segmentation
+        image_size = None
+        if grounding_data['objects'] and 'segmentation' in grounding_data['objects'][0]:
+            image_size = grounding_data['objects'][0]['segmentation']['size']
+            assert image_size is not None
+        else:
+            logger.warning(f"No segmentation data found for {uid}, skipping bbox processing")
+            return []
+        
+        # Process grounding phrases
+        for phrase_info in grounding_data['grounding_phrases']:
+            phrase = phrase_info['phrase']
+            object_ids = phrase_info['object_ids']
+            
+            # Find corresponding objects and their bboxes
+            for obj_id in object_ids:
+                for obj in grounding_data['objects']:
+                    if obj.get('id') == obj_id:
+                        bbox = obj.get('bbox', [])
+                        
+                        # Process bbox if valid
+                        if bbox and len(bbox) == 4:
+                            bbox = self.resize_and_crop_bbox(bbox, image_size, 512)
+                            bbox = self.normalize_bbox(bbox, [512, 512])
+                        else:
+                            bbox = []
+                        
+                        sample = {
+                            'uid': uid,
+                            'phrase': phrase,
+                            'object_id': obj_id,
+                            'bbox': bbox,
+                            'labels': obj.get('labels', []),
+                            'attributes': obj.get('attributes', []),
+                            'sample_type': 'grounding_phrase',
+                            'segmentation': obj.get('segmentation', []),
+                        }
+                        samples.append(sample)
+                        break
+        
+        # Process short captions with details
+        for caption_info in grounding_data['short_captions']:
+            caption = caption_info['caption']
+            for detail in caption_info.get('details', []):
+                phrase = detail['phrase']
+                object_id = detail.get('id')
+                bbox = detail.get('bbox', [])
+                
+                # Process bbox if valid
+                if bbox and len(bbox) == 4:
+                    bbox = self.resize_and_crop_bbox(bbox, image_size, 512)
+                    bbox = self.normalize_bbox(bbox, [512, 512])
+                else:
+                    bbox = []
+                
+                tokens_positive = detail.get('tokens_positive', [])
+                
+                sample = {
+                    'uid': uid,
+                    'phrase': phrase,
+                    'caption_id': object_id,
+                    'bbox': bbox,
+                    'tokens_positive': tokens_positive,
+                    'full_caption': caption,
+                    'sample_type': 'short_caption'
+                }
+                samples.append(sample)
+        
+        # # Process dense caption
+        # if grounding_data['dense_caption']:
+        #     dense_caption = grounding_data['dense_caption']
+        #     caption = dense_caption.get('caption', '')
+        #     for detail in dense_caption.get('details', []):
+        #         phrase = detail['phrase']
+        #         object_ids = detail.get('ids', [])
+        #         bboxes = detail.get('bbox', [])
+        #         tokens_positive = detail.get('tokens_positive', [])
+                
+        #         sample = {
+        #             'uid': uid,
+        #             'phrase': phrase,
+        #             'object_ids': object_ids,
+        #             'bboxes': bboxes,
+        #             'tokens_positive': tokens_positive,
+        #             'full_caption': caption,
+        #             'sample_type': 'dense_caption'
+        #         }
+        #         samples.append(sample)
+        
+        return samples
+    
+    def process_tar_file(self, rgb_caption_tar):
+        """Process a single tar file and create corresponding parquet"""
+        
+        # Extract base name (e.g., 'sa_000000' from 'sa_000000.tar')
+        base_name = os.path.splitext(os.path.basename(rgb_caption_tar))[0]
+        
+        # Construct paths for corresponding depth, normal, det_seg, grounding tar files
+        depth_tar = os.path.join(self.depth_dir, f"{base_name}.tar")
+        normal_tar = os.path.join(self.normal_dir, f"{base_name}.tar")
+        det_seg_tar = os.path.join(self.det_seg_dir, f"{base_name}.tar")
+        grounding_tar = os.path.join(self.grounding_dir, f"{base_name}.tar")
+        
+        logger.info(f"Processing {base_name}")
+        
+        # Check if corresponding files exist
+        if not os.path.exists(depth_tar):
+            logger.warning(f"Depth tar file not found: {depth_tar}")
+            return None
+        if not os.path.exists(normal_tar):
+            logger.warning(f"Normal tar file not found: {normal_tar}")
+            return None
+        if not os.path.exists(det_seg_tar):
+            logger.warning(f"Det seg tar file not found: {det_seg_tar}")
+            return None
+        # Initialize grounding_data
+        grounding_data = {}
+        
+        if not os.path.exists(grounding_tar):
+            logger.warning(f"Grounding tar file not found: {grounding_tar} - will process without grounding data")
+        else:
+            logger.info(f"Extracting from grounding: {grounding_tar}")
+            grounding_data = self.extract_from_tar(grounding_tar, 'grounding')
+        
+        # Extract data from all tar files
+        logger.info(f"Extracting from rgb_caption: {rgb_caption_tar}")
+        rgb_caption_data = self.extract_from_tar(rgb_caption_tar, 'rgb_caption')
+        
+        logger.info(f"Extracting from depth: {depth_tar}")
+        depth_data = self.extract_from_tar(depth_tar, 'depth')
+        
+        logger.info(f"Extracting from normal: {normal_tar}")
+        normal_data = self.extract_from_tar(normal_tar, 'normal')
+        
+        logger.info(f"Extracting from det_seg: {det_seg_tar}")
+        det_seg_data = self.extract_from_tar(det_seg_tar, 'det_seg')
+        
+        # Find common UIDs (grounding is optional)
+        common_uids = set(rgb_caption_data.keys()) & set(depth_data.keys()) & set(normal_data.keys()) & set(det_seg_data.keys())
+        
+        if not common_uids:
+            logger.warning(f"No common UIDs found for {base_name}")
+            return None
+        
+        logger.info(f"Found {len(common_uids)} common UIDs")
+        
+        # Sort UIDs for consistent ordering
+        sorted_uids = sorted(common_uids)
+        
+        # Create samples for each UID
+        all_samples = []
+        
+        for uid in sorted_uids:
+            rgb_info = rgb_caption_data[uid]
+            
+            # Skip if missing image or caption
+            if rgb_info['image'] is None or rgb_info['caption'] is None:
+                continue
+            
+            # Create base sample with all modalities
+            sample = {
+                'uid': uid,
+                'rgb': rgb_caption_data[uid]['image'],
+                'caption': rgb_caption_data[uid]['caption'],
+                'depth': depth_data[uid],
+                'normal': normal_data[uid],
+                'det_seg': det_seg_data[uid]
+            }
+            
+            # Add grounding data if available
+            if uid in grounding_data:
+                extracted_grounding = self.extract_grounding_data(grounding_data[uid])
+                if extracted_grounding:
+                    raw_image = Image.open(io.BytesIO(rgb_caption_data[uid]['image']))
+                    raw_image_size = raw_image.size
+                    assert raw_image_size == (512, 512), f"Raw image size is {raw_image_size} but expected (512, 512) for {uid}"
+                    grounding_samples = self.create_grounding_samples(extracted_grounding, uid)
+                    if grounding_samples:
+                        sample['grounding'] = grounding_samples
+            
+            # visualize_segmentation = sample['grounding'][1]['segmentation'] 
+            # self.visualize_rle_mask(visualize_segmentation,save_path=f"visualize_{uid}_1.png")
+            all_samples.append(sample)
+        
+        if not all_samples:
+            logger.warning(f"No samples found for {base_name}")
+            return None
+        
+        # Create DataFrame
+        df = pd.DataFrame(all_samples)
+        
+        # Debug information
+        logger.info(f"DataFrame shape: {df.shape}")
+        if not df.empty:
+            logger.info(f"Sample UIDs: {df['uid'].head().tolist()}")
+            # Count grounding samples
+            grounding_counts = []
+            for _, row in df.iterrows():
+                if 'grounding' in row and row['grounding'] is not None and isinstance(row['grounding'], list):
+                    grounding_counts.append(len(row['grounding']))
+                else:
+                    grounding_counts.append(0)
+            if grounding_counts:
+                logger.info(f"Grounding samples per row: {np.mean(grounding_counts):.2f} avg, {max(grounding_counts)} max")
+            else:
+                logger.info("No grounding samples found")
+        
+        # Save to parquet
+        output_path = os.path.join(self.parquet_dir, f"{base_name}.parquet")
+        
+        df.to_parquet(
+            output_path, 
+            index=False,
+            engine='pyarrow',
+            compression='snappy',
+            row_group_size=1000
+        )
+        logger.info(f"Saved {len(df)} grounding samples to {output_path}")
+        
+        return output_path
+    
+    def process_all_tars(self):
+        """Process all tar files and create corresponding parquet files"""
+        tar_files = self.get_tar_files()
+        
+        if not tar_files:
+            logger.error("No tar files found in rgb_caption directory")
+            return
+        
+        logger.info(f"Found {len(tar_files)} tar files to process")
+        
+        processed_files = []
+        for tar_file in tar_files:
+            try:
+                output_path = self.process_tar_file(tar_file)
+                if output_path:
+                    processed_files.append(output_path)
+            except Exception as e:
+                logger.error(f"Error processing {tar_file}: {e}")
+        
+        logger.info(f"Successfully processed {len(processed_files)} files")
+        return processed_files
+
+def main():
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Generate parquet files with grounding data')
+    parser.add_argument('--base_dir', type=str, default='./datasets/blip3o',
+                       help='Base directory (default: ./datasets/blip3o)')
+    parser.add_argument('--visualize_mask', action='store_true',
+                       help='Visualize RLE masks for debugging')
+    parser.add_argument('--visualize_tar', type=str, default=None,
+                       help='Specific tar file to visualize (e.g., sa_000001.tar)')
+    parser.add_argument('--visualize_uid', type=str, default=None,
+                       help='Specific UID to visualize (e.g., sa_1)')
+    parser.add_argument('--visualize_object_id', type=int, default=0,
+                       help='Object ID to visualize (default: 0)')
+    
+    args = parser.parse_args()
+    
+    # Set the base directory
+    base_dir = args.base_dir
+    
+    # Create processor
+    processor = GroundingProcessor(base_dir)
+    
+    # Visualize mask if requested
+    if args.visualize_mask and args.visualize_tar:
+        logger.info(f"Visualizing mask for {args.visualize_tar}, object {args.visualize_object_id}")
+        
+        # Find the grounding tar file for this UID
+        grounding_tar = os.path.join(processor.grounding_dir, f"{args.visualize_tar}.tar")
+        if os.path.exists(grounding_tar):
+            grounding_data = processor.extract_from_tar(grounding_tar, 'grounding')
+            if args.visualize_uid in grounding_data:
+                extracted_grounding = processor.extract_grounding_data(grounding_data[args.visualize_uid])
+                if extracted_grounding:
+                    save_path = f"mask_visualization_{args.visualize_uid}.png"
+                    processor.visualize_object_mask(extracted_grounding, args.visualize_uid, 
+                                                  args.visualize_object_id, save_path)
+                else:
+                    logger.error(f"No grounding data found for {args.visualize_uid}")
+            else:
+                logger.error(f"UID {args.visualize_uid} not found in grounding tar")
+        else:
+            logger.error(f"Grounding tar not found: {grounding_tar}")
+        return
+    
+    # Process all tar files
+    processed_files = processor.process_all_tars()
+    
+    if processed_files:
+        logger.info(f"Successfully created {len(processed_files)} grounding parquet files:")
+        for file in processed_files:
+            logger.info(f"  - {file}")
+    else:
+        logger.error("No grounding parquet files were created")
+
+if __name__ == "__main__":
+    main()

data/any2any_preprocess/generate_parquet_grounding_canny_dino.py

@@ -0,0 +1,627 @@
+import os
+import tarfile
+import pandas as pd
+import numpy as np
+from PIL import Image
+import io
+import glob
+import json
+from tqdm import tqdm
+import logging
+import pycocotools.mask as mask_util
+import matplotlib.pyplot as plt
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class GroundingProcessor:
+    def __init__(self, base_dir):
+        self.base_dir = base_dir
+        self.rgb_caption_dir = os.path.join(base_dir, 'datasets')
+        self.depth_dir = os.path.join(base_dir, 'datasets_depth')
+        self.normal_dir = os.path.join(base_dir, 'datasets_normal')
+        self.det_seg_dir = os.path.join(base_dir, 'datasets_seg_swinb')
+        self.grounding_dir = os.path.join(base_dir, 'datasets_grounding')  # Pre-processed grounding tars
+        self.canny_dir = os.path.join(base_dir, 'datasets_cannyedge')
+        self.dino_dir = os.path.join(base_dir, 'dinov2')
+        self.parquet_dir = os.path.join(base_dir, 'parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino')
+        
+        # Create parquet directory if it doesn't exist
+        os.makedirs(self.parquet_dir, exist_ok=True)
+        
+    def get_tar_files(self):
+        """Get all tar files from rgb_caption directory"""
+        tar_files = glob.glob(os.path.join(self.rgb_caption_dir, '*.tar'))
+        return sorted(tar_files)
+    
+    
+    def extract_from_tar(self, tar_path, file_type):
+        """
+        Extract data from tar file
+        
+        Args:
+            tar_path: Path to tar file
+            file_type: 'rgb_caption', 'depth', 'normal', 'det_seg', 'grounding', 'canny', 'dino'
+        """
+        data = {}
+        
+        try:
+            with tarfile.open(tar_path, 'r') as tar:
+                for member in tar.getmembers():
+                    if member.isfile():
+                        filename = member.name
+                        uid = os.path.splitext(filename)[0]  # Remove extension
+                        
+                        # Extract file content
+                        f = tar.extractfile(member)
+                        if f is None:
+                            continue
+                        
+                        content = f.read()
+                        
+                        if file_type == 'rgb_caption':
+                            # Handle RGB + caption files
+                            if filename.endswith('.jpg') or filename.endswith('.jpeg') or filename.endswith('.png'):
+                                data[uid] = {'image': content, 'caption': None}
+                            elif filename.endswith('.txt'):
+                                # Find corresponding image
+                                base_name = os.path.splitext(filename)[0]
+                                if base_name in data:
+                                    data[base_name]['caption'] = content.decode('utf-8').strip()
+                                else:
+                                    data[base_name] = {'image': None, 'caption': content.decode('utf-8').strip()}
+                        
+                        elif file_type == 'depth':
+                            # Handle depth maps (PNG format)
+                            if filename.endswith('.png'):
+                                data[uid] = content
+                        
+                        elif file_type == 'normal':
+                            # Handle normal maps (RGB format)
+                            if filename.endswith('.jpg') or filename.endswith('.jpeg') or filename.endswith('.png'):
+                                data[uid] = content
+                        
+                        elif file_type == 'det_seg':
+                            # Handle detection maps (JSON format)
+                            if filename.endswith('.json'):
+                                # Decode bytes to string and parse JSON
+                                json_str = content.decode('utf-8')
+                                json_data = json.loads(json_str)
+                                data[uid] = json_data
+                        
+                        elif file_type == 'grounding':
+                            # Handle grounding JSON files
+                            if filename.endswith('.json'):
+                                # Decode bytes to string and parse JSON
+                                json_str = content.decode('utf-8')
+                                json_data = json.loads(json_str)
+                                data[uid] = json_data
+                        
+                        elif file_type == 'canny':
+                            # Handle Canny maps (PNG format, similar to depth)
+                            if filename.endswith('.png'):
+                                data[uid] = content
+                        
+                        elif file_type == 'dino':
+                            # Handle DINOv2 features (NPY format)
+                            if filename.endswith('_tokens.npy'):
+                                # Load NPY data from bytes
+                                npy_data = np.load(io.BytesIO(content))
+                                uid = uid[:-len('_tokens')]
+                                # Flatten the array for parquet storage
+                                data[uid] = npy_data.flatten()
+        except Exception as e:
+            logger.error(f"Error processing {tar_path}: {e}")
+            return {}
+            
+        return data
+    
+    def extract_grounding_data(self, grounding_data):
+        """Extract grounding information from JSON data"""
+        if not grounding_data:
+            return None
+            
+        # Get the first (and usually only) key which is the image name
+        image_key = list(grounding_data.keys())[0]
+        image_data = grounding_data[image_key]
+        
+        # Extract objects and their bounding boxes
+        objects = image_data.get('objects', [])
+        floating_objects = image_data.get('floating_objects', [])
+        
+        # Extract grounding relationships
+        relationships = image_data.get('relationships', {})
+        grounding_phrases = relationships.get('grounding', [])
+        
+        # Extract captions with grounding details
+        short_captions = image_data.get('short_captions', [])
+        dense_caption = image_data.get('dense_caption', {})
+        
+        return {
+            'objects': objects,
+            'floating_objects': floating_objects,
+            'grounding_phrases': grounding_phrases,
+            'short_captions': short_captions,
+            'dense_caption': dense_caption
+        }
+    
+    def normalize_bbox(self, bbox, image_size):
+        """Normalize bbox to 0-999 based on image dimensions"""
+        if not bbox or len(bbox) != 4:
+            return []
+        
+        # RLE format: size = [width, height]
+        img_height, img_width = image_size
+        
+        x1, y1, x2, y2 = bbox
+        
+        # Normalize to 0-999 and clamp values
+        norm_x1 = max(0, min(999, int((x1 / img_width) * 999)))
+        norm_y1 = max(0, min(999, int((y1 / img_height) * 999)))
+        norm_x2 = max(0, min(999, int((x2 / img_width) * 999)))
+        norm_y2 = max(0, min(999, int((y2 / img_height) * 999)))
+        
+        return [norm_x1, norm_y1, norm_x2, norm_y2]
+
+
+    def resize_and_crop_bbox(self, bbox, orig_size, target_size=512):
+        """
+        Resize and center-crop a bounding box using CLIP/BLIP-style transform.
+
+        Args:
+            bbox: list or tuple [x1, y1, x2, y2] in pixel coords.
+            orig_size: [W, H] original image size.
+            target_size: int, final square crop size (default=512).
+
+        Returns:
+            [x1', y1', x2', y2'] after resize+crop.
+            Also returns (scale, left, top) for debugging.
+        """
+        if not bbox or len(bbox) != 4:
+            return []
+
+        H, W = orig_size
+        x1, y1, x2, y2 = bbox
+
+        # --- 1. Resize (keep aspect ratio, short side -> target_size) ---
+        scale = target_size / min(W, H)
+        new_W, new_H = W * scale, H * scale
+
+        # --- 2. Center crop to target_size×target_size ---
+        left = max(0, (new_W - target_size) / 2)
+        top = max(0, (new_H - target_size) / 2)
+
+        # --- 3. Apply transform to bbox ---
+        x1_new = x1 * scale - left
+        y1_new = y1 * scale - top
+        x2_new = x2 * scale - left
+        y2_new = y2 * scale - top
+
+        # Clip to [0, target_size]
+        x1_new, y1_new, x2_new, y2_new = np.clip(
+            [x1_new, y1_new, x2_new, y2_new],
+            0, target_size
+        )
+
+        return [x1_new, y1_new, x2_new, y2_new]
+
+
+    def visualize_rle_mask(self, rle_data, save_path=None):
+        """Visualize RLE mask from segmentation data"""
+        try:
+            # Decode RLE to binary mask
+            mask = mask_util.decode(rle_data)
+            
+            # Create visualization
+            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
+            
+            # Show mask
+            ax1.imshow(mask, cmap='gray')
+            ax1.set_title('RLE Mask')
+            ax1.axis('off')
+            
+            # Show mask with bbox overlay
+            ax2.imshow(mask, cmap='gray', alpha=0.7)
+            
+            # Add bbox if available
+            if 'bbox' in rle_data:
+                bbox = rle_data['bbox']
+                x1, y1, x2, y2 = bbox
+                width = x2 - x1
+                height = y2 - y1
+                rect = plt.Rectangle((x1, y1), width, height, 
+                                  fill=False, edgecolor='red', linewidth=2)
+                ax2.add_patch(rect)
+                ax2.set_title('Mask + Bbox')
+            else:
+                ax2.set_title('Mask Only')
+            ax2.axis('off')
+            
+            plt.tight_layout()
+            
+            if save_path:
+                plt.savefig(save_path, dpi=150, bbox_inches='tight')
+                logger.info(f"Mask visualization saved to {save_path}")
+            else:
+                plt.show()
+            
+            plt.close()
+            
+        except Exception as e:
+            logger.error(f"Error visualizing RLE mask: {e}")
+    
+    def visualize_object_mask(self, grounding_data, uid, object_id=0, save_path=None):
+        """Visualize a specific object's mask from grounding data"""
+        try:
+            if not grounding_data or 'objects' not in grounding_data:
+                logger.warning(f"No objects found for {uid}")
+                return
+            
+            objects = grounding_data['objects']
+            if object_id >= len(objects):
+                logger.warning(f"Object ID {object_id} not found for {uid}")
+                return
+            
+            obj = objects[object_id]
+            if 'segmentation' not in obj:
+                logger.warning(f"No segmentation found for object {object_id} in {uid}")
+                return
+            
+            segmentation = obj['segmentation']
+            bbox = obj.get('bbox', [])
+            
+            # Create RLE data for visualization
+            rle_data = {
+                'size': segmentation['size'],
+                'counts': segmentation['counts'],
+                'bbox': bbox
+            }
+            
+            # Visualize
+            if save_path:
+                save_path_with_id = save_path.replace('.png', f'_obj_{object_id}.png')
+            else:
+                save_path_with_id = None
+                
+            self.visualize_rle_mask(rle_data, save_path_with_id)
+            
+        except Exception as e:
+            logger.error(f"Error visualizing object mask for {uid}: {e}")
+    
+    def create_grounding_samples(self, grounding_data, uid):
+        """Create training samples from grounding data"""
+        if not grounding_data:
+            return []
+        
+        samples = []
+        
+        # Get image size from first object's segmentation
+        image_size = None
+        if grounding_data['objects'] and 'segmentation' in grounding_data['objects'][0]:
+            image_size = grounding_data['objects'][0]['segmentation']['size']
+            assert image_size is not None
+        else:
+            logger.warning(f"No segmentation data found for {uid}, skipping bbox processing")
+            return []
+        
+        # Process grounding phrases
+        for phrase_info in grounding_data['grounding_phrases']:
+            phrase = phrase_info['phrase']
+            object_ids = phrase_info['object_ids']
+            
+            # Find corresponding objects and their bboxes
+            for obj_id in object_ids:
+                for obj in grounding_data['objects']:
+                    if obj.get('id') == obj_id:
+                        bbox = obj.get('bbox', [])
+                        
+                        # Process bbox if valid
+                        if bbox and len(bbox) == 4:
+                            bbox = self.resize_and_crop_bbox(bbox, image_size, 512)
+                            bbox = self.normalize_bbox(bbox, [512, 512])
+                        else:
+                            bbox = []
+                        
+                        sample = {
+                            'uid': uid,
+                            'phrase': phrase,
+                            'object_id': obj_id,
+                            'bbox': bbox,
+                            'labels': obj.get('labels', []),
+                            'attributes': obj.get('attributes', []),
+                            'sample_type': 'grounding_phrase',
+                            'segmentation': obj.get('segmentation', []),
+                        }
+                        samples.append(sample)
+                        break
+        
+        # Process short captions with details
+        for caption_info in grounding_data['short_captions']:
+            caption = caption_info['caption']
+            for detail in caption_info.get('details', []):
+                phrase = detail['phrase']
+                object_id = detail.get('id')
+                bbox = detail.get('bbox', [])
+                
+                # Process bbox if valid
+                if bbox and len(bbox) == 4:
+                    bbox = self.resize_and_crop_bbox(bbox, image_size, 512)
+                    bbox = self.normalize_bbox(bbox, [512, 512])
+                else:
+                    bbox = []
+                
+                tokens_positive = detail.get('tokens_positive', [])
+                
+                sample = {
+                    'uid': uid,
+                    'phrase': phrase,
+                    'caption_id': object_id,
+                    'bbox': bbox,
+                    'tokens_positive': tokens_positive,
+                    'full_caption': caption,
+                    'sample_type': 'short_caption'
+                }
+                samples.append(sample)
+        
+        # # Process dense caption
+        # if grounding_data['dense_caption']:
+        #     dense_caption = grounding_data['dense_caption']
+        #     caption = dense_caption.get('caption', '')
+        #     for detail in dense_caption.get('details', []):
+        #         phrase = detail['phrase']
+        #         object_ids = detail.get('ids', [])
+        #         bboxes = detail.get('bbox', [])
+        #         tokens_positive = detail.get('tokens_positive', [])
+                
+        #         sample = {
+        #             'uid': uid,
+        #             'phrase': phrase,
+        #             'object_ids': object_ids,
+        #             'bboxes': bboxes,
+        #             'tokens_positive': tokens_positive,
+        #             'full_caption': caption,
+        #             'sample_type': 'dense_caption'
+        #         }
+        #         samples.append(sample)
+        
+        return samples
+    
+    def process_tar_file(self, rgb_caption_tar):
+        """Process a single tar file and create corresponding parquet"""
+        
+        # Extract base name (e.g., 'sa_000000' from 'sa_000000.tar')
+        base_name = os.path.splitext(os.path.basename(rgb_caption_tar))[0]
+        
+        # Construct paths for corresponding depth, normal, det_seg, grounding, canny, dino tar files
+        depth_tar = os.path.join(self.depth_dir, f"{base_name}.tar")
+        normal_tar = os.path.join(self.normal_dir, f"{base_name}.tar")
+        det_seg_tar = os.path.join(self.det_seg_dir, f"{base_name}.tar")
+        grounding_tar = os.path.join(self.grounding_dir, f"{base_name}.tar")
+        canny_tar = os.path.join(self.canny_dir, f"{base_name}.tar")
+        dino_tar = os.path.join(self.dino_dir, f"{base_name}.tar")
+        
+        logger.info(f"Processing {base_name}")
+        
+        # Check if corresponding files exist
+        if not os.path.exists(depth_tar):
+            logger.warning(f"Depth tar file not found: {depth_tar}")
+            return None
+        if not os.path.exists(normal_tar):
+            logger.warning(f"Normal tar file not found: {normal_tar}")
+            return None
+        if not os.path.exists(det_seg_tar):
+            logger.warning(f"Det seg tar file not found: {det_seg_tar}")
+            return None
+        if not os.path.exists(canny_tar):
+            logger.warning(f"Canny tar file not found: {canny_tar}")
+            return None
+        if not os.path.exists(dino_tar):
+            logger.warning(f"DINOv2 tar file not found: {dino_tar}")
+            return None
+        # Initialize grounding_data
+        grounding_data = {}
+        
+        if not os.path.exists(grounding_tar):
+            logger.warning(f"Grounding tar file not found: {grounding_tar} - will process without grounding data")
+        else:
+            logger.info(f"Extracting from grounding: {grounding_tar}")
+            grounding_data = self.extract_from_tar(grounding_tar, 'grounding')
+        
+        # Extract data from all tar files
+        logger.info(f"Extracting from rgb_caption: {rgb_caption_tar}")
+        rgb_caption_data = self.extract_from_tar(rgb_caption_tar, 'rgb_caption')
+        
+        logger.info(f"Extracting from depth: {depth_tar}")
+        depth_data = self.extract_from_tar(depth_tar, 'depth')
+        
+        logger.info(f"Extracting from normal: {normal_tar}")
+        normal_data = self.extract_from_tar(normal_tar, 'normal')
+        
+        logger.info(f"Extracting from det_seg: {det_seg_tar}")
+        det_seg_data = self.extract_from_tar(det_seg_tar, 'det_seg')
+        
+        logger.info(f"Extracting from canny: {canny_tar}")
+        canny_data = self.extract_from_tar(canny_tar, 'canny')
+        
+        logger.info(f"Extracting from dino: {dino_tar}")
+        dino_data = self.extract_from_tar(dino_tar, 'dino')
+        
+        # Find common UIDs (grounding is optional)
+        common_uids = set(rgb_caption_data.keys()) & set(depth_data.keys()) & set(normal_data.keys()) & set(det_seg_data.keys()) & set(canny_data.keys()) & set(dino_data.keys())
+        
+        if not common_uids:
+            logger.warning(f"No common UIDs found for {base_name}")
+            return None
+        
+        logger.info(f"Found {len(common_uids)} common UIDs")
+        
+        # Sort UIDs for consistent ordering
+        sorted_uids = sorted(common_uids)
+        
+        # Create samples for each UID
+        all_samples = []
+        
+        for uid in sorted_uids:
+            rgb_info = rgb_caption_data[uid]
+            
+            # Skip if missing image or caption
+            if rgb_info['image'] is None or rgb_info['caption'] is None:
+                continue
+            
+            # Create base sample with all modalities
+            sample = {
+                'uid': uid,
+                'rgb': rgb_caption_data[uid]['image'],
+                'caption': rgb_caption_data[uid]['caption'],
+                'depth': depth_data[uid],
+                'normal': normal_data[uid],
+                'det_seg': det_seg_data[uid],
+                'canny': canny_data[uid],
+                'dino': dino_data[uid]
+            }
+            
+            # Add grounding data if available
+            if uid in grounding_data:
+                extracted_grounding = self.extract_grounding_data(grounding_data[uid])
+                if extracted_grounding:
+                    raw_image = Image.open(io.BytesIO(rgb_caption_data[uid]['image']))
+                    raw_image_size = raw_image.size
+                    assert raw_image_size == (512, 512), f"Raw image size is {raw_image_size} but expected (512, 512) for {uid}"
+                    grounding_samples = self.create_grounding_samples(extracted_grounding, uid)
+                    if grounding_samples:
+                        sample['grounding'] = grounding_samples
+            
+            # visualize_segmentation = sample['grounding'][1]['segmentation'] 
+            # self.visualize_rle_mask(visualize_segmentation,save_path=f"visualize_{uid}_1.png")
+            all_samples.append(sample)
+        
+        if not all_samples:
+            logger.warning(f"No samples found for {base_name}")
+            return None
+        
+        # Create DataFrame
+        df = pd.DataFrame(all_samples)
+        
+        # Debug information
+        logger.info(f"DataFrame shape: {df.shape}")
+        if not df.empty:
+            logger.info(f"Sample UIDs: {df['uid'].head().tolist()}")
+            # Count grounding samples
+            grounding_counts = []
+            for _, row in df.iterrows():
+                if 'grounding' in row and row['grounding'] is not None and isinstance(row['grounding'], list):
+                    grounding_counts.append(len(row['grounding']))
+                else:
+                    grounding_counts.append(0)
+            if grounding_counts:
+                logger.info(f"Grounding samples per row: {np.mean(grounding_counts):.2f} avg, {max(grounding_counts)} max")
+            else:
+                logger.info("No grounding samples found")
+            
+            # Count DINOv2 features
+            dino_counts = []
+            for _, row in df.iterrows():
+                if 'dino' in row and row['dino'] is not None:
+                    dino_counts.append(1)
+                else:
+                    dino_counts.append(0)
+            if dino_counts:
+                logger.info(f"DINOv2 features loaded: {sum(dino_counts)}/{len(dino_counts)} samples")
+            else:
+                logger.info("No DINOv2 features found")
+        
+        # Save to parquet
+        output_path = os.path.join(self.parquet_dir, f"{base_name}.parquet")
+        
+        df.to_parquet(
+            output_path, 
+            index=False,
+            engine='pyarrow',
+            compression='snappy',
+            row_group_size=1000
+        )
+        logger.info(f"Saved {len(df)} grounding samples to {output_path}")
+        
+        return output_path
+    
+    def process_all_tars(self):
+        """Process all tar files and create corresponding parquet files"""
+        tar_files = self.get_tar_files()
+        
+        if not tar_files:
+            logger.error("No tar files found in rgb_caption directory")
+            return
+        
+        tar_files = tar_files[1373:]
+        logger.info(f"Found {len(tar_files)} tar files to process")
+        
+        processed_files = []
+        for tar_file in tar_files:
+            try:
+                output_path = self.process_tar_file(tar_file)
+                if output_path:
+                    processed_files.append(output_path)
+            except Exception as e:
+                logger.error(f"Error processing {tar_file}: {e}")
+        
+        logger.info(f"Successfully processed {len(processed_files)} files")
+        return processed_files
+
+def main():
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Generate parquet files with grounding data')
+    parser.add_argument('--base_dir', type=str, default='./datasets/blip3o',
+                       help='Base directory (default: ./datasets/blip3o)')
+    parser.add_argument('--visualize_mask', action='store_true',
+                       help='Visualize RLE masks for debugging')
+    parser.add_argument('--visualize_tar', type=str, default=None,
+                       help='Specific tar file to visualize (e.g., sa_000001.tar)')
+    parser.add_argument('--visualize_uid', type=str, default=None,
+                       help='Specific UID to visualize (e.g., sa_1)')
+    parser.add_argument('--visualize_object_id', type=int, default=0,
+                       help='Object ID to visualize (default: 0)')
+    
+    args = parser.parse_args()
+    
+    # Set the base directory
+    base_dir = args.base_dir
+    
+    # Create processor
+    processor = GroundingProcessor(base_dir)
+    
+    # Visualize mask if requested
+    if args.visualize_mask and args.visualize_tar:
+        logger.info(f"Visualizing mask for {args.visualize_tar}, object {args.visualize_object_id}")
+        
+        # Find the grounding tar file for this UID
+        grounding_tar = os.path.join(processor.grounding_dir, f"{args.visualize_tar}.tar")
+        if os.path.exists(grounding_tar):
+            grounding_data = processor.extract_from_tar(grounding_tar, 'grounding')
+            if args.visualize_uid in grounding_data:
+                extracted_grounding = processor.extract_grounding_data(grounding_data[args.visualize_uid])
+                if extracted_grounding:
+                    save_path = f"mask_visualization_{args.visualize_uid}.png"
+                    processor.visualize_object_mask(extracted_grounding, args.visualize_uid, 
+                                                  args.visualize_object_id, save_path)
+                else:
+                    logger.error(f"No grounding data found for {args.visualize_uid}")
+            else:
+                logger.error(f"UID {args.visualize_uid} not found in grounding tar")
+        else:
+            logger.error(f"Grounding tar not found: {grounding_tar}")
+        return
+    
+    # Process all tar files
+    processed_files = processor.process_all_tars()
+    
+    if processed_files:
+        logger.info(f"Successfully created {len(processed_files)} grounding parquet files:")
+        for file in processed_files:
+            logger.info(f"  - {file}")
+    else:
+        logger.error("No grounding parquet files were created")
+
+if __name__ == "__main__":
+    main()

data/any2any_preprocess/generate_parquet_json.py

@@ -0,0 +1,69 @@
+import os
+import json
+import argparse
+import pyarrow.parquet as pq
+from pathlib import Path
+
+def get_parquet_info(file_path):
+    """Get information about a parquet file"""
+    try:
+        parquet_file = pq.ParquetFile(file_path)
+        num_row_groups = parquet_file.num_row_groups
+        
+        # Calculate total number of rows
+        total_rows = 0
+        for i in range(num_row_groups):
+            total_rows += parquet_file.metadata.row_group(i).num_rows
+        
+        return {
+            "num_row_groups": num_row_groups,
+            "num_rows": total_rows
+        }
+    except Exception as e:
+        print(f"Error reading {file_path}: {e}")
+        return None
+
+def generate_parquet_info(parquet_dir=None, output_file=None):
+    """Generate information for all parquet files in the parquet folder"""
+    parquet_dir = Path(parquet_dir) if parquet_dir is not None else Path("./datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global_clip448_imagebind")
+    output_file = Path(output_file) if output_file is not None else Path("./datasets/blip3o/parquet_info/blip3o_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global_clip448_imagebind.json")
+
+    # Create output directory if it doesn't exist
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    
+    # Get all parquet files
+    parquet_files = list(parquet_dir.glob("*.parquet"))
+    print(f"Found {len(parquet_files)} parquet files")
+    
+    # Process each file
+    parquet_info = {}
+    
+    for i, parquet_file in enumerate(parquet_files):
+        print(f"Processing {i+1}/{len(parquet_files)}: {parquet_file.name}")
+        
+        # Get relative path from current directory using parquet_dir
+        relative_path = str((parquet_dir / parquet_file.name))
+        if not relative_path.startswith("./"):
+            relative_path = f"./{relative_path}"
+        
+        info = get_parquet_info(parquet_file)
+        if info:
+            parquet_info[relative_path] = info
+            print(f"Processed {parquet_file.name} successfully")
+            print(info)
+        else:
+            print(f"Failed to process {parquet_file.name}")
+    
+    # Save to JSON file
+    with open(output_file, 'w') as f:
+        json.dump(parquet_info, f, indent=2, sort_keys=True)
+    
+    print(f"Saved parquet information to {output_file}")
+    print(f"Processed {len(parquet_info)} files successfully")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--parquet_dir", default=None)
+    parser.add_argument("--output_file", default=None)
+    args = parser.parse_args()
+    generate_parquet_info(parquet_dir=args.parquet_dir, output_file=args.output_file)

data/any2any_preprocess/generate_parquet_vlm_sft.py

@@ -0,0 +1,278 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""Bundle vlm_sft (llava-onevision style jsonl + loose image files) into
+parquet shards loadable by SftParquetIterableDataset (data/vlm_dataset.py).
+
+Why: llava-onevision is ~4M small files, impractical to upload to / stream
+from S3. The shards follow the blip3o convention: few large parquet files +
+a parquet_info JSON, sharded by (file, row_group) at load time.
+
+Equivalence: the jsonl loader selects samples per subset via
+    rng = random.Random(); rng.seed(shuffle_seed); rng.shuffle(lines)
+    lines = lines[:num_used_data]
+(SftJSONLIterableDataset.get_data_paths). This script replicates that
+selection EXACTLY (taking shuffle_lines / shuffle_seed / num_used_data from
+the same training yaml), so the parquet rows are precisely the records the
+jsonl loader would have used, in the same shuffled order.
+
+Row schema:
+    record: str                original jsonl line, verbatim
+    image_bytes: list<binary>  raw image files (no re-encode), in the order
+                               the jsonl loader's valid_paths would produce
+
+Skipped at conversion (mirroring the jsonl loader's runtime skips):
+  - records with 'video' (frame sampling can't be baked into bytes)
+  - records whose 'image' field resolves to zero existing files
+
+Usage (single process, all subsets of the group):
+  python3 data/any2any_preprocess/generate_parquet_vlm_sft.py \
+      --config data/configs/modus_stage1.yaml --group vlm_sft \
+      --out-dir ./datasets/llava_onevision_vqa_parquet
+
+Parallelize across subsets with disjoint --subsets lists (e.g. sbatch array);
+each invocation only rewrites its own subsets' shards and merge-updates the
+shared parquet_info JSON. Already-converted subsets are skipped unless
+--overwrite is given.
+
+After conversion, point the training config at the new group:
+  vlm_sft_parquet:
+    dataset_names:
+    - llava_onevision_vqa
+    image_transform_args: { ...same as the old vlm_sft block... }
+    is_mandatory: true
+    weight: <same as before>
+  (no num_used_data / shuffle_lines needed — selection is baked in)
+"""
+
+import argparse
+import json
+import os
+import random
+import re
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import yaml
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
+from data.dataset_info import DATASET_INFO  # noqa: E402
+from data.parquet_utils import apply_data_root_override, read_file_bytes  # noqa: E402
+
+SCHEMA = pa.schema([
+    pa.field('record', pa.string()),
+    pa.field('image_bytes', pa.list_(pa.binary())),
+])
+
+
+def sanitize(name):
+    # subset names contain '(', ',', ')' etc. — make a filesystem-safe stem
+    return re.sub(r'[^0-9A-Za-z_.-]+', '_', name).strip('_')
+
+
+def _read_image_bytes(path):
+    if path.startswith('s3://'):
+        return read_file_bytes(path)
+    with open(path, 'rb') as f:
+        return f.read()
+
+
+def _load_record_images(args):
+    """Worker: (raw_line, image_dir) -> (raw_line, image_bytes list, status).
+    status: 'ok' | 'skip_video' | 'skip_missing' | 'skip_badjson'."""
+    raw_line, image_dir = args
+    try:
+        data_item = json.loads(raw_line)
+    except Exception:
+        return raw_line, None, 'skip_badjson'
+    if 'video' in data_item:
+        return raw_line, None, 'skip_video'
+    image_bytes = []
+    if 'image' in data_item:
+        images = data_item['image'] if isinstance(data_item['image'], list) else [data_item['image']]
+        for img_name in images:
+            if not (isinstance(img_name, str) and img_name.strip() != ''):
+                continue
+            img_path = os.path.join(image_dir, img_name)
+            try:
+                image_bytes.append(_read_image_bytes(img_path))
+            except (OSError, IOError):
+                continue  # mirrors the jsonl loader dropping invalid paths
+        if len(image_bytes) == 0:
+            return raw_line, None, 'skip_missing'
+    return raw_line, image_bytes, 'ok'
+
+
+def select_lines(jsonl_path, num_used_data, shuffle_lines, shuffle_seed):
+    """Replicates SftJSONLIterableDataset.get_data_paths line selection."""
+    jsonl_path = apply_data_root_override(jsonl_path)
+    if jsonl_path.startswith('s3://'):
+        raw_data = read_file_bytes(jsonl_path).decode('utf-8').splitlines(keepends=True)
+    else:
+        with open(jsonl_path, 'r') as f:
+            raw_data = f.readlines()
+    if shuffle_lines:
+        rng = random.Random()
+        rng.seed(shuffle_seed)
+        rng.shuffle(raw_data)
+    return raw_data[:num_used_data]
+
+
+def convert_subset(subset, jsonl_path, image_dir, num_used_data, shuffle_lines,
+                   shuffle_seed, train_dir, rows_per_file, row_group_size,
+                   io_threads, info_key_dir):
+    """Convert one subset to parquet shards. Returns (shard_info, stats):
+    shard_info: {info_key: {'num_row_groups': int, 'num_rows': int}}."""
+    image_dir = apply_data_root_override(image_dir)
+    lines = select_lines(jsonl_path, num_used_data, shuffle_lines, shuffle_seed)
+    print(f'[convert] {subset}: selected {len(lines)} lines (quota {num_used_data})')
+
+    stem = sanitize(subset)
+    stats = {'ok': 0, 'skip_video': 0, 'skip_missing': 0, 'skip_badjson': 0}
+    shard_info = {}
+    part = 0
+    records, images = [], []
+
+    def flush():
+        nonlocal part, records, images
+        if not records:
+            return
+        fname = f'{stem}__part{part:04d}.parquet'
+        out_path = os.path.join(train_dir, fname)
+        table = pa.Table.from_arrays(
+            [pa.array(records, type=pa.string()),
+             pa.array(images, type=pa.list_(pa.binary()))],
+            schema=SCHEMA,
+        )
+        pq.write_table(table, out_path, compression='snappy',
+                       row_group_size=row_group_size)
+        meta = pq.ParquetFile(out_path).metadata
+        shard_info[os.path.join(info_key_dir, fname)] = {
+            'num_row_groups': meta.num_row_groups,
+            'num_rows': meta.num_rows,
+        }
+        part += 1
+        records, images = [], []
+
+    with ThreadPoolExecutor(max_workers=io_threads) as ex:
+        # batch the submissions: Executor.map submits everything up front and
+        # holds completed results until iterated, so mapping the whole subset
+        # at once would keep every image of the subset in memory.
+        for start in range(0, len(lines), rows_per_file):
+            chunk = lines[start:start + rows_per_file]
+            for raw_line, image_bytes, status in ex.map(
+                _load_record_images, ((l, image_dir) for l in chunk)
+            ):
+                stats[status] += 1
+                if status != 'ok':
+                    continue
+                records.append(raw_line)
+                images.append(image_bytes)
+                if len(records) >= rows_per_file:
+                    flush()
+    flush()
+    return shard_info, stats
+
+
+def main():
+    # DATASET_INFO uses './datasets/...' paths relative to the repo root;
+    # run from there regardless of the caller's cwd (e.g. container default
+    # workdir). Pass absolute --config/--out-dir or repo-root-relative ones.
+    os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--config', required=True,
+                        help='training data yaml (e.g. data/configs/modus_stage1.yaml)')
+    parser.add_argument('--group', default='vlm_sft',
+                        help='dataset group in the yaml / DATASET_INFO')
+    parser.add_argument('--out-dir', default='./datasets/llava_onevision_vqa_parquet')
+    parser.add_argument('--subsets', nargs='*', default=None,
+                        help='only convert these subsets (default: all in the yaml)')
+    parser.add_argument('--rows-per-file', type=int, default=5000)
+    parser.add_argument('--row-group-size', type=int, default=200)
+    parser.add_argument('--io-threads', type=int, default=16)
+    parser.add_argument('--overwrite', action='store_true',
+                        help='re-convert subsets whose shards already exist')
+    args = parser.parse_args()
+
+    with open(args.config, 'r') as f:
+        cfg = yaml.safe_load(f)[args.group]
+    dataset_names = cfg['dataset_names']
+    num_used_data = cfg['num_used_data']
+    shuffle_lines = cfg.get('shuffle_lines', False)
+    shuffle_seed = cfg.get('shuffle_seed', 0)
+    if len(dataset_names) != len(num_used_data):
+        # The jsonl loader zips the two lists, silently truncating the longer
+        # one (this happens in the real configs: subsets commented out of
+        # dataset_names leave stale num_used_data rows). Replicate that
+        # pairing exactly — equivalence with current training behavior.
+        print(f'[warn] dataset_names ({len(dataset_names)}) and num_used_data '
+              f'({len(num_used_data)}) differ; zip-truncating like the jsonl loader')
+
+    stems = [sanitize(s) for s in dataset_names]
+    assert len(set(stems)) == len(stems), 'sanitized subset names collide'
+
+    train_dir = os.path.join(args.out_dir, 'train')
+    info_dir = os.path.join(args.out_dir, 'parquet_info')
+    os.makedirs(train_dir, exist_ok=True)
+    os.makedirs(info_dir, exist_ok=True)
+    info_path = os.path.join(info_dir, 'llava_onevision_vqa.json')
+    # info keys use the configured-style relative path; loaders fall back to
+    # basename matching anyway, so the prefix only matters cosmetically.
+    info_key_dir = train_dir
+
+    selected = args.subsets if args.subsets else dataset_names
+    unknown = [s for s in selected if s not in dataset_names]
+    assert not unknown, f'subsets not in the yaml group: {unknown}'
+
+    totals = {'ok': 0, 'skip_video': 0, 'skip_missing': 0, 'skip_badjson': 0}
+    for subset, n_used in zip(dataset_names, num_used_data):
+        if subset not in selected:
+            continue
+        stem = sanitize(subset)
+        # A subset counts as done only if its shards are in the info JSON,
+        # which is merge-updated AFTER the whole subset finishes. Files on
+        # disk without an info entry are partial output of an interrupted
+        # run — delete and re-convert.
+        done_in_info = set()
+        if os.path.exists(info_path) and not args.overwrite:
+            with open(info_path, 'r') as f:
+                done_in_info = {os.path.basename(k) for k in json.load(f)
+                                if os.path.basename(k).startswith(f'{stem}__part')}
+        existing = [f for f in os.listdir(train_dir)
+                    if f.startswith(f'{stem}__part') and f.endswith('.parquet')]
+        if done_in_info and set(existing) >= done_in_info:
+            print(f'[skip] {subset}: {len(done_in_info)} shard(s) already converted')
+            continue
+        for f in existing:
+            os.remove(os.path.join(train_dir, f))
+
+        meta_info = DATASET_INFO[args.group][subset]
+        shard_info, stats = convert_subset(
+            subset, meta_info['jsonl_path'], meta_info['data_dir'], n_used,
+            shuffle_lines, shuffle_seed, train_dir, args.rows_per_file,
+            args.row_group_size, args.io_threads, info_key_dir,
+        )
+        for k in totals:
+            totals[k] += stats[k]
+        print(f'[done] {subset}: {stats}')
+
+        # merge-update the shared parquet_info JSON after each subset so
+        # partial runs / parallel invocations stay resumable.
+        merged = {}
+        if os.path.exists(info_path):
+            with open(info_path, 'r') as f:
+                merged = json.load(f)
+        merged.update(shard_info)
+        tmp_path = info_path + f'.tmp.{os.getpid()}'
+        with open(tmp_path, 'w') as f:
+            json.dump(merged, f, indent=1)
+        os.replace(tmp_path, info_path)
+
+    print(f'[all done] totals: {totals}')
+    print(f'parquet_info: {info_path}')
+
+
+if __name__ == '__main__':
+    main()

data/any2any_preprocess/hf_upload.py

@@ -0,0 +1,26 @@
+"""Upload a folder to a HuggingFace dataset repo.
+
+Token is read from the file at $HF_TOKEN_FILE (never passed on the command line
+so it does not appear in logs / process listings).
+
+Usage:
+    HF_TOKEN_FILE=/users/mye/.hf_token python hf_upload.py <repo_id> <folder>
+"""
+import os
+import sys
+
+from huggingface_hub import HfApi
+
+repo_id = sys.argv[1]
+folder = sys.argv[2]
+token = open(os.environ['HF_TOKEN_FILE']).read().strip()
+
+api = HfApi(token=token)
+api.create_repo(repo_id, repo_type='dataset', exist_ok=True)
+api.upload_folder(
+    repo_id=repo_id,
+    repo_type='dataset',
+    folder_path=folder,
+    commit_message='Add MODUS 16-mod preview (500 SA-1B samples + hero + card)',
+)
+print(f'DONE: https://huggingface.co/datasets/{repo_id}', flush=True)

data/any2any_preprocess/parquet_visualize.py

@@ -0,0 +1,408 @@
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+import os
+import numpy as np
+from pathlib import Path
+from PIL import Image, ImageDraw, ImageFont
+import json
+import io
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.patches import Rectangle
+import cv2
+
+def make_json_serializable(obj):
+    """
+    Convert numpy arrays and other non-serializable objects to JSON-serializable format.
+    """
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    elif isinstance(obj, dict):
+        return {key: make_json_serializable(value) for key, value in obj.items()}
+    elif isinstance(obj, list):
+        return [make_json_serializable(item) for item in obj]
+    elif isinstance(obj, (np.integer, np.floating)):
+        return obj.item()
+    elif hasattr(obj, 'tolist'):  # Handle other numpy-like objects
+        return obj.tolist()
+    elif hasattr(obj, 'item'):  # Handle other numpy scalar types
+        return obj.item()
+    else:
+        return obj
+
+def decode_rle_to_mask(rle_counts, size):
+    """
+    Decode RLE (Run Length Encoding) to binary mask.
+    This is a simplified RLE decoder for COCO format.
+    
+    Args:
+        rle_counts (str): RLE encoded string
+        size (list): [width, height] of the mask
+    
+    Returns:
+        numpy.ndarray: Binary mask of shape (height, width)
+    """
+    try:
+        # Try to use pycocotools if available
+        import pycocotools.mask as mask_util
+        
+        # Create RLE dict
+        rle = {
+            'counts': rle_counts,
+            'size': size
+        }
+        
+        # Decode to binary mask
+        binary_mask = mask_util.decode(rle)
+        return binary_mask
+    except ImportError:
+        # Fallback: create a simple mask visualization without decoding
+        print("Warning: pycocotools not available, creating placeholder mask")
+        height, width = size[1], size[0]  # size is [width, height]
+        return np.zeros((height, width), dtype=np.uint8)
+
+def simple_rle_decode(rle_counts, size):
+    """
+    Simple RLE decoder for COCO format masks.
+    This is a basic implementation that should work for most cases.
+    """
+    try:
+        height, width = size[1], size[0]  # size is [width, height]
+        mask = np.zeros(height * width, dtype=np.uint8)
+        
+        # Parse the RLE counts
+        counts = []
+        current_num = ""
+        for char in rle_counts:
+            if char.isdigit() or char in '-+':
+                current_num += char
+            else:
+                if current_num:
+                    counts.append(int(current_num))
+                    current_num = ""
+        if current_num:
+            counts.append(int(current_num))
+        
+        # Decode RLE
+        pos = 0
+        for i, count in enumerate(counts):
+            if i % 2 == 0:  # Even indices are background (0)
+                pos += count
+            else:  # Odd indices are foreground (1)
+                mask[pos:pos+count] = 1
+                pos += count
+        
+        # Reshape to 2D
+        mask = mask.reshape((height, width))
+        return mask
+    except Exception as e:
+        print(f"Error in simple RLE decode: {e}")
+        height, width = size[1], size[0]
+        return np.zeros((height, width), dtype=np.uint8)
+
+def visualize_detection_data(det_data, rgb_image, output_dir, filename_base):
+    """
+    Visualize detection data with bounding boxes and masks overlaid on RGB image.
+    
+    Args:
+        det_data (dict): Detection data containing instances
+        rgb_image (PIL.Image): RGB image to overlay on
+        output_dir (str): Output directory
+        filename_base (str): Base filename for saving
+    """
+    try:
+        # Make a deep copy of the original data to avoid modifying it
+        import copy
+        original_det_data = make_json_serializable(copy.deepcopy(det_data))
+        # Convert PIL image to numpy array
+        img_array = np.array(rgb_image)
+        height, width = img_array.shape[:2]
+        
+        # Create figure with subplots
+        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+        
+        # Original image
+        axes[0].imshow(img_array)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+        
+        # Image with bounding boxes
+        axes[1].imshow(img_array)
+        axes[1].set_title('Bounding Boxes')
+        axes[1].axis('off')
+        
+        # Image with masks
+        axes[2].imshow(img_array)
+        axes[2].set_title('Segmentation Masks')
+        axes[2].axis('off')
+        
+        # Colors for different categories
+        colors = ['red', 'blue', 'green', 'yellow', 'purple', 'orange', 'cyan', 'magenta']
+        
+        # Process each instance
+        for i, instance in enumerate(det_data.get('instances', [])):
+            color = colors[i % len(colors)]
+            category = instance.get('category', 'unknown')
+            score = instance.get('score', 0.0)
+            bbox = instance.get('bbox', [])
+            
+            # Draw bounding box
+            if len(bbox) == 4:
+                x1, y1, x2, y2 = bbox
+                rect = Rectangle((x1, y1), x2-x1, y2-y1, 
+                               linewidth=2, edgecolor=color, facecolor='none')
+                axes[1].add_patch(rect)
+                
+                # Add label
+                label = f"{category}: {score:.2f}"
+                axes[1].text(x1, y1-5, label, color=color, fontsize=8, 
+                           bbox=dict(boxstyle="round,pad=0.3", facecolor='white', alpha=0.7))
+            
+            # Draw segmentation mask
+            segmentation = instance.get('segmentation', {})
+            mask_drawn = False
+            
+            if 'counts' in segmentation and 'size' in segmentation:
+                # Try multiple decoding methods
+                mask = None
+                
+                # Method 1: Try pycocotools
+                try:
+                    mask = decode_rle_to_mask(segmentation['counts'], segmentation['size'])
+                    if mask.size > 0 and np.any(mask):
+                        print(f"Successfully decoded mask for {category} using pycocotools")
+                except:
+                    pass
+                
+                # Method 2: Try simple RLE decoder
+                if mask is None or not np.any(mask):
+                    try:
+                        mask = simple_rle_decode(segmentation['counts'], segmentation['size'])
+                        if mask.size > 0 and np.any(mask):
+                            print(f"Successfully decoded mask for {category} using simple decoder")
+                    except Exception as e:
+                        print(f"Simple RLE decode failed for {category}: {e}")
+                
+                # If we have a valid mask, visualize it
+                if mask is not None and mask.size > 0 and np.any(mask):
+                    # Create a more visible mask visualization
+                    # Method 1: Colored overlay
+                    mask_colored = np.zeros((height, width, 4))
+                    color_rgb = np.array(plt.cm.colors.to_rgb(color)) * 255
+                    mask_colored[:, :, :3] = color_rgb
+                    mask_colored[:, :, 3] = mask * 0.6  # More opaque
+                    
+                    # Apply mask overlay
+                    axes[2].imshow(mask_colored.astype(np.uint8), alpha=0.7)
+                    
+                    # Method 2: Also draw contour
+                    from matplotlib.patches import Polygon
+                    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                    for contour in contours:
+                        if len(contour) > 2:
+                            contour = contour.reshape(-1, 2)
+                            polygon = Polygon(contour, fill=False, edgecolor=color, linewidth=2)
+                            axes[2].add_patch(polygon)
+                    
+                    mask_drawn = True
+            
+            # Fallback: if no mask was drawn, create a simple mask from bounding box
+            if not mask_drawn and len(bbox) == 4:
+                x1, y1, x2, y2 = bbox
+                # Create a simple rectangular mask
+                mask = np.zeros((height, width), dtype=np.uint8)
+                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+                # Ensure coordinates are within bounds
+                x1, y1 = max(0, x1), max(0, y1)
+                x2, y2 = min(width, x2), min(height, y2)
+                
+                if x2 > x1 and y2 > y1:
+                    mask[y1:y2, x1:x2] = 1
+                    
+                    # Create colored mask overlay
+                    mask_colored = np.zeros((height, width, 4))
+                    color_rgb = np.array(plt.cm.colors.to_rgb(color)) * 255
+                    mask_colored[:, :, :3] = color_rgb
+                    mask_colored[:, :, 3] = mask * 0.4  # Semi-transparent
+                    
+                    # Apply mask overlay
+                    axes[2].imshow(mask_colored.astype(np.uint8), alpha=0.6)
+                    
+                    # Also draw the bounding box outline
+                    rect = Rectangle((x1, y1), x2-x1, y2-y1, 
+                                   linewidth=2, edgecolor=color, facecolor='none')
+                    axes[2].add_patch(rect)
+        
+        # Save the visualization
+        output_path = os.path.join(output_dir, f"{filename_base}_detection_visualization.png")
+        plt.tight_layout()
+        plt.savefig(output_path, dpi=150, bbox_inches='tight')
+        plt.close()
+        
+        print(f"  Saved detection visualization: {filename_base}_detection_visualization.png")
+        
+        # Also save the JSON data for reference
+        json_path = os.path.join(output_dir, f"{filename_base}_det_seg.json")
+        try:
+            with open(json_path, 'w') as f:
+                json.dump(original_det_data, f, indent=2)
+            print(f"  Saved detection JSON: {filename_base}_det_seg.json")
+        except Exception as json_error:
+            print(f"  JSON save error: {json_error}")
+            # Try to save a simplified version
+            try:
+                simplified_data = {
+                    'image_id': original_det_data.get('image_id', 'unknown'),
+                    'instances': []
+                }
+                for instance in original_det_data.get('instances', []):
+                    simplified_instance = {
+                        'category': instance.get('category', 'unknown'),
+                        'score': float(instance.get('score', 0.0)),
+                        'bbox': [float(x) for x in instance.get('bbox', [])],
+                        'area': float(instance.get('area', 0.0))
+                    }
+                    simplified_data['instances'].append(simplified_instance)
+                
+                with open(json_path, 'w') as f:
+                    json.dump(simplified_data, f, indent=2)
+                print(f"  Saved simplified detection JSON: {filename_base}_det_seg.json")
+            except Exception as e2:
+                print(f"  Failed to save even simplified JSON: {e2}")
+        
+    except Exception as e:
+        print(f"  Error visualizing detection data: {e}")
+
+def visualize_parquet_examples(input_file, output_dir, max_rows=None):
+    """
+    Visualize examples from a parquet file by saving each row as separate files.
+    Saves images as PNG and text data as TXT files with uid_column_name naming format.
+    
+    Args:
+        input_file (str): Path to the input parquet file
+        output_dir (str): Directory to save the visualization files
+        max_rows (int): Maximum number of rows to process (None for all rows)
+    """
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+    
+    print(f"Processing {input_file}...")
+    
+    # Read the parquet file
+    df = pd.read_parquet(input_file)
+    print(f"Total rows in file: {len(df)}")
+    print(f"Columns: {list(df.columns)}")
+    
+    # Limit rows if max_rows is specified
+    if max_rows is not None:
+        df = df.head(max_rows)
+        print(f"Processing only first {max_rows} rows")
+    
+    # Display first few rows to understand the data structure
+    print("\nFirst few rows:")
+    print(df.head(3))
+    
+    # Process each row
+    for index, row in df.iterrows():
+        print(f"Processing row {index}...")
+        
+        filename_base = None
+        
+        # Process each column in the row
+        for col_name, value in row.items():
+            if col_name == 'uid':
+                filename_base = value
+            elif col_name in ['rgb', 'depth', 'normal'] and filename_base is not None:
+                # Handle image data (binary)
+                if value is not None and isinstance(value, bytes):
+                    try:
+                        # Convert binary data to PIL Image
+                        image = Image.open(io.BytesIO(value))
+                        
+                        # Create filename
+                        filename = f"{filename_base}_{col_name}.png"
+                        filepath = os.path.join(output_dir, filename)
+                        
+                        # Save as PNG
+                        image.save(filepath, 'PNG')
+                        print(f"  Saved {col_name} image: {filename}")
+                        
+                    except Exception as e:
+                        print(f"  Error processing {col_name} image: {e}")
+                        
+            elif col_name == 'caption' and filename_base is not None:
+                # Handle text data
+                if value is not None:
+                    try:
+                        filename = f"{filename_base}_captions.txt"
+                        filepath = os.path.join(output_dir, filename)
+                        
+                        # Save text data
+                        with open(filepath, 'w', encoding='utf-8') as f:
+                            if isinstance(value, str):
+                                f.write(value)
+                            else:
+                                f.write(str(value))
+                        print(f"  Saved caption: {filename}")
+                        
+                    except Exception as e:
+                        print(f"  Error processing caption: {e}")
+            elif col_name == 'det_seg' and filename_base is not None:
+                # Handle detection maps (JSON format)
+                if value is not None:
+                    try:
+                        # Get the RGB image for visualization
+                        rgb_image = None
+                        if 'rgb' in row and row['rgb'] is not None:
+                            rgb_image = Image.open(io.BytesIO(row['rgb']))
+                        
+                        # Visualize detection data with bounding boxes and masks
+                        if rgb_image is not None:
+                            visualize_detection_data(value, rgb_image, output_dir, filename_base)
+                        else:
+                            # If no RGB image, just save the JSON
+                            filename = f"{filename_base}_det_seg.json"
+                            filepath = os.path.join(output_dir, filename)
+                            with open(filepath, 'w') as f:
+                                serializable_data = make_json_serializable(value)
+                                json.dump(serializable_data, f, indent=2)
+                            print(f"  Saved detection JSON: {filename}")
+                        
+                    except Exception as e:
+                        print(f"  Error processing detection data: {e}")
+                        
+            
+    
+    print(f"\nVisualization completed! Files saved to: {output_dir}")
+    
+    # List the created files
+    created_files = os.listdir(output_dir)
+    print(f"Created {len(created_files)} files:")
+    for file in sorted(created_files)[:10]:  # Show first 10 files
+        print(f"  {file}")
+    if len(created_files) > 10:
+        print(f"  ... and {len(created_files) - 10} more files")
+
+def main():
+    # Configuration
+    input_file = "./datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg/sa_000002.parquet"
+    output_dir = "./datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg_visualize"
+    max_rows = 30  # Process only first 30 rows
+    
+    # Check if input file exists
+    if not os.path.exists(input_file):
+        print(f"Error: Input file {input_file} not found!")
+        return
+    
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Visualize the parquet examples
+    visualize_parquet_examples(input_file, output_dir, max_rows)
+    
+    print("\nVisualization completed successfully!")
+
+if __name__ == "__main__":
+    main()

data/any2any_preprocess/process_grounding.py

@@ -0,0 +1,346 @@
+import os
+import tarfile
+import json
+import glob
+import io
+from tqdm import tqdm
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class GroundingProcessor:
+    def __init__(self, base_dir):
+        self.base_dir = base_dir
+        self.rgb_caption_dir = os.path.join(base_dir, 'datasets/blip3o/datasets')
+        self.grounding_source_dir = os.path.join(base_dir, 'datasets/glamm')
+        self.grounding_output_dir = os.path.join(base_dir, 'datasets/blip3o/datasets_grounding')
+        self.grounding_cache_dir = os.path.join(base_dir, 'datasets/blip3o/grounding_cache')
+        
+        # Create output directories if they don't exist
+        os.makedirs(self.grounding_output_dir, exist_ok=True)
+        os.makedirs(self.grounding_cache_dir, exist_ok=True)
+        
+    def get_rgb_tar_files(self):
+        """Get all RGB tar files from datasets directory"""
+        tar_files = glob.glob(os.path.join(self.rgb_caption_dir, '*.tar'))
+        return sorted(tar_files)
+    
+    def get_grounding_tar_files(self):
+        """Get all grounding tar files from part directories"""
+        grounding_tars = []
+        for part in ['part_1', 'part_2', 'part_3', 'part_4']:
+            part_dir = os.path.join(self.grounding_source_dir, part)
+            if os.path.exists(part_dir):
+                part_tars = glob.glob(os.path.join(part_dir, '*.tar.gz'))
+                grounding_tars.extend(part_tars)
+        return sorted(grounding_tars)
+    
+    def extract_uids_from_rgb_tar(self, tar_path):
+        """Extract all UIDs from RGB tar file"""
+        uids = set()
+        
+        try:
+            with tarfile.open(tar_path, 'r') as tar:
+                for member in tar.getmembers():
+                    if member.isfile():
+                        filename = member.name
+                        uid = os.path.splitext(filename)[0]  # Remove extension
+                        uids.add(uid)
+        except Exception as e:
+            logger.error(f"Error extracting UIDs from {tar_path}: {e}")
+            return set()
+        
+        return uids
+    
+    def build_grounding_index(self):
+        """Build an index of all grounding data: uid -> (tar_path, member_name)"""
+        logger.info("Building grounding data index...")
+        grounding_index = {}
+        
+        grounding_tars = self.get_grounding_tar_files()
+        
+        for tar_path in tqdm(grounding_tars, desc="Indexing grounding tars"):
+            try:
+                with tarfile.open(tar_path, 'r') as tar:
+                    for member in tar.getmembers():
+                        if member.isfile() and member.name.endswith('.json'):
+                            uid = os.path.splitext(member.name)[0]
+                            grounding_index[uid] = (tar_path, member.name)
+            except Exception as e:
+                logger.warning(f"Error indexing {tar_path}: {e}")
+                continue
+        
+        logger.info(f"Built index with {len(grounding_index)} grounding entries")
+        return grounding_index
+    
+    def pre_extract_all_grounding_data_to_disk(self):
+        """Pre-extract all grounding data to disk cache for fastest access"""
+        logger.info("Pre-extracting all grounding data to disk cache...")
+        
+        grounding_tars = self.get_grounding_tar_files()
+        cache_index = {}
+        
+        for tar_path in tqdm(grounding_tars, desc="Pre-extracting grounding data to disk"):
+            try:
+                with tarfile.open(tar_path, 'r') as tar:
+                    for member in tar.getmembers():
+                        if member.isfile() and member.name.endswith('.json'):
+                            uid = os.path.splitext(member.name)[0]
+                            
+                            # Extract JSON content
+                            f = tar.extractfile(member)
+                            if f is not None:
+                                content = f.read()
+                                json_str = content.decode('utf-8')
+                                json_data = json.loads(json_str)
+                                
+                                # Save to disk cache
+                                cache_file = os.path.join(self.grounding_cache_dir, f"{uid}.json")
+                                with open(cache_file, 'w') as cache_f:
+                                    json.dump(json_data, cache_f, indent=2)
+                                
+                                cache_index[uid] = cache_file
+            except Exception as e:
+                logger.warning(f"Error pre-extracting from {tar_path}: {e}")
+                continue
+        
+        # Save index file
+        index_file = os.path.join(self.grounding_cache_dir, "index.json")
+        with open(index_file, 'w') as f:
+            json.dump(cache_index, f, indent=2)
+        
+        logger.info(f"Pre-extracted {len(cache_index)} grounding entries to disk cache")
+        return cache_index
+    
+    def load_grounding_cache_index(self):
+        """Load the grounding cache index"""
+        index_file = os.path.join(self.grounding_cache_dir, "index.json")
+        if os.path.exists(index_file):
+            with open(index_file, 'r') as f:
+                return json.load(f)
+        return {}
+    
+    def load_grounding_data_from_cache(self, uids):
+        """Load grounding data from disk cache for specific UIDs"""
+        uid_to_grounding = {}
+        
+        for uid in uids:
+            cache_file = os.path.join(self.grounding_cache_dir, f"{uid}.json")
+            if os.path.exists(cache_file):
+                try:
+                    with open(cache_file, 'r') as f:
+                        json_data = json.load(f)
+                        uid_to_grounding[uid] = json_data
+                except Exception as e:
+                    logger.warning(f"Error loading {cache_file}: {e}")
+                    continue
+        
+        return uid_to_grounding
+    
+    def find_grounding_data_for_uids_fast(self, uids, grounding_data):
+        """Find grounding JSON data for given UIDs using pre-extracted data - FASTEST"""
+        uid_to_grounding = {}
+        
+        for uid in uids:
+            if uid in grounding_data:
+                uid_to_grounding[uid] = grounding_data[uid]
+        
+        return uid_to_grounding
+    
+    def find_grounding_data_for_uids(self, uids, grounding_index):
+        """Find grounding JSON data for given UIDs using pre-built index with batch processing"""
+        uid_to_grounding = {}
+        
+        # Group UIDs by tar_path to minimize tar file opens
+        tar_to_uids = {}
+        for uid in uids:
+            if uid in grounding_index:
+                tar_path, member_name = grounding_index[uid]
+                if tar_path not in tar_to_uids:
+                    tar_to_uids[tar_path] = []
+                tar_to_uids[tar_path].append((uid, member_name))
+        
+        # Process each tar file once, extracting all needed UIDs
+        for tar_path, uid_members in tar_to_uids.items():
+            try:
+                with tarfile.open(tar_path, 'r') as tar:
+                    for uid, member_name in uid_members:
+                        try:
+                            member = tar.getmember(member_name)
+                            f = tar.extractfile(member)
+                            if f is not None:
+                                content = f.read()
+                                json_str = content.decode('utf-8')
+                                json_data = json.loads(json_str)
+                                uid_to_grounding[uid] = json_data
+                        except Exception as e:
+                            logger.warning(f"Error extracting {uid} from {tar_path}: {e}")
+                            continue
+            except Exception as e:
+                logger.warning(f"Error opening {tar_path}: {e}")
+                continue
+        
+        return uid_to_grounding
+    
+    def create_grounding_tar(self, rgb_tar_path, uid_to_grounding):
+        """Create a grounding tar file with the same structure as RGB tar"""
+        
+        # Extract base name (e.g., 'sa_000000' from 'sa_000000.tar')
+        base_name = os.path.splitext(os.path.basename(rgb_tar_path))[0]
+        output_tar_path = os.path.join(self.grounding_output_dir, f"{base_name}.tar")
+        
+        logger.info(f"Creating grounding tar: {output_tar_path}")
+        
+        try:
+            with tarfile.open(output_tar_path, 'w') as output_tar:
+                for uid, grounding_data in uid_to_grounding.items():
+                    # Create JSON content
+                    json_str = json.dumps(grounding_data, indent=2)
+                    json_bytes = json_str.encode('utf-8')
+                    
+                    # Add to tar with same naming convention as RGB
+                    tarinfo = tarfile.TarInfo(name=f"{uid}.json")
+                    tarinfo.size = len(json_bytes)
+                    output_tar.addfile(tarinfo, io.BytesIO(json_bytes))
+            
+            logger.info(f"Created grounding tar with {len(uid_to_grounding)} UIDs: {output_tar_path}")
+            return output_tar_path
+            
+        except Exception as e:
+            logger.error(f"Error creating grounding tar {output_tar_path}: {e}")
+            return None
+    
+    def process_single_rgb_tar(self, rgb_tar_path, grounding_data=None, grounding_index=None, use_cache=False):
+        """Process a single RGB tar file and create corresponding grounding tar"""
+        
+        # Extract base name
+        base_name = os.path.splitext(os.path.basename(rgb_tar_path))[0]
+        
+        logger.info(f"Processing {base_name}")
+        
+        # Extract UIDs from RGB tar
+        uids = self.extract_uids_from_rgb_tar(rgb_tar_path)
+        
+        if not uids:
+            logger.warning(f"No UIDs found in {rgb_tar_path}")
+            return None
+        
+        logger.info(f"Found {len(uids)} UIDs in {base_name}")
+        
+        # Find grounding data for these UIDs
+        if use_cache:
+            # Use disk cache (fastest and memory efficient)
+            uid_to_grounding = self.load_grounding_data_from_cache(uids)
+        elif grounding_data is not None:
+            # Use pre-extracted data (fastest)
+            uid_to_grounding = self.find_grounding_data_for_uids_fast(uids, grounding_data)
+        elif grounding_index is not None:
+            # Use index-based lookup (faster than original)
+            uid_to_grounding = self.find_grounding_data_for_uids(uids, grounding_index)
+        else:
+            logger.error("No grounding data or index provided")
+            return None
+        
+        if not uid_to_grounding:
+            logger.warning(f"No grounding data found for UIDs in {base_name}")
+            return None
+        
+        logger.info(f"Found grounding data for {len(uid_to_grounding)}/{len(uids)} UIDs")
+        
+        # Create grounding tar file
+        output_path = self.create_grounding_tar(rgb_tar_path, uid_to_grounding)
+        
+        return output_path
+    
+    def process_all_rgb_tars(self, use_cache=True, pre_extract_cache=False):
+        """Process all RGB tar files and create corresponding grounding tars"""
+        rgb_tars = self.get_rgb_tar_files()
+        
+        if not rgb_tars:
+            logger.error("No RGB tar files found")
+            return
+        
+        logger.info(f"Found {len(rgb_tars)} RGB tar files to process")
+        
+        if use_cache:
+            # Check if cache exists
+            cache_index = self.load_grounding_cache_index()
+            
+            if not cache_index:
+                if pre_extract_cache:
+                    logger.info("Pre-extracting grounding data to disk cache...")
+                    self.pre_extract_all_grounding_data_to_disk()
+                else:
+                    logger.error("No cache found. Use --pre_extract_cache to create cache first.")
+                    return
+            
+            logger.info("Using disk cache method (fastest and memory efficient)")
+            processed_files = []
+            for rgb_tar in tqdm(rgb_tars, desc="Processing RGB tars"):
+                try:
+                    output_path = self.process_single_rgb_tar(rgb_tar, use_cache=True)
+                    if output_path:
+                        processed_files.append(output_path)
+                except Exception as e:
+                    logger.error(f"Error processing {rgb_tar}: {e}")
+        else:
+            # Use index-based method (faster than original, less memory)
+            logger.info("Using index-based method (faster, less memory)")
+            grounding_index = self.build_grounding_index()
+            
+            processed_files = []
+            for rgb_tar in tqdm(rgb_tars, desc="Processing RGB tars"):
+                try:
+                    output_path = self.process_single_rgb_tar(rgb_tar, grounding_index=grounding_index)
+                    if output_path:
+                        processed_files.append(output_path)
+                except Exception as e:
+                    logger.error(f"Error processing {rgb_tar}: {e}")
+        
+        logger.info(f"Successfully processed {len(processed_files)} grounding tar files")
+        return processed_files
+
+def main():
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Process grounding data')
+    parser.add_argument('--use_cache', action='store_true', default=True,
+                       help='Use disk cache method (fastest and memory efficient)')
+    parser.add_argument('--pre_extract_cache', action='store_true',
+                       help='Pre-extract grounding data to disk cache')
+    parser.add_argument('--base_dir', type=str, default='./',
+                       help='Base directory (default: ./)')
+    
+    args = parser.parse_args()
+    
+    # Set the base directory
+    base_dir = args.base_dir
+    
+    # Create processor
+    processor = GroundingProcessor(base_dir)
+    
+    # Process all RGB tar files
+    processed_files = processor.process_all_rgb_tars(
+        use_cache=args.use_cache,
+        pre_extract_cache=args.pre_extract_cache
+    )
+    
+    if processed_files:
+        logger.info(f"Successfully created {len(processed_files)} grounding tar files:")
+        for file in processed_files:
+            logger.info(f"  - {file}")
+    else:
+        logger.error("No grounding tar files were created")
+
+if __name__ == "__main__":
+    main()
+
+
+
+# This will take time but only needs to be done once
+# python data/any2any_preprocess/process_grounding.py --pre_extract_cache
+
+# This will be extremely fast
+# python data/any2any_preprocess/process_grounding.py --use_cache

data/any2any_preprocess/recompress_normal_jpeg.py

@@ -0,0 +1,90 @@
+"""Recompress the `normal` column of the 16-mod parquet from PNG to JPEG q95.
+
+The `normal` column stores raw image bytes (Marigold normal maps, pass-through
+PNG) and accounts for ~65% of on-disk size while being ~6x larger than the
+JPEG-encoded `rgb` column. Re-encoding normal as JPEG q95 shrinks it to
+roughly rgb size with negligible quality loss.
+
+Every other column is passed through untouched; the Arrow schema is identical
+to the source, so downstream dataloader / HF viewer see no structural change.
+
+Usage:
+    python recompress_normal_jpeg.py --out_dir <new_folder> file1.parquet [file2 ...]
+"""
+import argparse
+import io
+import os
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from PIL import Image
+
+JPEG_QUALITY = 95
+
+
+def recompress_normal_bytes(b):
+    if b is None:
+        return None
+    if isinstance(b, memoryview):
+        b = b.tobytes()
+    img = Image.open(io.BytesIO(b)).convert('RGB')
+    out = io.BytesIO()
+    img.save(out, format='JPEG', quality=JPEG_QUALITY)
+    return out.getvalue()
+
+
+def process_file(src, out_dir):
+    dst = os.path.join(out_dir, os.path.basename(src))
+    pf = pq.ParquetFile(src)
+    schema = pf.schema_arrow
+    assert 'normal' in schema.names, f'no normal column in {src}'
+    norm_idx = schema.names.index('normal')
+
+    reported_fmt = False
+    writer = pq.ParquetWriter(dst, schema, compression='snappy')
+    try:
+        for rg in range(pf.num_row_groups):
+            table = pf.read_row_group(rg)
+            normal_col = table.column('normal').to_pylist()
+
+            if not reported_fmt:
+                for b in normal_col:
+                    if b is not None:
+                        bb = b.tobytes() if isinstance(b, memoryview) else b
+                        fmt = Image.open(io.BytesIO(bb)).format
+                        print(f'    source normal format = {fmt}', flush=True)
+                        reported_fmt = True
+                        break
+
+            new_normal = [recompress_normal_bytes(b) for b in normal_col]
+            new_col = pa.array(new_normal, type=schema.field('normal').type)
+            table = table.set_column(norm_idx, schema.field('normal'), new_col)
+            writer.write_table(table)
+    finally:
+        writer.close()
+
+    s0, s1 = os.path.getsize(src), os.path.getsize(dst)
+    print(f'    {os.path.basename(src)}: {s0/1e9:.2f} GB -> {s1/1e9:.2f} GB '
+          f'({100*(1-s1/s0):.0f}% smaller)', flush=True)
+    return s0, s1
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--out_dir', required=True)
+    ap.add_argument('files', nargs='+')
+    args = ap.parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    tot0 = tot1 = 0
+    for f in args.files:
+        print(f'[{os.path.basename(f)}]', flush=True)
+        s0, s1 = process_file(f, args.out_dir)
+        tot0 += s0
+        tot1 += s1
+    print(f'TOTAL: {tot0/1e9:.2f} GB -> {tot1/1e9:.2f} GB '
+          f'({100*(1-tot1/tot0):.0f}% smaller)', flush=True)
+
+
+if __name__ == '__main__':
+    main()

data/any2any_preprocess/run_full_rebuild.py

@@ -0,0 +1,82 @@
+"""Slurm driver: rebuild a strided slice of the full dataset.
+
+Robust against memory leaks: each small chunk of files is processed by a FRESH
+subprocess (build_full_release.py --workers 1) that exits afterwards, so the OS
+reclaims all memory — no accumulation across the run. `concurrency` such
+subprocesses run at once. Resumable via process_one's skip.
+"""
+import argparse
+import os
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+SRC = ('datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg_grounding'
+       '_canny_dino_global_clip448_imagebind_samseg_samedge_cocodet')
+OUT = 'datasets/blip3o/modus_full'
+WORKER = 'data/any2any_preprocess/build_full_release.py'
+CHUNK = 4
+
+
+def config_of(fname):
+    if fname.startswith('sa_'):
+        return 'sa1b'
+    if fname.startswith('webdataset_JDB_'):
+        return 'journeydb'
+    if fname.startswith('webdataset_shard_'):
+        return 'cc12m'
+    return None
+
+
+def run_chunk(cfg, files):
+    cmd = [sys.executable, WORKER, '--config', cfg,
+           '--out_dir', os.path.join(OUT, cfg), '--workers', '1'] + files
+    r = subprocess.run(cmd, capture_output=True, text=True)
+    return cfg, len(files), r.returncode, (r.stderr or '')[-300:]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--num-tasks', type=int, required=True)
+    ap.add_argument('--task-id', type=int, required=True)
+    ap.add_argument('--concurrency', type=int, default=48)
+    args = ap.parse_args()
+
+    for cfg in ('sa1b', 'journeydb', 'cc12m'):
+        os.makedirs(os.path.join(OUT, cfg), exist_ok=True)
+
+    items = []
+    for fn in sorted(os.listdir(SRC)):
+        if fn.endswith('.parquet') and config_of(fn):
+            items.append((config_of(fn), os.path.join(SRC, fn)))
+    mine = items[args.task_id::args.num_tasks]
+
+    # group into per-config chunks
+    chunks = []
+    bycfg = {}
+    for cfg, src in mine:
+        bycfg.setdefault(cfg, []).append(src)
+    for cfg, fs in bycfg.items():
+        for i in range(0, len(fs), CHUNK):
+            chunks.append((cfg, fs[i:i + CHUNK]))
+
+    print(f'task {args.task_id}/{args.num_tasks}: {len(mine)} files in '
+          f'{len(chunks)} chunks, concurrency={args.concurrency}', flush=True)
+    done_files = fail = 0
+    with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
+        futs = [ex.submit(run_chunk, cfg, fs) for cfg, fs in chunks]
+        for i, fu in enumerate(as_completed(futs), 1):
+            cfg, n, rc, err = fu.result()
+            if rc == 0:
+                done_files += n
+            else:
+                fail += n
+                print(f'  CHUNK FAIL {cfg} rc={rc}: {err}', flush=True)
+            if i % 20 == 0:
+                print(f'[{i}/{len(chunks)} chunks] files_ok={done_files} '
+                      f'fail={fail}', flush=True)
+    print(f'TASK COMPLETE: {done_files} files ok, {fail} failed', flush=True)
+
+
+if __name__ == '__main__':
+    main()

data/any2any_preprocess/upload_full.py

@@ -0,0 +1,26 @@
+"""Resumable upload of the full-release folder to the HF dataset repo.
+
+Uses upload_large_folder (xet, multi-commit, resumable). Uploads folder_path
+contents to the repo root, preserving the sa1b/ journeydb/ cc12m/ structure.
+Re-running skips already-uploaded files.
+
+Usage:
+  HF_TOKEN_FILE=/users/mye/.hf_token python upload_full.py <repo_id> <folder> [allow_glob]
+"""
+import os
+import sys
+
+from huggingface_hub import HfApi
+
+repo_id = sys.argv[1]
+folder = sys.argv[2]
+allow = list(sys.argv[3:])  # zero or more allow_patterns
+token = open(os.environ['HF_TOKEN_FILE']).read().strip()
+
+api = HfApi(token=token)
+api.create_repo(repo_id, repo_type='dataset', exist_ok=True)
+kw = dict(repo_id=repo_id, repo_type='dataset', folder_path=folder, print_report=True)
+if allow:
+    kw['allow_patterns'] = allow
+api.upload_large_folder(**kw)
+print(f'UPLOAD PASS DONE: {folder} {allow} -> {repo_id}', flush=True)

data/any2any_preprocess/vqa_convert_parquet_to_jsonl.py

@@ -0,0 +1,312 @@
+#!/usr/bin/env python3
+"""
+Convert LLaVA OneVision parquet files to JSONL format matching llava_ov_si.jsonl
+
+Single-folder usage (backwards compatible):
+    python vqa_convert_parquet_to_jsonl.py <parquet_dir> <output_jsonl> --image-dir <image_dir> [--max-rows N]
+
+Multi-folder usage (process all immediate subfolders of <parquet_root>):
+    python vqa_convert_parquet_to_jsonl.py <parquet_root> dummy.jsonl --process-all --output-root <output_root> [--max-rows N]
+
+Notes:
+- In multi-folder mode, outputs are written per subfolder under <output_root>/<subfolder>/{images,label/labels.jsonl}
+- The script prints the sorted folder index and name, and a global image counter every N (default 1000)
+"""
+
+import argparse
+import json
+import os
+import glob
+from pathlib import Path
+import pandas as pd
+from PIL import Image
+import io
+from tqdm import tqdm
+def _save_pil_image_as_jpeg(pil_image: Image.Image, image_path: str):
+    """
+    Ensure the PIL image is compatible with JPEG (no alpha), then save.
+    """
+    try:
+        # Convert to RGB if image has alpha channel or unsupported mode for JPEG
+        if pil_image.mode not in ("RGB", "L"):
+            pil_image = pil_image.convert("RGB")
+        pil_image.save(image_path, format="JPEG")
+    except Exception:
+        # As a fallback, force convert to RGB and retry
+        pil_image = pil_image.convert("RGB")
+        pil_image.save(image_path, format="JPEG")
+
+
+
+def convert_parquet_to_jsonl(
+    parquet_dir: str,
+    output_jsonl: str,
+    image_dir: str = None,
+    max_rows: int = None,
+    sample_id_offset: int = 0,
+    progress: dict | None = None,
+):
+    """
+    Convert parquet files to JSONL format compatible with llava_ov_si.jsonl
+
+    Args:
+        parquet_dir: Directory containing parquet files (can be subdirectory)
+        output_jsonl: Output JSONL file path
+        image_dir: Directory to save images (optional, if None saves with dataset)
+        max_rows: Maximum number of rows to process (for testing)
+        sample_id_offset: Starting ID offset (for merging multiple datasets)
+        progress: Optional dict to track global progress across folders. Expected keys:
+                  {"count": int, "next_milestone": int, "interval": int}
+    """
+
+    # Find all parquet files in the directory
+    parquet_pattern = os.path.join(parquet_dir, "train-*.parquet")
+    parquet_files = sorted(glob.glob(parquet_pattern))
+
+    if not parquet_files:
+        print(f"Error: No parquet files found in {parquet_dir}")
+        return
+
+    print(f"Found {len(parquet_files)} parquet file(s)")
+
+    # Create image directory if specified
+    if image_dir:
+        os.makedirs(image_dir, exist_ok=True)
+
+    all_entries = []
+    current_id = sample_id_offset
+
+    for parquet_file in parquet_files:
+        print(f"\nProcessing {parquet_file}...")
+        df = pd.read_parquet(parquet_file)
+
+        for idx, row in tqdm(df.iterrows(), total=len(df)):
+            if max_rows and len(all_entries) >= max_rows:
+                break
+
+            # Extract conversations - handle numpy array conversion
+            conversations = row["conversations"]
+            if hasattr(conversations, "tolist"):
+                conversations = conversations.tolist()
+
+            # Handle image - extract from binary data and save
+            image_filename = None
+            if row["image"] is not None:
+                try:
+                    # The image data is stored as binary bytes
+                    image_data = row["image"]
+
+                    # Generate filename from ID
+                    image_filename = f"{current_id}.jpg"
+
+                    if image_dir:
+                        image_path = os.path.join(image_dir, image_filename)
+                    else:
+                        # Save in same directory as parquet file
+                        image_dir_local = os.path.dirname(parquet_file)
+                        image_path = os.path.join(
+                            image_dir_local, image_filename
+                        )
+
+                    # Convert image data to PIL Image and save
+                    if isinstance(image_data, dict):
+                        # Image is stored as a dict with 'bytes' key containing image data
+                        if (
+                            "bytes" in image_data
+                            and image_data["bytes"] is not None
+                        ):
+                            image_bytes = image_data["bytes"]
+                            pil_image = Image.open(io.BytesIO(image_bytes))
+                            _save_pil_image_as_jpeg(pil_image, image_path)
+                        else:
+                            print(
+                                f"Warning: Dict image data doesn't contain 'bytes' key for row {idx}"
+                            )
+                            continue
+                    elif isinstance(image_data, bytes):
+                        # Try to open as PIL Image from bytes
+                        pil_image = Image.open(io.BytesIO(image_data))
+                        _save_pil_image_as_jpeg(pil_image, image_path)
+                    elif isinstance(image_data, str):
+                        # If it's a string, try to decode it
+                        image_bytes = image_data.encode(
+                            "latin-1"
+                        )  # Convert string to bytes
+                        pil_image = Image.open(io.BytesIO(image_bytes))
+                        _save_pil_image_as_jpeg(pil_image, image_path)
+                    else:
+                        print(
+                            f"Warning: Unknown image data type {type(image_data)} for row {idx}"
+                        )
+                        continue
+
+                except Exception as e:
+                    print(f"Error processing image for row {idx}: {e}")
+                    continue
+
+            # Create JSONL entry matching llava_ov_si.jsonl format
+            entry = {
+                "id": current_id,
+                "image": image_filename,
+                "conversations": conversations,
+            }
+
+            all_entries.append(entry)
+            current_id += 1
+
+            # Update global progress if provided
+            if progress is not None:
+                progress["count"] += 1
+                if progress["count"] >= progress.get("next_milestone", 1000):
+                    print(
+                        f"[global] Processed {progress['count']} images total"
+                    )
+                    interval = progress.get("interval", 1000)
+                    progress["next_milestone"] = progress["count"] + interval
+
+            if max_rows and len(all_entries) >= max_rows:
+                break
+
+        if max_rows and len(all_entries) >= max_rows:
+            break
+
+    # Write JSONL file
+    print(f"\nWriting {len(all_entries)} entries to {output_jsonl}...")
+    with open(output_jsonl, "w", encoding="utf-8") as f:
+        for entry in all_entries:
+            f.write(json.dumps(entry, ensure_ascii=False) + "\n")
+
+    print(f"✓ Converted {len(all_entries)} samples to {output_jsonl}")
+    if image_dir:
+        print(f"✓ Images saved to {image_dir}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert LLaVA OneVision parquet files to JSONL format"
+    )
+    parser.add_argument(
+        "parquet_dir",
+        help="Directory containing parquet files, or a root directory with subfolders",
+    )
+    parser.add_argument("output_jsonl", help="Output JSONL file path (single-folder mode)")
+    parser.add_argument(
+        "--image-dir",
+        default=None,
+        help="Directory to save images (default: same as parquet directory)",
+    )
+    parser.add_argument(
+        "--max-rows",
+        type=int,
+        default=None,
+        help="Maximum number of rows to process (for testing)",
+    )
+    parser.add_argument(
+        "--id-offset",
+        type=int,
+        default=0,
+        help="Starting ID offset (for merging datasets)",
+    )
+    parser.add_argument(
+        "--process-all",
+        action="store_true",
+        help="Process all immediate subfolders in parquet_dir and write per-folder outputs under --output-root",
+    )
+    parser.add_argument(
+        "--output-root",
+        default=None,
+        help="Base output directory when using --process-all",
+    )
+    parser.add_argument(
+        "--print-every",
+        type=int,
+        default=1000,
+        help="Global image progress interval for multi-folder mode",
+    )
+    parser.add_argument(
+        "--start-index",
+        type=int,
+        default=1,
+        help="1-based index of subfolder to start processing from (resume)",
+    )
+
+    args = parser.parse_args()
+
+    if args.process_all:
+        if not args.output_root:
+            raise SystemExit("--output-root is required when using --process-all")
+
+        # Discover immediate subfolders and sort
+        if not os.path.isdir(args.parquet_dir):
+            raise SystemExit(f"Input root '{args.parquet_dir}' is not a directory")
+
+        subfolders = [
+            d
+            for d in sorted(os.listdir(args.parquet_dir))
+            if os.path.isdir(os.path.join(args.parquet_dir, d))
+        ]
+
+        # Filter to only those having parquet files
+        filtered_subfolders = []
+        for d in subfolders:
+            pattern = os.path.join(args.parquet_dir, d, "train-*.parquet")
+            if glob.glob(pattern):
+                filtered_subfolders.append(d)
+        subfolders = filtered_subfolders
+
+        total = len(subfolders)
+        print(f"Found {total} subfolder(s) to process under {args.parquet_dir}")
+
+        # Compute starting index (1-based) and slice subfolders for resume
+        start_index = max(1, int(args.start_index)) if args.start_index else 1
+        if start_index > total:
+            print(f"Start index {start_index} is beyond total {total}; nothing to do.")
+            return
+        subfolders_to_process = subfolders[start_index - 1 :]
+
+        # Prepare global progress
+        progress = {
+            "count": 0,
+            "next_milestone": args.print_every,
+            "interval": args.print_every,
+        }
+        # When resuming at folder index 41, hardcode the global counter baseline
+        if start_index == 41:
+            progress["count"] = 1380000
+            progress["next_milestone"] = progress["count"] + progress["interval"]
+
+        for idx, sub in enumerate(subfolders_to_process, start=start_index):
+            print(f"\n[{idx}/{total}] Processing folder: {sub}")
+            sub_in = os.path.join(args.parquet_dir, sub)
+            sub_out_images = os.path.join(args.output_root, sub, "images")
+            sub_out_label_dir = os.path.join(args.output_root, sub, "label")
+            os.makedirs(sub_out_images, exist_ok=True)
+            os.makedirs(sub_out_label_dir, exist_ok=True)
+            sub_out_jsonl = os.path.join(sub_out_label_dir, "labels.jsonl")
+
+            convert_parquet_to_jsonl(
+                parquet_dir=sub_in,
+                output_jsonl=sub_out_jsonl,
+                image_dir=sub_out_images,
+                max_rows=args.max_rows,
+                sample_id_offset=args.id_offset,
+                progress=progress,
+            )
+    else:
+        convert_parquet_to_jsonl(
+            parquet_dir=args.parquet_dir,
+            output_jsonl=args.output_jsonl,
+            image_dir=args.image_dir,
+            max_rows=args.max_rows,
+            sample_id_offset=args.id_offset,
+        )
+
+
+if __name__ == "__main__":
+    main()
+
+
+# python data/any2any_preprocess/vqa_convert_parquet_to_jsonl.py datasets/llava_onevision/ai2d\(cauldron\,llava_format\)/ datasets/llava_onevision_vqa/label/ai2d_converted.jsonl --image-dir datasets/llava_onevision_vqa/images
+
+# python data/any2any_preprocess/vqa_convert_parquet_to_jsonl.py datasets/llava_onevision dummy.jsonl --process-all --output-root datasets/llava_onevision_vqa
+# python data/any2any_preprocess/vqa_convert_parquet_to_jsonl.py datasets/llava_onevision dummy.jsonl --process-all --output-root datasets/llava_onevision_vqa --start-index 41

data/bundled_parquet_info/README.md

@@ -0,0 +1,43 @@
+# bundled_parquet_info — repo-tracked parquet metadata (auto-fallback)
+
+Repo-tracked copies of the parquet-info JSONs used by BAGEL training.
+**The code reads this directory automatically**: `data/dataset_base.py`
+falls back to `data/bundled_parquet_info/<basename>` whenever the
+configured live path (`./datasets/blip3o/parquet_info/…`, a symlink chain
+that only exists on the original cluster) is missing. When the live path
+exists it still wins — behavior on the original cluster is unchanged.
+
+## Files
+
+| File | Files listed | Rows | Used by |
+|---|---|---|---|
+| `…_grounding_canny_dino_global_clip448_imagebind.json` | 2891 (1000 `sa_*` + 1891 `webdataset_*`) | 29.2M | 13mod: all target groups except grounding (any2rgb / depth / normal / caption / canny / dino / dinolocal / clip / imagebind / imagebindlocal) |
+| `…_grounding2_canny_dino_global_clip448_imagebind.json` | 1000 (`sa_*` only) | 11.1M | 13mod: `unified_any2grounding` only |
+| `…_clip448_imagebind_samseg_samedge_cocodet.json` | — | — | 16mod runs (adds cocodet / samseg / samedge) |
+| `llava_onevision_vqa.json` | — | — | VLM-SFT (llava-onevision) |
+
+**Why grounding has two JSONs:** grounding annotations (GLaMM phrase+bbox,
+matched by UID) only exist for the SA-1B portion of the parquet set. The
+`webdataset_*` shards have no `grounding` column at all (verified via
+parquet footer schema). So when grounding is the generation TARGET, the
+loader must restrict itself to the `sa_*` files — that restricted file
+list is the `grounding2` JSON. Both JSONs point into the SAME `data_dir`.
+
+## Setting up on a new machine
+
+1. Get the parquet data itself (≈ several TB, NOT in this repo):
+   `parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global_clip448_imagebind/`
+   (2891 parquet files).
+2. Place (or symlink) it at
+   `datasets/blip3o/parquet_rgb_caption_depth_normal_det_seg_grounding_canny_dino_global_clip448_imagebind/`
+   relative to the repo root.
+3. parquet-info JSONs: **no action needed** — the loader automatically
+   falls back to this directory when the live
+   `datasets/blip3o/parquet_info/` path is absent
+   (`data/dataset_base.py`, "Fall back to the copy bundled in the repo").
+4. Note the JSON keys are relative paths
+   (`./datasets/blip3o/parquet_.../sa_000000.parquet`), so always launch
+   training from the repo root.
+
+If the parquet set ever changes, regenerate with
+`data/any2any_preprocess/generate_parquet_json.py` and refresh these copies.