Initial sapiens2-normal Space (HF download at startup, all 4 sizes)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+*.pyc
+default.profraw
+.DS_Store
+*.log

README.md

@@ -1,12 +1,23 @@
 ---
 title: Sapiens2 Normal
-emoji: 🏢
-colorFrom: green
-colorTo: gray
+emoji: 🧊
+colorFrom: purple
+colorTo: blue
 sdk: gradio
-sdk_version: 6.13.0
+sdk_version: 4.42.0
 app_file: app.py
+python_version: "3.12"
 pinned: false
+license: other
+license_name: sapiens2-license
+license_link: https://github.com/facebookresearch/sapiens2/blob/main/LICENSE.md
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Sapiens2: Surface Normal Estimation
+### ICLR 2026
+
+Per-pixel surface-normal estimation (3-channel unit vectors in camera frame).
+
+- **Code:** [github.com/facebookresearch/sapiens2](https://github.com/facebookresearch/sapiens2)
+- **Models:** [Sapiens2 collection](https://huggingface.co/facebook/sapiens2)
+- **Paper:** https://openreview.net/pdf?id=IVAlYCqdvW

app.py

@@ -0,0 +1,167 @@
+"""Sapiens2 surface-normal Gradio Space.
+
+Image → per-pixel surface normals. Visualized by RGB-encoding the unit-length
+(x, y, z) normal: r = (x + 1) / 2, g = (y + 1) / 2, b = (z + 1) / 2.
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import tempfile
+
+import cv2
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from huggingface_hub import hf_hub_download
+from sapiens.dense.models import NormalEstimator, init_model  # NormalEstimator triggers registry
+_ = NormalEstimator
+
+
+# -----------------------------------------------------------------------------
+# Config
+
+ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
+CONFIGS_DIR = os.path.join(ASSETS_DIR, "configs")
+
+NORMAL_MODELS = {
+    "0.4B": {
+        "repo": "facebook/sapiens2-normal-0.4b",
+        "filename": "sapiens2_0.4b_normal.safetensors",
+        "config": os.path.join(CONFIGS_DIR, "sapiens2_0.4b_normal_metasim_render_people-1024x768.py"),
+    },
+    "0.8B": {
+        "repo": "facebook/sapiens2-normal-0.8b",
+        "filename": "sapiens2_0.8b_normal.safetensors",
+        "config": os.path.join(CONFIGS_DIR, "sapiens2_0.8b_normal_metasim_render_people-1024x768.py"),
+    },
+    "1B": {
+        "repo": "facebook/sapiens2-normal-1b",
+        "filename": "sapiens2_1b_normal.safetensors",
+        "config": os.path.join(CONFIGS_DIR, "sapiens2_1b_normal_metasim_render_people-1024x768.py"),
+    },
+    "5B": {
+        "repo": "facebook/sapiens2-normal-5b",
+        "filename": "sapiens2_5b_normal.safetensors",
+        "config": os.path.join(CONFIGS_DIR, "sapiens2_5b_normal_metasim_render_people-1024x768.py"),
+    },
+}
+DEFAULT_SIZE = "1B"
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# -----------------------------------------------------------------------------
+# Model cache
+
+_normal_model_cache: dict = {}
+
+
+def _get_normal_model(size: str):
+    if size not in _normal_model_cache:
+        spec = NORMAL_MODELS[size]
+        ckpt = hf_hub_download(repo_id=spec["repo"], filename=spec["filename"])
+        model = init_model(spec["config"], ckpt, device=DEVICE)
+        _normal_model_cache[size] = model
+    return _normal_model_cache[size]
+
+
+print("[startup] pre-loading all normal sizes ...")
+for _size in NORMAL_MODELS:
+    _get_normal_model(_size)
+print("[startup] ready.")
+
+
+# -----------------------------------------------------------------------------
+# Inference
+
+def _estimate_normal(image_bgr: np.ndarray, model) -> np.ndarray:
+    h0, w0 = image_bgr.shape[:2]
+    data = model.pipeline(dict(img=image_bgr))
+    data = model.data_preprocessor(data)
+    inputs = data["inputs"]
+    if inputs.ndim == 3:
+        inputs = inputs.unsqueeze(0)
+
+    with torch.no_grad():
+        normals = model(inputs)  # (1, 3, H, W)
+
+    # Unit-length normalization, interpolate to original size, cast to numpy
+    normals = normals / normals.norm(dim=1, keepdim=True).clamp_min(1e-6)
+    normals = F.interpolate(normals, size=(h0, w0), mode="bilinear", align_corners=False)
+    normals = normals[0].cpu().float().numpy()  # (3, H, W) in [-1, 1]
+    return normals.transpose(1, 2, 0)  # (H, W, 3)
+
+
+def _normal_to_rgb(normal_hwc: np.ndarray) -> np.ndarray:
+    rgb = (((normal_hwc + 1.0) / 2.0) * 255.0).clip(0, 255).astype(np.uint8)
+    return rgb[:, :, ::-1]  # match training viz channel order
+
+
+# -----------------------------------------------------------------------------
+# Gradio handler
+
+@spaces.GPU(duration=120)
+def predict(image: Image.Image, size: str):
+    if image is None:
+        return None, None
+
+    image_rgb = np.array(image.convert("RGB"))
+    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
+
+    model = _get_normal_model(size)
+    normals = _estimate_normal(image_bgr, model)  # (H, W, 3) in [-1, 1]
+    rgb = _normal_to_rgb(normals)
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f:
+        np.save(f.name, normals.astype(np.float32))
+        npy_path = f.name
+
+    return Image.fromarray(rgb), npy_path
+
+
+# -----------------------------------------------------------------------------
+# UI
+
+EXAMPLES = sorted(
+    os.path.join(ASSETS_DIR, "images", n)
+    for n in os.listdir(os.path.join(ASSETS_DIR, "images"))
+    if n.lower().endswith((".jpg", ".jpeg", ".png"))
+)
+
+with gr.Blocks(title="Sapiens2 Normal", theme=gr.themes.Default()) as demo:
+    gr.Markdown(
+        "# Sapiens2: Surface Normal Estimation\n"
+        "### ICLR 2026\n"
+        "Per-pixel surface-normal estimation. Output is RGB-encoded (x, y, z → R, G, B).\n\n"
+        "[Code](https://github.com/facebookresearch/sapiens2) · "
+        "[Models](https://huggingface.co/facebook/sapiens2) · "
+        "[Paper](https://openreview.net/pdf?id=IVAlYCqdvW)"
+    )
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Image(label="Input", type="pil")
+            size = gr.Radio(
+                choices=list(NORMAL_MODELS.keys()),
+                value=DEFAULT_SIZE,
+                label="Model size",
+            )
+            run = gr.Button("Run", variant="primary")
+            gr.Examples(examples=EXAMPLES, inputs=inp, examples_per_page=14)
+        with gr.Column():
+            out_img = gr.Image(label="Surface normal (RGB-encoded)", type="pil")
+            out_npy = gr.File(label="Raw normals (.npy float32 [-1, 1])")
+
+    run.click(predict, inputs=[inp, size], outputs=[out_img, out_npy])
+
+
+if __name__ == "__main__":
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    demo.launch(share=False)

assets/configs/sapiens2_0.4b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.4b"
+embed_dim = 1024
+num_layers = 24
+num_heads = 16
+
+layer_decay_rate = 0.8
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.4b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=8,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=2.0, norm_type=2.0)

assets/configs/sapiens2_0.8b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 1e4
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.8b"
+embed_dim = 1280
+num_layers = 32
+num_heads = 16
+
+layer_decay_rate = 0.85
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.8b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=8,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

assets/configs/sapiens2_1b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,306 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+# num_iters = 1e4  ## light finetune
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_1b"
+embed_dim = 1536
+num_layers = 40
+num_heads = 24
+layer_decay_rate = 0.9
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_1b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        # with_cls_token=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

assets/configs/sapiens2_5b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,312 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+# num_iters = 1e4  ## light finetune
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_5b"
+embed_dim = 2432
+num_layers = 56
+num_heads = 32
+layer_decay_rate = 0.94
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_5b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+        # parallelism_cfg=dict(
+        #     dp_shard_size=2,  # Fully Sharded Data Parallel degree
+        #     dp_replicate_size=1,  # Data Parallel degree
+        #     tp_size=1,  # Tensor Parallel degree
+        #     cp_size=4,  # Context Parallel degree
+        # ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[1536, 768, 512, 256],  ## 1K resolution
+        conv_out_channels=[128, 64, 32],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    # lr=5e-4,
+    lr=1e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

requirements.txt

@@ -0,0 +1,21 @@
+gradio==4.42.0
+spaces
+
+torch==2.7.1
+torchvision==0.22.1
+
+numpy
+opencv-python
+pillow
+matplotlib
+safetensors
+huggingface_hub
+
+# Sapiens2 deps (sapiens2 source is vendored under ./sapiens/, not pip-installed).
+tqdm
+scipy
+iopath
+prettytable
+termcolor
+accelerate
+rich

sapiens/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .version import __version__
+from .engine import *
+from .backbones import *
+from .dense import *
+from .pose import *
+from .registry import *
+
+__all__ = ["__version__"]

sapiens/backbones/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .sapiens import Sapiens
+from .sapiens2 import Sapiens2
+
+__all__ = ["Sapiens", "Sapiens2"]

sapiens/backbones/sapiens.py

@@ -0,0 +1,611 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sapiens.engine.models.base_model import BaseModel
+from sapiens.registry import MODELS
+from torch.nn import Linear, Sequential
+
+
+# ----------------------------------------------------------------------------
+def to_2tuple(x):
+    if isinstance(x, (str, bytes)):
+        return (x, x)
+    if isinstance(x, Sequence):
+        x = tuple(x)
+        if len(x) == 2:
+            return x
+        raise ValueError("Expected scalar or length-2 iterable")
+    return (x, x)
+
+
+def resize_pos_embed(
+    pos_embed, src_shape, dst_shape, mode="bicubic", num_extra_tokens=1
+):
+    if src_shape[0] == dst_shape[0] and src_shape[1] == dst_shape[1]:
+        return pos_embed
+    assert pos_embed.ndim == 3, "shape of pos_embed must be [1, L, C]"
+    _, L, C = pos_embed.shape
+    src_h, src_w = src_shape
+    assert L == src_h * src_w + num_extra_tokens, (
+        f"The length of `pos_embed` ({L}) doesn't match the expected "
+        f"shape ({src_h}*{src_w}+{num_extra_tokens}). Please check the"
+        "`img_size` argument."
+    )
+    extra_tokens = pos_embed[:, :num_extra_tokens]
+
+    src_weight = pos_embed[:, num_extra_tokens:]
+    src_weight = src_weight.reshape(1, src_h, src_w, C).permute(0, 3, 1, 2)
+
+    # The cubic interpolate algorithm only accepts float32
+    dst_weight = F.interpolate(
+        src_weight.float(), size=dst_shape, align_corners=False, mode=mode
+    )
+    dst_weight = torch.flatten(dst_weight, 2).transpose(1, 2)
+    dst_weight = dst_weight.to(src_weight.dtype)
+
+    return torch.cat((extra_tokens, dst_weight), dim=1)
+
+
+# ----------------------------------------------------------------------------
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        embed_dims=768,
+        kernel_size=16,
+        stride=16,
+        padding="corner",
+        dilation=1,
+        bias=True,
+        input_size=None,
+    ):
+        super().__init__()
+
+        self.embed_dims = embed_dims
+        if stride is None:
+            stride = kernel_size
+
+        kernel_size = to_2tuple(kernel_size)
+        stride = to_2tuple(stride)
+        dilation = to_2tuple(dilation)
+        padding = 0
+        padding = to_2tuple(padding)
+
+        self.projection = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=embed_dims,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+        if input_size:
+            input_size = to_2tuple(input_size)
+            self.init_input_size = input_size
+
+            h_out = (
+                input_size[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1
+            ) // stride[0] + 1
+            w_out = (
+                input_size[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1
+            ) // stride[1] + 1
+            self.init_out_size = (h_out, w_out)
+        else:
+            self.init_input_size = None
+            self.init_out_size = None
+
+    def forward(self, x):
+        x = self.projection(x)
+        out_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)
+        return x, out_size
+
+
+# ----------------------------------------------------------------------------
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inplace: bool = False,
+        data_format: str = "channels_last",
+        scale: float = 1e-5,
+    ):
+        super().__init__()
+        assert data_format in (
+            "channels_last",
+            "channels_first",
+        ), "'data_format' could only be channels_last or channels_first."
+        self.inplace = inplace
+        self.data_format = data_format
+        self.weight = nn.Parameter(torch.ones(dim) * scale)
+
+    def forward(self, x) -> torch.Tensor:
+        if self.data_format == "channels_first":
+            shape = tuple((1, -1, *(1 for _ in range(x.dim() - 2))))
+        else:
+            shape = tuple((*(1 for _ in range(x.dim() - 1)), -1))
+        if self.inplace:
+            return x.mul_(self.weight.view(*shape))
+        else:
+            return x * self.weight.view(*shape)
+
+
+# ----------------------------------------------------------------------------
+class FFN(nn.Module):
+    def __init__(
+        self,
+        embed_dims=256,
+        feedforward_channels=1024,
+        num_fcs=2,
+        ffn_drop=0.0,
+        add_identity=True,
+        layer_scale_init_value=0.0,
+    ):
+        super().__init__()
+        assert num_fcs >= 2, f"num_fcs should be no less than 2. got {num_fcs}."
+        self.embed_dims = embed_dims
+        self.feedforward_channels = feedforward_channels
+        self.num_fcs = num_fcs
+
+        layers = []
+        in_channels = embed_dims
+        for _ in range(num_fcs - 1):
+            layers.append(
+                Sequential(
+                    Linear(in_channels, feedforward_channels),
+                    nn.GELU(),
+                    nn.Dropout(ffn_drop),
+                )
+            )
+            in_channels = feedforward_channels
+        layers.append(Linear(feedforward_channels, embed_dims))
+        layers.append(nn.Dropout(ffn_drop))
+        self.layers = Sequential(*layers)
+        self.dropout_layer = nn.Identity()
+        self.add_identity = add_identity
+
+        if layer_scale_init_value > 0:
+            self.gamma2 = LayerScale(embed_dims, scale=layer_scale_init_value)
+        else:
+            self.gamma2 = nn.Identity()
+
+    def forward(self, x, identity=None):
+        out = self.layers(x)
+        out = self.gamma2(out)
+        if not self.add_identity:
+            return out
+        if identity is None:
+            identity = x
+        return identity + out
+
+
+# ----------------------------------------------------------------------------
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        input_dims=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        qkv_bias=True,
+        proj_bias=True,
+        v_shortcut=False,
+    ):
+        super(MultiheadAttention, self).__init__()
+
+        self.input_dims = input_dims or embed_dims
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.v_shortcut = v_shortcut
+
+        self.head_dims = embed_dims // num_heads
+        self.scaled_dot_product_attention = F.scaled_dot_product_attention
+
+        self.qkv = nn.Linear(self.input_dims, embed_dims * 3, bias=qkv_bias)
+        self.attn_drop = attn_drop
+        self.proj = nn.Linear(embed_dims, embed_dims, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.gamma1 = nn.Identity()
+
+    def forward(self, x):
+        B, N, _ = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, self.head_dims)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn_drop = self.attn_drop if self.training else 0.0
+        x = self.scaled_dot_product_attention(q, k, v, dropout_p=attn_drop)
+        x = x.transpose(1, 2).reshape(B, N, self.embed_dims)
+
+        x = self.proj(x)
+        x = self.gamma1(self.proj_drop(x))
+
+        if self.v_shortcut:
+            x = v.squeeze(1) + x
+        return x
+
+
+# ----------------------------------------------------------------------------
+class TransformerEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        feedforward_channels,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        num_fcs=2,
+        qkv_bias=True,
+    ):
+        super(TransformerEncoderLayer, self).__init__()
+
+        self.embed_dims = embed_dims
+        self.ln1 = nn.LayerNorm(self.embed_dims, eps=1e-6, elementwise_affine=True)
+        self.attn = MultiheadAttention(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            attn_drop=attn_drop_rate,
+            proj_drop=drop_rate,
+            qkv_bias=qkv_bias,
+        )
+
+        self.ln2 = nn.LayerNorm(self.embed_dims, eps=1e-6, elementwise_affine=True)
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=feedforward_channels,
+            num_fcs=num_fcs,
+            ffn_drop=drop_rate,
+            add_identity=True,
+        )
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+    @property
+    def norm2(self):
+        return self.ln2
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = self.ffn(self.ln2(x), identity=x)
+        return x
+
+
+# ----------------------------------------------------------------------------
+@MODELS.register_module()
+class Sapiens(BaseModel):
+    arch_zoo = {
+        **dict.fromkeys(  ## this is vit-large
+            ["0.3b", "sapiens_0.3b"],
+            {
+                "embed_dims": 1024,
+                "num_layers": 24,
+                "num_heads": 16,
+                "feedforward_channels": 1024 * 4,
+            },
+        ),
+        **dict.fromkeys(  ## this is vit-huge
+            ["0.6b", "sapiens_0.6b"],
+            {
+                "embed_dims": 1280,
+                "num_layers": 32,
+                "num_heads": 16,
+                "feedforward_channels": 1280 * 4,
+            },
+        ),
+        **dict.fromkeys(  ## this is vit-g
+            ["1b", "sapiens_1b"],
+            {
+                "embed_dims": 1536,
+                "num_layers": 40,
+                "num_heads": 24,
+                "feedforward_channels": 1536 * 4,
+            },
+        ),
+        **dict.fromkeys(
+            ["2b", "sapiens_2b"],
+            {
+                "embed_dims": 1920,
+                "num_layers": 48,
+                "num_heads": 32,
+                "feedforward_channels": 1920 * 4,
+            },
+        ),
+    }
+    num_extra_tokens = 1  # class token
+    OUT_TYPES = {"raw", "cls_token", "featmap"}
+
+    def __init__(
+        self,
+        arch="base",
+        img_size=1024,
+        patch_size=16,
+        in_channels=3,
+        out_indices=-1,
+        drop_rate=0.0,
+        qkv_bias=True,
+        final_norm=True,
+        out_type="cls_token",
+        with_cls_token=True,
+        frozen_stages=-1,
+        interpolate_mode="bicubic",
+        patch_cfg=dict(),
+        layer_cfgs=dict(),
+        init_cfg=None,
+    ):
+        super(Sapiens, self).__init__(init_cfg=init_cfg)
+
+        arch = arch.lower()
+        assert arch in set(self.arch_zoo), (
+            f"Arch {arch} is not in default archs {set(self.arch_zoo)}"
+        )
+        self.arch_settings = self.arch_zoo[arch]
+
+        self.embed_dims = self.arch_settings["embed_dims"]
+        self.num_layers = self.arch_settings["num_layers"]
+        self.img_size = to_2tuple(img_size)
+        self.patch_size = patch_size
+
+        # Set patch embedding
+        _patch_cfg = dict(
+            in_channels=in_channels,
+            input_size=img_size,
+            embed_dims=self.embed_dims,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+        _patch_cfg.update(patch_cfg)
+        self.patch_embed = PatchEmbed(**_patch_cfg)
+        self.patch_resolution = self.patch_embed.init_out_size
+        num_patches = self.patch_resolution[0] * self.patch_resolution[1]
+
+        # Set out type
+        if out_type not in self.OUT_TYPES:
+            raise ValueError(
+                f"Unsupported `out_type` {out_type}, please "
+                f"choose from {self.OUT_TYPES}"
+            )
+        self.out_type = out_type
+
+        # Set cls token
+        self.with_cls_token = with_cls_token
+        if with_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims))
+        elif out_type != "cls_token":
+            self.cls_token = None
+            self.num_extra_tokens = 0
+        else:
+            raise ValueError('with_cls_token must be True when `out_type="cls_token"`.')
+
+        # Set position embedding
+        self.interpolate_mode = interpolate_mode
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + self.num_extra_tokens, self.embed_dims)
+        )
+        self.drop_after_pos = nn.Dropout(p=drop_rate)
+
+        if isinstance(out_indices, int):
+            out_indices = [out_indices]
+        assert isinstance(out_indices, Sequence), (
+            f'"out_indices" must by a sequence or int, get {type(out_indices)} instead.'
+        )
+        for i, index in enumerate(out_indices):
+            if index < 0:
+                out_indices[i] = self.num_layers + index
+            assert 0 <= out_indices[i] <= self.num_layers, (
+                f"Invalid out_indices {index}"
+            )
+        self.out_indices = out_indices
+
+        self.layers = nn.Sequential()
+        if isinstance(layer_cfgs, dict):
+            layer_cfgs = [layer_cfgs] * self.num_layers
+        for i in range(self.num_layers):
+            _layer_cfg = dict(
+                embed_dims=self.embed_dims,
+                num_heads=self.arch_settings["num_heads"],
+                feedforward_channels=self.arch_settings["feedforward_channels"],
+                drop_rate=drop_rate,
+                qkv_bias=qkv_bias,
+            )
+            _layer_cfg.update(layer_cfgs[i])
+            self.layers.append(TransformerEncoderLayer(**_layer_cfg))
+
+        self.frozen_stages = frozen_stages
+        self.pre_norm = nn.Identity()
+
+        self.final_norm = final_norm
+        if final_norm:
+            self.ln1 = nn.LayerNorm(self.embed_dims, eps=1e-6, elementwise_affine=True)
+
+        # freeze stages only when self.frozen_stages > 0
+        if self.frozen_stages > 0:
+            self._freeze_stages()
+
+        self._register_load_state_dict_pre_hook(self._prepare_pos_embed)
+
+        self.init_weights()
+
+        return
+
+    def _prepare_pos_embed(self, state_dict, prefix, *args, **kwargs):
+        name = prefix + "pos_embed"
+        if name not in state_dict.keys():
+            return
+
+        ckpt_pos_embed_shape = state_dict[name].shape
+
+        from sapiens.engine.logger import Logger
+
+        logger = Logger.get_current_instance()
+        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+
+        # Handle class token removal if needed
+        if not self.with_cls_token:
+            if ckpt_pos_embed_shape[1] == self.pos_embed.shape[1] + 1:
+                # Remove cls token from state dict if it's not used
+                state_dict[name] = state_dict[name][:, 1:]
+                ckpt_pos_embed_shape = state_dict[name].shape
+            elif ckpt_pos_embed_shape[1] % 2 == 1:
+                # Remove class token when interpolation is required
+                if rank == 0:
+                    logger.info(
+                        "Note: removing the class token from pretrained weights"
+                    )
+                state_dict[name] = state_dict[name][:, 1:]
+                ckpt_pos_embed_shape = state_dict[name].shape
+
+        # Skip if shapes already match
+        if self.pos_embed.shape == ckpt_pos_embed_shape:
+            return
+
+        if rank == 0:
+            logger.info(
+                f"Resize the pos_embed shape from {ckpt_pos_embed_shape} "
+                f"to {self.pos_embed.shape}."
+            )
+
+        # Calculate grid dimensions
+        pos_h, pos_w = self.patch_embed.init_out_size
+        assert pos_h >= pos_w  # for vertical aspect ratio or square
+
+        # Number of non-extra tokens in checkpoint
+        num_vis = ckpt_pos_embed_shape[1] - self.num_extra_tokens
+
+        # Determine original grid shape
+        side = int(math.sqrt(num_vis))
+        factor = int(math.sqrt((num_vis * self.patch_size * self.patch_size) // 12))
+
+        # Set old grid based on aspect ratio detection
+        if side * side == num_vis:
+            old_grid = (side, side)  # square grid
+        elif 4 * factor * 3 * factor == num_vis * self.patch_size * self.patch_size:
+            old_grid = (
+                (factor * 4) // self.patch_size,
+                (factor * 3) // self.patch_size,
+            )  # 4:3 ratio
+        else:
+            if rank == 0:
+                logger.warning(
+                    f"Original pos_embed tokens ({num_vis}) not square or 4:3 does not match current size"
+                )
+            state_dict[name] = self.pos_embed
+            return
+
+        # Resize position embedding
+        new_grid = (pos_h, pos_w)
+        state_dict[name] = resize_pos_embed(
+            state_dict[name],
+            old_grid,
+            new_grid,
+            mode=self.interpolate_mode,
+            num_extra_tokens=self.num_extra_tokens,
+        )
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+    @property
+    def norm2(self):
+        return self.ln2
+
+    @staticmethod
+    def resize_pos_embed(*args, **kwargs):
+        """Interface for backward-compatibility."""
+        return resize_pos_embed(*args, **kwargs)
+
+    def _freeze_stages(self):
+        # freeze position embedding
+        if self.pos_embed is not None:
+            self.pos_embed.requires_grad = False
+
+        # set dropout to eval model
+        self.drop_after_pos.eval()
+        # freeze patch embedding
+        self.patch_embed.eval()
+        for param in self.patch_embed.parameters():
+            param.requires_grad = False
+        # freeze pre-norm
+        for param in self.pre_norm.parameters():
+            param.requires_grad = False
+        # freeze cls_token
+        if self.cls_token is not None:
+            self.cls_token.requires_grad = False
+        # freeze layers
+        for i in range(1, self.frozen_stages + 1):
+            m = self.layers[i - 1]
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+        # freeze the last layer norm
+        if self.frozen_stages == len(self.layers):
+            if self.final_norm:
+                self.ln1.eval()
+                for param in self.ln1.parameters():
+                    param.requires_grad = False
+
+            if self.out_type == "avg_featmap":
+                self.ln2.eval()
+                for param in self.ln2.parameters():
+                    param.requires_grad = False
+
+    def forward(self, x):
+        B = x.shape[0]
+        x, patch_resolution = self.patch_embed(x)
+
+        if self.cls_token is not None:
+            cls_token = self.cls_token.expand(B, -1, -1)
+            x = torch.cat((cls_token, x), dim=1)
+
+        x = x + resize_pos_embed(
+            self.pos_embed,
+            self.patch_resolution,
+            patch_resolution,
+            mode=self.interpolate_mode,
+            num_extra_tokens=self.num_extra_tokens,
+        )
+        x = self.drop_after_pos(x)
+
+        x = self.pre_norm(x)  ## B x (num tokens) x embed_dim
+
+        outs = []
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+
+            if i == len(self.layers) - 1 and self.final_norm:
+                x = self.ln1(x)
+
+            if i in self.out_indices:
+                outs.append(self._format_output(x, patch_resolution))
+
+        return tuple(outs)
+
+    def _format_output(self, x, hw):
+        if self.out_type == "raw":
+            return x
+        if self.out_type == "cls_token":
+            return x[:, 0]
+
+        patch_token = x[:, self.num_extra_tokens :]
+        if self.out_type == "featmap":
+            B = x.size(0)
+            # (B, N, C) -> (B, H, W, C) -> (B, C, H, W)
+            return patch_token.reshape(B, *hw, -1).permute(0, 3, 1, 2)

sapiens/backbones/sapiens2.py

@@ -0,0 +1,916 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sapiens.engine.models.base_model import BaseModel
+from sapiens.registry import MODELS
+from torch import Tensor
+from torch.nn.init import trunc_normal_
+from torch.utils.checkpoint import checkpoint
+
+
+# ----------------------------------------------------------------------------
+def to_2tuple(x):
+    if isinstance(x, (str, bytes)):
+        return (x, x)
+    if isinstance(x, Sequence):
+        x = tuple(x)
+        if len(x) == 2:
+            return x
+        raise ValueError("Expected scalar or length-2 iterable")
+    return (x, x)
+
+
+class RopePositionEmbedding(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        *,
+        num_heads: int,
+        base: float | None = 100.0,
+        min_period: float | None = None,
+        max_period: float | None = None,
+        normalize_coords: Literal["min", "max", "separate"] = "separate",
+        shift_coords: float | None = None,
+        jitter_coords: float | None = None,
+        rescale_coords: float | None = None,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        assert embed_dim % (4 * num_heads) == 0
+        both_periods = min_period is not None and max_period is not None
+        if (base is None and not both_periods) or (base is not None and both_periods):
+            raise ValueError(
+                "Either `base` or `min_period`+`max_period` must be provided."
+            )
+
+        D_head = embed_dim // num_heads
+        self.base = base
+        self.min_period = min_period
+        self.max_period = max_period
+        self.D_head = D_head
+        self.normalize_coords = normalize_coords
+        self.shift_coords = shift_coords
+        self.jitter_coords = jitter_coords
+        self.rescale_coords = rescale_coords
+
+        # Needs persistent=True because we do teacher.load_state_dict(student.state_dict()) to initialize the teacher
+        self.dtype = dtype or torch.float32  # Don't rely on self.periods.dtype
+        self.register_buffer(
+            "periods",
+            torch.empty(D_head // 4, device=device, dtype=self.dtype),
+            persistent=True,
+        )
+        self._init_weights()
+
+    def forward(self, *, H: int, W: int) -> tuple[Tensor, Tensor]:
+        device = self.periods.device
+        dtype = self.dtype
+        dd = {"device": device, "dtype": dtype}
+        # Prepare coords in range [-1, +1]
+        if self.normalize_coords == "max":
+            max_HW = max(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / max_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / max_HW  # [W]
+        elif self.normalize_coords == "min":
+            min_HW = min(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / min_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / min_HW  # [W]
+        elif self.normalize_coords == "separate":
+            coords_h = torch.arange(0.5, H, **dd) / H  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / W  # [W]
+        else:
+            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
+        coords = torch.stack(
+            torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1
+        )  # [H, W, 2]
+        coords = coords.flatten(0, 1)  # [HW, 2]
+        coords = 2.0 * coords - 1.0  # Shift range [0, 1] to [-1, +1]
+
+        # Shift coords by adding a uniform value in [-shift, shift]
+        if self.training and self.shift_coords is not None:
+            shift_hw = torch.empty(2, **dd).uniform_(
+                -self.shift_coords, self.shift_coords
+            )
+            coords += shift_hw[None, :]
+
+        # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter]
+        if self.training and self.jitter_coords is not None:
+            jitter_max = np.log(self.jitter_coords)
+            jitter_min = -jitter_max
+            jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp()
+            coords *= jitter_hw[None, :]
+
+        # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale]
+        if self.training and self.rescale_coords is not None:
+            rescale_max = np.log(self.rescale_coords)
+            rescale_min = -rescale_max
+            rescale_hw = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp()
+            coords *= rescale_hw
+
+        # Prepare angles and sin/cos
+        angles = (
+            2 * math.pi * coords[:, :, None] / self.periods[None, None, :]
+        )  # [HW, 2, D//4]
+        angles = angles.flatten(1, 2)  # [HW, D//2]
+        angles = angles.tile(2)  # [HW, D]
+        cos = torch.cos(angles)  # [HW, D]
+        sin = torch.sin(angles)  # [HW, D]
+
+        return (sin, cos)  # 2 * [HW, D]
+
+    def _init_weights(self):
+        device = self.periods.device
+        dtype = self.dtype
+        if self.base is not None:
+            periods = self.base ** (
+                2
+                * torch.arange(self.D_head // 4, device=device, dtype=dtype)
+                / (self.D_head // 2)
+            )  # [D//4]
+        else:
+            base = self.max_period / self.min_period
+            exponents = torch.linspace(
+                0, 1, self.D_head // 4, device=device, dtype=dtype
+            )  # [D//4] range [0, 1]
+            periods = base**exponents  # range [1, max_period / min_period]
+            periods = periods / base  # range [min_period / max_period, 1]
+            periods = periods * self.max_period  # range [min_period, max_period]
+        self.periods.data = periods
+
+
+# -------------------------------------------------------------------------------
+class Tokenizer(nn.Module):
+    """Stacked window self‑attention that emits one token per window
+    by re‑using TransformerEncoderLayer blocks."""
+
+    def __init__(
+        self,
+        embed_dims: int,
+        window_size: int = 4,
+        num_heads: int = 4,
+        num_tokenizer_layers: int = 1,
+        qkv_bias: bool = True,
+        use_qk_norm: bool = False,
+        chunk_size: int = 1024,  # max windows per chunk
+    ):
+        super().__init__()
+        self.ws = window_size
+        self.chunk_size = chunk_size
+
+        # local absolute positional embeddings for [CLS] + patch tokens
+        self.local_pos_embed = nn.Parameter(
+            torch.zeros(1, 1 + window_size * window_size, embed_dims)
+        )
+        trunc_normal_(self.local_pos_embed, std=0.02)
+
+        # build N identical TransformerEncoderLayer blocks
+        self.blocks = nn.ModuleList(
+            [
+                TransformerEncoderLayer2(
+                    embed_dims=embed_dims,
+                    num_heads=num_heads,
+                    feedforward_channels=embed_dims * 4,  # standard FFN size
+                    qkv_bias=qkv_bias,
+                    use_qk_norm=use_qk_norm,
+                )
+                for _ in range(num_tokenizer_layers)
+            ]
+        )
+
+        # shared CLS token for pooling
+        self.w_cls = nn.Parameter(torch.zeros(1, 1, embed_dims))
+        trunc_normal_(self.w_cls, std=0.02)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        hw: Tuple[int, int],
+    ) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        """Args:
+           x  : B, N, C   (N = H*W)
+           hw : (H, W) before reduction
+        Returns:
+           x_ : B, (H/ws)*(W/ws), C
+           hw_: (H/ws, W/ws)
+        """
+        B, N, C = x.shape
+        H, W = hw
+        ws = self.ws
+        assert H % ws == 0 and W % ws == 0, (
+            f"Image size {H}×{W} must be divisible by window {ws}."
+        )
+
+        # reshape tokens → non‑overlapping windows
+        x = x.view(B, H, W, C)
+
+        ph, pw = H // ws, W // ws  ## ints in eager mode
+        ph, pw = int(ph), int(pw)  ## ints in scripting mode
+        x = x.view(B, ph, ws, pw, ws, C)  # B, H/ws, ws, W/ws, ws, C
+        x = x.permute(0, 1, 3, 2, 4, 5)  # B, H/ws, W/ws, ws, ws, C
+        x = x.contiguous().view(B * ph * pw, ws * ws, C)  # (B*H/ws*W/ws), ws², C))
+
+        total_windows = x.size(0)
+        chunk_size = int(min(self.chunk_size, total_windows))
+        token_out = x.new_empty(total_windows, C)
+
+        use_ckpt = self.training and torch.is_grad_enabled()
+
+        def _run_blocks(t: torch.Tensor) -> torch.Tensor:
+            for blk in self.blocks:
+                t = blk(t)
+            return t
+
+        for i in range(0, total_windows, chunk_size):
+            chunk = x[i : i + chunk_size]  # (m, ws², C)
+            m = chunk.size(0)
+            cls = self.w_cls.expand(m, -1, -1)  # (m, 1, C)
+            chunk = torch.cat([cls, chunk], dim=1)  # (m, 1+ws², C)
+            chunk = chunk + self.local_pos_embed  # add local PE
+
+            if use_ckpt:
+                chunk = checkpoint(_run_blocks, chunk, use_reentrant=False)
+            else:
+                chunk = _run_blocks(chunk)
+
+            token_out[i : i + m] = chunk[:, 0]  # take CLS out
+
+        token = token_out.view(B, ph * pw, C)  # (B, (H/ws)*(W
+        return token, (ph, pw)
+
+
+# -------------------------------------------------------------------------------
+class GroupedQueryAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        num_kv_heads=None,
+        input_dims=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        qkv_bias=True,
+        qk_scale=None,
+        proj_bias=True,
+        use_qk_norm=True,
+        v_shortcut=False,
+        layer_scale_init_value=0.0,
+    ):
+        super().__init__()
+        # Core dims
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        assert self.num_heads % self.num_kv_heads == 0, (
+            "num_kv_heads must divide num_heads"
+        )
+        self.head_dim = embed_dims // num_heads
+        self.input_dims = input_dims or embed_dims
+        # Features
+        self.attn_drop = attn_drop
+        self.v_shortcut = v_shortcut
+        self.use_qk_norm = use_qk_norm
+
+        # Attention operation selection
+        if qk_scale is not None:
+            scale = qk_scale
+        else:
+            scale = self.head_dim**-0.5
+
+        assert qk_scale is None, "qk_scale is not supported"
+        self.attn_op = F.scaled_dot_product_attention
+
+        # Q/K/V projections
+        self.wq = nn.Linear(self.input_dims, embed_dims, bias=qkv_bias)
+        self.wk = nn.Linear(
+            self.input_dims, self.num_kv_heads * self.head_dim, bias=qkv_bias
+        )
+        self.wv = nn.Linear(
+            self.input_dims, self.num_kv_heads * self.head_dim, bias=qkv_bias
+        )
+
+        if self.use_qk_norm:
+            self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6)
+            self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6)
+
+        # Output projection + dropout
+        self.proj = nn.Linear(embed_dims, embed_dims, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        # Optional LayerScale
+        if layer_scale_init_value > 0:
+            self.gamma = LayerScale(embed_dims, scale=layer_scale_init_value)
+        else:
+            self.gamma = nn.Identity()
+
+    def apply_rope(
+        self, q: Tensor, k: Tensor, rope: Tensor | Tuple[Tensor, Tensor]
+    ) -> Tuple[Tensor, Tensor]:
+        # All operations will use the dtype of rope, the output is cast back to the dtype of q and k
+        q_dtype = q.dtype
+        k_dtype = k.dtype
+        sin, cos = rope
+        rope_dtype = sin.dtype
+        q = q.to(dtype=rope_dtype)
+        k = k.to(dtype=rope_dtype)
+        N = q.shape[-2]
+        prefix = N - sin.shape[-2]  ## extra tokens
+        assert prefix >= 0
+        q_prefix = q[:, :, :prefix, :]
+        q = self._rope_apply(q[:, :, prefix:, :], sin, cos)  # [B, head, hw, D//head]
+        q = torch.cat((q_prefix, q), dim=-2)  # [B, head, N, D//head]
+        k_prefix = k[:, :, :prefix, :]
+        k = self._rope_apply(k[:, :, prefix:, :], sin, cos)  # [B, head, hw, D//head]
+        k = torch.cat((k_prefix, k), dim=-2)  # [B, head, N, D//head]
+        q = q.to(dtype=q_dtype)
+        k = k.to(dtype=k_dtype)
+        return q, k
+
+    def _rope_rotate_half(self, x: Tensor) -> Tensor:
+        # x:   [ x0  x1  x2  x3  x4  x5]
+        # out: [-x3 -x4 -x5  x0  x1  x2]
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat([-x2, x1], dim=-1)
+
+    def _rope_apply(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+        # x:   [..., D], eg [x0,     x1,   x2,   x3,   x4,   x5]
+        # sin: [..., D], eg [sin0, sin1, sin2, sin0, sin1, sin2]
+        # cos: [..., D], eg [cos0, cos1, cos2, cos0, cos1, cos2]
+        return (x * cos) + (self._rope_rotate_half(x) * sin)
+
+    def forward(self, x, rope=None):
+        B, N, _ = x.shape
+        # Q: (B, N, num_heads, head_dim)
+        q = self.wq(x).view(B, N, self.num_heads, self.head_dim)
+        # K/V: (B, N, num_kv_heads, head_dim)
+        k = self.wk(x).view(B, N, self.num_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, N, self.num_kv_heads, self.head_dim)
+
+        # (B, heads, N, head_dim)
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        # Repeat KV heads if group ratio >1
+        if self.num_kv_heads != self.num_heads:
+            factor = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(factor, dim=1)
+            v = v.repeat_interleave(factor, dim=1)
+
+        if rope is not None:
+            q, k = self.apply_rope(q, k, rope)
+
+        # Scaled dot-product attention
+        attn_out = self.attn_op(
+            q, k, v, dropout_p=self.attn_drop if self.training else 0.0
+        )  # (B, num_heads, N, head_dim)
+
+        # Merge heads -> (B, N, embed_dims)
+        out = attn_out.permute(0, 2, 1, 3).reshape(B, N, self.embed_dims)
+
+        # Output projection + drop + layer scale
+        out = self.proj(out)
+        out = self.gamma(self.proj_drop(out))
+
+        # Optional V-shortcut (only when MQA)
+        if self.v_shortcut and self.num_kv_heads == 1:
+            raise NotImplementedError
+        return out
+
+
+# -------------------------------------------------------------------------------
+class TransformerEncoderLayer2(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        num_kv_heads=None,
+        feedforward_channels=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        layer_scale_init_value=0.0,
+        use_qk_norm=True,
+        qkv_bias=True,
+    ):
+        super(TransformerEncoderLayer2, self).__init__()
+
+        self.embed_dims = embed_dims
+        self.ln1 = nn.RMSNorm(self.embed_dims, eps=1e-6)
+        self.attn = GroupedQueryAttention(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            attn_drop=attn_drop_rate,
+            proj_drop=drop_rate,
+            qkv_bias=qkv_bias,
+            layer_scale_init_value=layer_scale_init_value,
+            use_qk_norm=use_qk_norm,
+        )
+
+        self.ln2 = nn.RMSNorm(self.embed_dims, eps=1e-6)
+        self.ffn = SwiGLUFFN(
+            embed_dims=embed_dims,
+            feedforward_channels=feedforward_channels,
+        )
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+    @property
+    def norm2(self):
+        return self.ln2
+
+    def forward(self, x, rope=None):
+        x = x + self.attn(self.ln1(x), rope=rope)
+        x = self.ffn(self.ln2(x), identity=x)
+        return x
+
+
+##-----------------------------------
+@MODELS.register_module()
+class Sapiens2(BaseModel):
+    arch_zoo = {
+        **dict.fromkeys(
+            ["sapiens2_0.1b"],
+            {
+                "embed_dims": 768,
+                "num_layers": 12,
+                "num_heads": 12,
+                "feedforward_channels": 768 * 4,
+                "num_tokenizer_layers": 2,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_0.4b"],
+            {
+                "embed_dims": 1024,
+                "num_layers": 24,
+                "num_heads": 16,
+                "feedforward_channels": 1024 * 4,
+                "num_tokenizer_layers": 2,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_0.8b"],
+            {
+                "embed_dims": 1280,
+                "num_layers": 32,
+                "num_heads": 16,
+                "feedforward_channels": 1280 * 4,
+                "num_tokenizer_layers": 3,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_1b"],
+            {
+                "embed_dims": 1536,
+                "num_layers": 40,
+                "num_heads": 24,
+                "feedforward_channels": 1536 * 4,
+                "num_tokenizer_layers": 4,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_5b"],
+            {
+                "embed_dims": 2432,
+                "num_layers": 56,
+                "num_heads": 32,
+                "feedforward_channels": 2432 * 4,
+                "num_tokenizer_layers": 6,
+            },
+        ),
+    }
+
+    num_extra_tokens = 1  # class token
+    OUT_TYPES = {"raw", "cls_token", "featmap"}
+
+    def __init__(
+        self,
+        arch="sapiens2_1b",
+        img_size=(1024, 768),
+        patch_size=16,
+        in_channels=3,
+        out_indices=-1,
+        drop_rate=0.0,
+        window_size=4,
+        use_tokenizer=False,  ## 4k resolution
+        use_qk_norm=True,
+        qkv_bias=True,
+        final_norm=True,
+        out_type="raw",
+        with_cls_token=True,
+        layer_scale_init_value=1e-4,  ## non zero init to activate layerscale
+        frozen_stages=-1,
+        patch_cfg=dict(),
+        layer_cfgs=dict(),
+        pos_embed_rope_base: float = 100.0,
+        pos_embed_rope_min_period: float | None = None,
+        pos_embed_rope_max_period: float | None = None,
+        pos_embed_rope_normalize_coords: Literal["min", "max", "separate"] = "separate",
+        pos_embed_rope_shift_coords: float | None = None,
+        pos_embed_rope_jitter_coords: float | None = None,
+        pos_embed_rope_rescale_coords: float | None = None,
+        pos_embed_rope_dtype: str = "bf16",
+        n_storage_tokens: int = 8,
+        init_cfg=None,
+    ):
+        super(Sapiens2, self).__init__(init_cfg=init_cfg)
+
+        arch = arch.lower()
+        assert arch in set(self.arch_zoo), (
+            f"Arch {arch} is not in default archs {set(self.arch_zoo)}"
+        )
+        self.arch_settings = self.arch_zoo[arch]
+
+        self.embed_dims = self.arch_settings["embed_dims"]
+        self.num_layers = self.arch_settings["num_layers"]
+        self.patch_size = patch_size
+
+        self.window_size = window_size
+        img_size = to_2tuple(img_size)
+        encoder_img_size = (
+            (img_size[0] // window_size, img_size[1] // window_size)
+            if use_tokenizer
+            else img_size
+        )
+        self.img_size = to_2tuple(encoder_img_size)
+
+        # Set patch embedding
+        _patch_cfg = dict(
+            in_channels=in_channels,
+            input_size=self.img_size,
+            embed_dims=self.embed_dims,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+        _patch_cfg.update(patch_cfg)
+        self.patch_embed = PatchEmbed(**_patch_cfg)
+        self.patch_resolution = self.patch_embed.init_out_size
+        num_patches = self.patch_resolution[0] * self.patch_resolution[1]
+
+        self.rope_embed = RopePositionEmbedding(
+            embed_dim=self.embed_dims,
+            num_heads=self.arch_settings["num_heads"],
+            base=pos_embed_rope_base,
+            min_period=pos_embed_rope_min_period,
+            max_period=pos_embed_rope_max_period,
+            normalize_coords=pos_embed_rope_normalize_coords,
+            shift_coords=pos_embed_rope_shift_coords,
+            jitter_coords=pos_embed_rope_jitter_coords,
+            rescale_coords=pos_embed_rope_rescale_coords,
+            dtype=torch.bfloat16 if pos_embed_rope_dtype == "bf16" else torch.float32,
+        )
+
+        # Set out type
+        if out_type not in self.OUT_TYPES:
+            raise ValueError(
+                f"Unsupported `out_type` {out_type}, please "
+                f"choose from {self.OUT_TYPES}"
+            )
+        self.out_type = out_type
+
+        if use_tokenizer == True:
+            self.tokenizer = Tokenizer(
+                embed_dims=self.embed_dims,
+                window_size=self.window_size,
+                num_heads=self.arch_settings["num_heads"],
+                num_tokenizer_layers=self.arch_settings["num_tokenizer_layers"],
+                qkv_bias=True,
+                use_qk_norm=False,
+            )
+        else:
+            self.tokenizer = None
+
+        # Set cls + storage tokens
+        self.with_cls_token = with_cls_token
+        if with_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims))
+        elif out_type != "cls_token":
+            self.cls_token = None
+            self.num_extra_tokens = 0
+        else:
+            raise ValueError('with_cls_token must be True when `out_type="cls_token"`.')
+
+        ## registers
+        self.n_storage_tokens = int(n_storage_tokens)
+        self.storage_tokens = (
+            nn.Parameter(torch.zeros(1, self.n_storage_tokens, self.embed_dims))
+            if self.n_storage_tokens > 0
+            else None
+        )
+        # how many non-patch tokens are at the front
+        self.num_extra_tokens = (
+            1 if self.cls_token is not None else 0
+        ) + self.n_storage_tokens
+
+        if isinstance(out_indices, int):
+            out_indices = [out_indices]
+        assert isinstance(out_indices, Sequence), (
+            f'"out_indices" must by a sequence or int, get {type(out_indices)} instead.'
+        )
+        for i, index in enumerate(out_indices):
+            if index < 0:
+                out_indices[i] = self.num_layers + index
+            assert 0 <= out_indices[i] <= self.num_layers, (
+                f"Invalid out_indices {index}"
+            )
+        self.out_indices = out_indices
+
+        self.blocks = nn.Sequential()
+        if isinstance(layer_cfgs, dict):
+            layer_cfgs = [layer_cfgs] * self.num_layers
+
+        mhsa_early, mhsa_late = 8, 8
+        for i in range(self.num_layers):
+            if i < mhsa_early or i >= self.num_layers - mhsa_late:
+                num_kv_heads = None  ## use MHSA
+            else:
+                num_kv_heads = self.arch_settings["num_heads"] // 2  # Use GQA
+
+            _layer_cfg = dict(
+                embed_dims=self.embed_dims,
+                num_heads=self.arch_settings["num_heads"],
+                num_kv_heads=num_kv_heads,
+                feedforward_channels=self.arch_settings["feedforward_channels"],
+                use_qk_norm=use_qk_norm,
+                layer_scale_init_value=layer_scale_init_value,
+                drop_rate=drop_rate,
+                qkv_bias=qkv_bias,
+            )
+            _layer_cfg.update(layer_cfgs[i])
+            self.blocks.append(TransformerEncoderLayer2(**_layer_cfg))
+
+        self.frozen_stages = frozen_stages
+
+        self.final_norm = final_norm
+        if final_norm:
+            self.ln1 = nn.RMSNorm(self.embed_dims, eps=1e-6)
+
+        # freeze stages only when self.frozen_stages > 0
+        if self.frozen_stages > 0:
+            self._freeze_stages()
+
+        ## load init weights
+        self.init_weights()
+
+        return
+
+    def init_weights(self):
+        if self.init_cfg is not None:
+            super(Sapiens2, self).init_weights()
+            return
+
+        # Initialize class token and storagr token embeddings
+        if self.with_cls_token:
+            trunc_normal_(self.cls_token, std=0.02)
+
+        if self.storage_tokens is not None:
+            trunc_normal_(self.storage_tokens, std=0.02)
+
+        # Apply custom initialization to all submodules
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # Use a truncated normal distribution for linear layer weights
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+        elif isinstance(m, (nn.LayerNorm, nn.RMSNorm)):
+            # Initialize normalization layers to act as an identity function
+            if hasattr(m, "bias") and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+            if hasattr(m, "weight") and m.weight is not None:
+                nn.init.constant_(m.weight, 1.0)
+
+        elif isinstance(m, nn.Conv2d):
+            # Initialize conv layer weights like linear layers
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def _freeze_stages(self):
+        ## freeze tokenizer
+        if self.frozen_stages >= 1 and self.tokenizer is not None:
+            self.tokenizer.eval()
+            for param in self.tokenizer.parameters():
+                param.requires_grad = False
+
+        # freeze patch embedding
+        self.patch_embed.eval()
+        for param in self.patch_embed.parameters():
+            param.requires_grad = False
+        # freeze cls_token
+        if self.cls_token is not None:
+            self.cls_token.requires_grad = False
+        if self.storage_tokens is not None:
+            self.storage_tokens.requires_grad = False
+        # freeze layers
+        for i in range(1, self.frozen_stages + 1):
+            m = self.blocks[i - 1]
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+
+        # freeze the last layer norm
+        if self.frozen_stages == len(self.blocks):
+            if self.final_norm:
+                self.ln1.eval()
+                for param in self.ln1.parameters():
+                    param.requires_grad = False
+
+    def forward(self, x):
+        B = x.shape[0]
+
+        x, patch_resolution = self.patch_embed(x)  # (B, 256*256, C)
+        if self.tokenizer is not None:
+            x, patch_resolution = self.tokenizer(x, patch_resolution)
+
+        # prepend [CLS] and storage tokens
+        prepend = []
+        if self.cls_token is not None:
+            prepend.append(self.cls_token.expand(B, -1, -1))
+        if self.storage_tokens is not None:
+            prepend.append(self.storage_tokens.expand(B, -1, -1))
+        if len(prepend) > 0:
+            x = torch.cat(prepend + [x], dim=1)
+
+        rope_sincos = self.rope_embed(H=patch_resolution[0], W=patch_resolution[1])
+        outs = []
+        for i, layer in enumerate(self.blocks):
+            x = layer(x, rope=rope_sincos)
+
+            if i == len(self.blocks) - 1 and self.final_norm:
+                x = self.ln1(x)
+
+            if i in self.out_indices:
+                outs.append(self._format_output(x, patch_resolution))
+
+        return tuple(outs)
+
+    def _format_output(self, x, hw):
+        if self.out_type == "raw":
+            return x
+        if self.out_type == "cls_token":
+            return x[:, 0]
+
+        patch_token = x[:, self.num_extra_tokens :]
+        if self.out_type == "featmap":
+            B = x.size(0)
+            # (B, N, C) -> (B, H, W, C) -> (B, C, H, W)
+            return patch_token.reshape(B, *hw, -1).permute(0, 3, 1, 2)
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+
+# ----------------------------------------------------------------------------
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inplace: bool = False,
+        data_format: str = "channels_last",
+        scale: float = 1e-5,
+    ):
+        super().__init__()
+        assert data_format in (
+            "channels_last",
+            "channels_first",
+        ), "'data_format' could only be channels_last or channels_first."
+        self.inplace = inplace
+        self.data_format = data_format
+        self.weight = nn.Parameter(torch.ones(dim) * scale)
+
+    def forward(self, x) -> torch.Tensor:
+        if self.data_format == "channels_first":
+            shape = tuple((1, -1, *(1 for _ in range(x.dim() - 2))))
+        else:
+            shape = tuple((*(1 for _ in range(x.dim() - 1)), -1))
+        if self.inplace:
+            return x.mul_(self.weight.view(*shape))
+        else:
+            return x * self.weight.view(*shape)
+
+
+# ----------------------------------------------------------------------------
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        embed_dims=768,
+        kernel_size=16,
+        stride=16,
+        padding="corner",
+        dilation=1,
+        bias=True,
+        input_size=None,
+    ):
+        super().__init__()
+
+        self.embed_dims = embed_dims
+        if stride is None:
+            stride = kernel_size
+
+        kernel_size = to_2tuple(kernel_size)
+        stride = to_2tuple(stride)
+        dilation = to_2tuple(dilation)
+        padding = 0
+        padding = to_2tuple(padding)
+
+        self.projection = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=embed_dims,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+        if input_size:
+            input_size = to_2tuple(input_size)
+            self.init_input_size = input_size
+            h_out = (
+                input_size[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1
+            ) // stride[0] + 1
+            w_out = (
+                input_size[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1
+            ) // stride[1] + 1
+            self.init_out_size = (h_out, w_out)
+        else:
+            self.init_input_size = None
+            self.init_out_size = None
+
+    def forward(self, x):
+        x = self.projection(x)
+        out_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)
+        return x, out_size
+
+
+# ----------------------------------------------------------------------------
+class SwiGLUFFN(nn.Module):
+    """SwiGLU FFN layer.
+    https://github.com/facebookresearch/dinov2/blob/main/dinov2/layers/swiglu_ffn.py
+    """  # noqa
+
+    def __init__(
+        self,
+        embed_dims: int,
+        feedforward_channels: Optional[int] = None,
+        out_dims: Optional[int] = None,
+        layer_scale_init_value: float = 0.0,
+        bias: bool = True,
+        add_identity: bool = True,
+    ) -> None:
+        super().__init__()
+        self.embed_dims = embed_dims
+        self.out_dims = out_dims or embed_dims
+        hidden_dims = feedforward_channels or embed_dims
+
+        self.w12 = nn.Linear(self.embed_dims, 2 * hidden_dims, bias=bias)
+        self.w3 = nn.Linear(hidden_dims, self.out_dims, bias=bias)
+
+        if layer_scale_init_value > 0:
+            self.gamma2 = LayerScale(dim=embed_dims, scale=layer_scale_init_value)
+        else:
+            self.gamma2 = nn.Identity()
+
+        self.add_identity = add_identity
+
+    def forward(
+        self, x: torch.Tensor, identity: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        out = self.w3(hidden)
+        out = self.gamma2(out)
+
+        if self.out_dims != self.embed_dims or not self.add_identity:
+            # due to the dimension inconsistence or user setting
+            # not to apply residual operation
+            return out
+
+        if identity is None:
+            identity = x
+        return identity + out

sapiens/backbones/standalone/sapiens.py

@@ -0,0 +1,648 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Linear, Sequential
+
+
+# ----------------------------------------------------------------------------
+def to_2tuple(x):
+    if isinstance(x, (str, bytes)):
+        return (x, x)
+    if isinstance(x, Sequence):
+        x = tuple(x)
+        if len(x) == 2:
+            return x
+        raise ValueError("Expected scalar or length-2 iterable")
+    return (x, x)
+
+
+def resize_pos_embed(
+    pos_embed, src_shape, dst_shape, mode="bicubic", num_extra_tokens=1
+):
+    if src_shape[0] == dst_shape[0] and src_shape[1] == dst_shape[1]:
+        return pos_embed
+    assert pos_embed.ndim == 3, "shape of pos_embed must be [1, L, C]"
+    _, L, C = pos_embed.shape
+    src_h, src_w = src_shape
+    assert L == src_h * src_w + num_extra_tokens, (
+        f"The length of `pos_embed` ({L}) doesn't match the expected "
+        f"shape ({src_h}*{src_w}+{num_extra_tokens}). Please check the"
+        "`img_size` argument."
+    )
+    extra_tokens = pos_embed[:, :num_extra_tokens]
+
+    src_weight = pos_embed[:, num_extra_tokens:]
+    src_weight = src_weight.reshape(1, src_h, src_w, C).permute(0, 3, 1, 2)
+
+    # The cubic interpolate algorithm only accepts float32
+    dst_weight = F.interpolate(
+        src_weight.float(), size=dst_shape, align_corners=False, mode=mode
+    )
+    dst_weight = torch.flatten(dst_weight, 2).transpose(1, 2)
+    dst_weight = dst_weight.to(src_weight.dtype)
+
+    return torch.cat((extra_tokens, dst_weight), dim=1)
+
+
+# ----------------------------------------------------------------------------
+class AdaptivePadding(nn.Module):
+    def __init__(self, kernel_size=1, stride=1, dilation=1, padding="corner"):
+        super().__init__()
+        assert padding in ("same", "corner")
+
+        kernel_size = to_2tuple(kernel_size)
+        stride = to_2tuple(stride)
+        dilation = to_2tuple(dilation)
+
+        self.padding = padding
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+
+    def get_pad_shape(self, input_shape):
+        input_h, input_w = input_shape
+        kernel_h, kernel_w = self.kernel_size
+        stride_h, stride_w = self.stride
+        output_h = math.ceil(input_h / stride_h)
+        output_w = math.ceil(input_w / stride_w)
+        pad_h = max(
+            (output_h - 1) * stride_h + (kernel_h - 1) * self.dilation[0] + 1 - input_h,
+            0,
+        )
+        pad_w = max(
+            (output_w - 1) * stride_w + (kernel_w - 1) * self.dilation[1] + 1 - input_w,
+            0,
+        )
+        return pad_h, pad_w
+
+    def forward(self, x):
+        pad_h, pad_w = self.get_pad_shape(x.size()[-2:])
+        if pad_h > 0 or pad_w > 0:
+            if self.padding == "corner":
+                x = F.pad(x, [0, pad_w, 0, pad_h])
+            elif self.padding == "same":
+                x = F.pad(
+                    x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]
+                )
+        return x
+
+
+# ----------------------------------------------------------------------------
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        embed_dims=768,
+        kernel_size=16,
+        stride=16,
+        padding="corner",
+        dilation=1,
+        bias=True,
+        input_size=None,
+    ):
+        super().__init__()
+
+        self.embed_dims = embed_dims
+        if stride is None:
+            stride = kernel_size
+
+        kernel_size = to_2tuple(kernel_size)
+        stride = to_2tuple(stride)
+        dilation = to_2tuple(dilation)
+
+        if isinstance(padding, str):
+            self.adaptive_padding = AdaptivePadding(
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                padding=padding,
+            )
+            padding = 0
+        else:
+            self.adaptive_padding = None
+        padding = to_2tuple(padding)
+
+        self.projection = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=embed_dims,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+        if input_size:
+            input_size = to_2tuple(input_size)
+            self.init_input_size = input_size
+            if self.adaptive_padding:
+                pad_h, pad_w = self.adaptive_padding.get_pad_shape(input_size)
+                input_h, input_w = input_size
+                input_h = input_h + pad_h
+                input_w = input_w + pad_w
+                input_size = (input_h, input_w)
+
+            h_out = (
+                input_size[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1
+            ) // stride[0] + 1
+            w_out = (
+                input_size[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1
+            ) // stride[1] + 1
+            self.init_out_size = (h_out, w_out)
+        else:
+            self.init_input_size = None
+            self.init_out_size = None
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]:
+        if self.adaptive_padding:
+            x = self.adaptive_padding(x)
+
+        x = self.projection(x)
+        out_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)
+        return x, out_size
+
+
+# ----------------------------------------------------------------------------
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inplace: bool = False,
+        data_format: str = "channels_last",
+        scale: float = 1e-5,
+    ):
+        super().__init__()
+        assert data_format in (
+            "channels_last",
+            "channels_first",
+        ), "'data_format' could only be channels_last or channels_first."
+        self.inplace = inplace
+        self.data_format = data_format
+        self.weight = nn.Parameter(torch.ones(dim) * scale)
+
+    def forward(self, x) -> torch.Tensor:
+        if self.data_format == "channels_first":
+            shape = tuple((1, -1, *(1 for _ in range(x.dim() - 2))))
+        else:
+            shape = tuple((*(1 for _ in range(x.dim() - 1)), -1))
+        if self.inplace:
+            return x.mul_(self.weight.view(*shape))
+        else:
+            return x * self.weight.view(*shape)
+
+
+# ----------------------------------------------------------------------------
+class FFN(nn.Module):
+    def __init__(
+        self,
+        embed_dims=256,
+        feedforward_channels=1024,
+        num_fcs=2,
+        ffn_drop=0.0,
+        add_identity=True,
+        layer_scale_init_value=0.0,
+    ):
+        super().__init__()
+        assert num_fcs >= 2, f"num_fcs should be no less than 2. got {num_fcs}."
+        self.embed_dims = embed_dims
+        self.feedforward_channels = feedforward_channels
+        self.num_fcs = num_fcs
+
+        layers = []
+        in_channels = embed_dims
+        for _ in range(num_fcs - 1):
+            layers.append(
+                Sequential(
+                    Linear(in_channels, feedforward_channels),
+                    nn.GELU(),
+                    nn.Dropout(ffn_drop),
+                )
+            )
+            in_channels = feedforward_channels
+        layers.append(Linear(feedforward_channels, embed_dims))
+        layers.append(nn.Dropout(ffn_drop))
+        self.layers = Sequential(*layers)
+        self.dropout_layer = nn.Identity()
+        self.add_identity = add_identity
+
+        if layer_scale_init_value > 0:
+            self.gamma2 = LayerScale(embed_dims, scale=layer_scale_init_value)
+        else:
+            self.gamma2 = nn.Identity()
+
+    def forward(self, x, identity=None):
+        out = self.layers(x)
+        out = self.gamma2(out)
+        if not self.add_identity:
+            return out
+        if identity is None:
+            identity = x
+        return identity + out
+
+
+# ----------------------------------------------------------------------------
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        input_dims=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        qkv_bias=True,
+        proj_bias=True,
+        v_shortcut=False,
+    ):
+        super(MultiheadAttention, self).__init__()
+
+        self.input_dims = input_dims or embed_dims
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.v_shortcut = v_shortcut
+
+        self.head_dims = embed_dims // num_heads
+        self.scaled_dot_product_attention = F.scaled_dot_product_attention
+
+        self.qkv = nn.Linear(self.input_dims, embed_dims * 3, bias=qkv_bias)
+        self.attn_drop = attn_drop
+        self.proj = nn.Linear(embed_dims, embed_dims, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.gamma1 = nn.Identity()
+
+    def forward(self, x):
+        B, N, _ = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, self.head_dims)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn_drop = self.attn_drop if self.training else 0.0
+        x = self.scaled_dot_product_attention(q, k, v, dropout_p=attn_drop)
+        x = x.transpose(1, 2).reshape(B, N, self.embed_dims)
+
+        x = self.proj(x)
+        x = self.gamma1(self.proj_drop(x))
+
+        if self.v_shortcut:
+            x = v.squeeze(1) + x
+        return x
+
+
+# ----------------------------------------------------------------------------
+class TransformerEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        feedforward_channels,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        num_fcs=2,
+        qkv_bias=True,
+    ):
+        super(TransformerEncoderLayer, self).__init__()
+
+        self.embed_dims = embed_dims
+        self.ln1 = nn.LayerNorm(self.embed_dims, eps=1e-6, elementwise_affine=True)
+        self.attn = MultiheadAttention(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            attn_drop=attn_drop_rate,
+            proj_drop=drop_rate,
+            qkv_bias=qkv_bias,
+        )
+
+        self.ln2 = nn.LayerNorm(self.embed_dims, eps=1e-6, elementwise_affine=True)
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=feedforward_channels,
+            num_fcs=num_fcs,
+            ffn_drop=drop_rate,
+            add_identity=True,
+        )
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+    @property
+    def norm2(self):
+        return self.ln2
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = self.ffn(self.ln2(x), identity=x)
+        return x
+
+
+# ----------------------------------------------------------------------------
+class Sapiens(nn.Module):
+    arch_zoo = {
+        **dict.fromkeys(  ## this is vit-large
+            ["0.3b", "sapiens_0.3b"],
+            {
+                "embed_dims": 1024,
+                "num_layers": 24,
+                "num_heads": 16,
+                "feedforward_channels": 1024 * 4,
+            },
+        ),
+        **dict.fromkeys(  ## this is vit-huge
+            ["0.6b", "sapiens_0.6b"],
+            {
+                "embed_dims": 1280,
+                "num_layers": 32,
+                "num_heads": 16,
+                "feedforward_channels": 1280 * 4,
+            },
+        ),
+        **dict.fromkeys(  ## this is vit-g
+            ["1b", "sapiens_1b"],
+            {
+                "embed_dims": 1536,
+                "num_layers": 40,
+                "num_heads": 24,
+                "feedforward_channels": 1536 * 4,
+            },
+        ),
+        **dict.fromkeys(
+            ["2b", "sapiens_2b"],
+            {
+                "embed_dims": 1920,
+                "num_layers": 48,
+                "num_heads": 32,
+                "feedforward_channels": 1920 * 4,
+            },
+        ),
+    }
+    num_extra_tokens = 1  # class token
+    OUT_TYPES = {"raw", "cls_token", "featmap", "avg_featmap"}
+
+    def __init__(
+        self,
+        arch="base",
+        img_size=224,
+        patch_size=16,
+        in_channels=3,
+        out_indices=-1,
+        drop_rate=0.0,
+        qkv_bias=True,
+        final_norm=True,
+        out_type="cls_token",
+        with_cls_token=True,
+        frozen_stages=-1,
+        interpolate_mode="bicubic",
+        patch_cfg=dict(),
+        layer_cfgs=dict(),
+    ):
+        super(Sapiens, self).__init__()
+
+        arch = arch.lower()
+        assert arch in set(self.arch_zoo), (
+            f"Arch {arch} is not in default archs {set(self.arch_zoo)}"
+        )
+        self.arch_settings = self.arch_zoo[arch]
+
+        self.embed_dims = self.arch_settings["embed_dims"]
+        self.num_layers = self.arch_settings["num_layers"]
+        self.img_size = to_2tuple(img_size)
+        self.patch_size = patch_size
+
+        # Set patch embedding
+        _patch_cfg = dict(
+            in_channels=in_channels,
+            input_size=img_size,
+            embed_dims=self.embed_dims,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+        _patch_cfg.update(patch_cfg)
+        self.patch_embed = PatchEmbed(**_patch_cfg)
+        self.patch_resolution = self.patch_embed.init_out_size
+        num_patches = self.patch_resolution[0] * self.patch_resolution[1]
+
+        # Set out type
+        if out_type not in self.OUT_TYPES:
+            raise ValueError(
+                f"Unsupported `out_type` {out_type}, please "
+                f"choose from {self.OUT_TYPES}"
+            )
+        self.out_type = out_type
+
+        # Set cls token
+        self.with_cls_token = with_cls_token
+        if with_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims))
+        elif out_type != "cls_token":
+            self.cls_token = None
+            self.num_extra_tokens = 0
+        else:
+            raise ValueError('with_cls_token must be True when `out_type="cls_token"`.')
+
+        # Set position embedding
+        self.interpolate_mode = interpolate_mode
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + self.num_extra_tokens, self.embed_dims)
+        )
+        self.drop_after_pos = nn.Dropout(p=drop_rate)
+
+        if isinstance(out_indices, int):
+            out_indices = [out_indices]
+        assert isinstance(out_indices, Sequence), (
+            f'"out_indices" must by a sequence or int, get {type(out_indices)} instead.'
+        )
+        for i, index in enumerate(out_indices):
+            if index < 0:
+                out_indices[i] = self.num_layers + index
+            assert 0 <= out_indices[i] <= self.num_layers, (
+                f"Invalid out_indices {index}"
+            )
+        self.out_indices = out_indices
+
+        self.layers = nn.Sequential()
+        if isinstance(layer_cfgs, dict):
+            layer_cfgs = [layer_cfgs] * self.num_layers
+        for i in range(self.num_layers):
+            _layer_cfg = dict(
+                embed_dims=self.embed_dims,
+                num_heads=self.arch_settings["num_heads"],
+                feedforward_channels=self.arch_settings["feedforward_channels"],
+                drop_rate=drop_rate,
+                qkv_bias=qkv_bias,
+            )
+            _layer_cfg.update(layer_cfgs[i])
+            self.layers.append(TransformerEncoderLayer(**_layer_cfg))
+
+        self.frozen_stages = frozen_stages
+        self.pre_norm = nn.Identity()
+
+        self.final_norm = final_norm
+        if final_norm:
+            self.ln1 = nn.LayerNorm(self.embed_dims, eps=1e-6, elementwise_affine=True)
+
+        # freeze stages only when self.frozen_stages > 0
+        if self.frozen_stages > 0:
+            self._freeze_stages()
+
+        self._register_load_state_dict_pre_hook(self._prepare_pos_embed)
+
+        return
+
+    def _prepare_pos_embed(self, state_dict, prefix, *args, **kwargs):
+        name = prefix + "pos_embed"
+        if name not in state_dict.keys():
+            return
+
+        ckpt_pos_embed_shape = state_dict[name].shape
+
+        # Handle class token removal if needed
+        if not self.with_cls_token:
+            if ckpt_pos_embed_shape[1] == self.pos_embed.shape[1] + 1:
+                # Remove cls token from state dict if it's not used
+                state_dict[name] = state_dict[name][:, 1:]
+                ckpt_pos_embed_shape = state_dict[name].shape
+            elif ckpt_pos_embed_shape[1] % 2 == 1:
+                # Remove class token when interpolation is required
+                state_dict[name] = state_dict[name][:, 1:]
+                ckpt_pos_embed_shape = state_dict[name].shape
+
+        # Skip if shapes already match
+        if self.pos_embed.shape == ckpt_pos_embed_shape:
+            return
+
+        # Calculate grid dimensions
+        pos_h, pos_w = self.patch_embed.init_out_size
+        assert pos_h >= pos_w  # for vertical aspect ratio or square
+
+        # Number of non-extra tokens in checkpoint
+        num_vis = ckpt_pos_embed_shape[1] - self.num_extra_tokens
+
+        # Determine original grid shape
+        side = int(math.sqrt(num_vis))
+        factor = int(math.sqrt((num_vis * self.patch_size * self.patch_size) // 12))
+
+        # Set old grid based on aspect ratio detection
+        if side * side == num_vis:
+            old_grid = (side, side)  # square grid
+        elif 4 * factor * 3 * factor == num_vis * self.patch_size * self.patch_size:
+            old_grid = (
+                (factor * 4) // self.patch_size,
+                (factor * 3) // self.patch_size,
+            )  # 4:3 ratio
+        else:
+            state_dict[name] = self.pos_embed
+            return
+
+        # Resize position embedding
+        new_grid = (pos_h, pos_w)
+        state_dict[name] = resize_pos_embed(
+            state_dict[name],
+            old_grid,
+            new_grid,
+            mode=self.interpolate_mode,
+            num_extra_tokens=self.num_extra_tokens,
+        )
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+    @property
+    def norm2(self):
+        return self.ln2
+
+    @staticmethod
+    def resize_pos_embed(*args, **kwargs):
+        """Interface for backward-compatibility."""
+        return resize_pos_embed(*args, **kwargs)
+
+    def _freeze_stages(self):
+        # freeze position embedding
+        if self.pos_embed is not None:
+            self.pos_embed.requires_grad = False
+
+        # set dropout to eval model
+        self.drop_after_pos.eval()
+        # freeze patch embedding
+        self.patch_embed.eval()
+        for param in self.patch_embed.parameters():
+            param.requires_grad = False
+        # freeze pre-norm
+        for param in self.pre_norm.parameters():
+            param.requires_grad = False
+        # freeze cls_token
+        if self.cls_token is not None:
+            self.cls_token.requires_grad = False
+        # freeze layers
+        for i in range(1, self.frozen_stages + 1):
+            m = self.layers[i - 1]
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+        # freeze the last layer norm
+        if self.frozen_stages == len(self.layers):
+            if self.final_norm:
+                self.ln1.eval()
+                for param in self.ln1.parameters():
+                    param.requires_grad = False
+
+            if self.out_type == "avg_featmap":
+                self.ln2.eval()
+                for param in self.ln2.parameters():
+                    param.requires_grad = False
+
+    def forward(self, x):
+        B = x.shape[0]
+        x, patch_resolution = self.patch_embed(x)
+
+        if self.cls_token is not None:
+            cls_token = self.cls_token.expand(B, -1, -1)
+            x = torch.cat((cls_token, x), dim=1)
+
+        x = x + resize_pos_embed(
+            self.pos_embed,
+            self.patch_resolution,
+            patch_resolution,
+            mode=self.interpolate_mode,
+            num_extra_tokens=self.num_extra_tokens,
+        )
+        x = self.drop_after_pos(x)
+
+        x = self.pre_norm(x)  ## B x (num tokens) x embed_dim
+
+        outs = []
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+
+            if i == len(self.layers) - 1 and self.final_norm:
+                x = self.ln1(x)
+
+            if i in self.out_indices:
+                outs.append(self._format_output(x, patch_resolution))
+
+        return tuple(outs)
+
+    def _format_output(self, x, hw):
+        if self.out_type == "raw":
+            return x
+        if self.out_type == "cls_token":
+            return x[:, 0]
+
+        patch_token = x[:, self.num_extra_tokens :]
+        if self.out_type == "featmap":
+            B = x.size(0)
+            # (B, N, C) -> (B, H, W, C) -> (B, C, H, W)
+            return patch_token.reshape(B, *hw, -1).permute(0, 3, 1, 2)

sapiens/backbones/standalone/sapiens2.py

@@ -0,0 +1,908 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.init import trunc_normal_
+from torch.utils.checkpoint import checkpoint
+
+
+# ----------------------------------------------------------------------------
+def to_2tuple(x):
+    if isinstance(x, (str, bytes)):
+        return (x, x)
+    if isinstance(x, Sequence):
+        x = tuple(x)
+        if len(x) == 2:
+            return x
+        raise ValueError("Expected scalar or length-2 iterable")
+    return (x, x)
+
+
+class RopePositionEmbedding(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        *,
+        num_heads: int,
+        base: float | None = 100.0,
+        min_period: float | None = None,
+        max_period: float | None = None,
+        normalize_coords: Literal["min", "max", "separate"] = "separate",
+        shift_coords: float | None = None,
+        jitter_coords: float | None = None,
+        rescale_coords: float | None = None,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        assert embed_dim % (4 * num_heads) == 0
+        both_periods = min_period is not None and max_period is not None
+        if (base is None and not both_periods) or (base is not None and both_periods):
+            raise ValueError(
+                "Either `base` or `min_period`+`max_period` must be provided."
+            )
+
+        D_head = embed_dim // num_heads
+        self.base = base
+        self.min_period = min_period
+        self.max_period = max_period
+        self.D_head = D_head
+        self.normalize_coords = normalize_coords
+        self.shift_coords = shift_coords
+        self.jitter_coords = jitter_coords
+        self.rescale_coords = rescale_coords
+
+        # Needs persistent=True because we do teacher.load_state_dict(student.state_dict()) to initialize the teacher
+        self.dtype = dtype or torch.float32  # Don't rely on self.periods.dtype
+        self.register_buffer(
+            "periods",
+            torch.empty(D_head // 4, device=device, dtype=self.dtype),
+            persistent=True,
+        )
+        self._init_weights()
+
+    def forward(self, *, H: int, W: int) -> tuple[Tensor, Tensor]:
+        device = self.periods.device
+        dtype = self.dtype
+        dd = {"device": device, "dtype": dtype}
+        # Prepare coords in range [-1, +1]
+        if self.normalize_coords == "max":
+            max_HW = max(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / max_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / max_HW  # [W]
+        elif self.normalize_coords == "min":
+            min_HW = min(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / min_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / min_HW  # [W]
+        elif self.normalize_coords == "separate":
+            coords_h = torch.arange(0.5, H, **dd) / H  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / W  # [W]
+        else:
+            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
+        coords = torch.stack(
+            torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1
+        )  # [H, W, 2]
+        coords = coords.flatten(0, 1)  # [HW, 2]
+        coords = 2.0 * coords - 1.0  # Shift range [0, 1] to [-1, +1]
+
+        # Shift coords by adding a uniform value in [-shift, shift]
+        if self.training and self.shift_coords is not None:
+            shift_hw = torch.empty(2, **dd).uniform_(
+                -self.shift_coords, self.shift_coords
+            )
+            coords += shift_hw[None, :]
+
+        # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter]
+        if self.training and self.jitter_coords is not None:
+            jitter_max = np.log(self.jitter_coords)
+            jitter_min = -jitter_max
+            jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp()
+            coords *= jitter_hw[None, :]
+
+        # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale]
+        if self.training and self.rescale_coords is not None:
+            rescale_max = np.log(self.rescale_coords)
+            rescale_min = -rescale_max
+            rescale_hw = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp()
+            coords *= rescale_hw
+
+        # Prepare angles and sin/cos
+        angles = (
+            2 * math.pi * coords[:, :, None] / self.periods[None, None, :]
+        )  # [HW, 2, D//4]
+        angles = angles.flatten(1, 2)  # [HW, D//2]
+        angles = angles.tile(2)  # [HW, D]
+        cos = torch.cos(angles)  # [HW, D]
+        sin = torch.sin(angles)  # [HW, D]
+
+        return (sin, cos)  # 2 * [HW, D]
+
+    def _init_weights(self):
+        device = self.periods.device
+        dtype = self.dtype
+        if self.base is not None:
+            periods = self.base ** (
+                2
+                * torch.arange(self.D_head // 4, device=device, dtype=dtype)
+                / (self.D_head // 2)
+            )  # [D//4]
+        else:
+            base = self.max_period / self.min_period
+            exponents = torch.linspace(
+                0, 1, self.D_head // 4, device=device, dtype=dtype
+            )  # [D//4] range [0, 1]
+            periods = base**exponents  # range [1, max_period / min_period]
+            periods = periods / base  # range [min_period / max_period, 1]
+            periods = periods * self.max_period  # range [min_period, max_period]
+        self.periods.data = periods
+
+
+# -------------------------------------------------------------------------------
+class Tokenizer(nn.Module):
+    """Stacked window self‑attention that emits one token per window
+    by re‑using TransformerEncoderLayer blocks."""
+
+    def __init__(
+        self,
+        embed_dims: int,
+        window_size: int = 4,
+        num_heads: int = 4,
+        num_tokenizer_layers: int = 1,
+        qkv_bias: bool = True,
+        use_qk_norm: bool = False,
+        chunk_size: int = 1024,  # max windows per chunk
+    ):
+        super().__init__()
+        self.ws = window_size
+        self.chunk_size = chunk_size
+
+        # local absolute positional embeddings for [CLS] + patch tokens
+        self.local_pos_embed = nn.Parameter(
+            torch.zeros(1, 1 + window_size * window_size, embed_dims)
+        )
+        trunc_normal_(self.local_pos_embed, std=0.02)
+
+        # build N identical TransformerEncoderLayer blocks
+        self.blocks = nn.ModuleList(
+            [
+                TransformerEncoderLayer2(
+                    embed_dims=embed_dims,
+                    num_heads=num_heads,
+                    feedforward_channels=embed_dims * 4,  # standard FFN size
+                    qkv_bias=qkv_bias,
+                    use_qk_norm=use_qk_norm,
+                )
+                for _ in range(num_tokenizer_layers)
+            ]
+        )
+
+        # shared CLS token for pooling
+        self.w_cls = nn.Parameter(torch.zeros(1, 1, embed_dims))
+        trunc_normal_(self.w_cls, std=0.02)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        hw: Tuple[int, int],
+    ) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        """Args:
+           x  : B, N, C   (N = H*W)
+           hw : (H, W) before reduction
+        Returns:
+           x_ : B, (H/ws)*(W/ws), C
+           hw_: (H/ws, W/ws)
+        """
+        B, N, C = x.shape
+        H, W = hw
+        ws = self.ws
+        assert H % ws == 0 and W % ws == 0, (
+            f"Image size {H}×{W} must be divisible by window {ws}."
+        )
+
+        # reshape tokens → non‑overlapping windows
+        x = x.view(B, H, W, C)
+
+        ph, pw = H // ws, W // ws  ## ints in eager mode
+        ph, pw = int(ph), int(pw)  ## ints in scripting mode
+        x = x.view(B, ph, ws, pw, ws, C)  # B, H/ws, ws, W/ws, ws, C
+        x = x.permute(0, 1, 3, 2, 4, 5)  # B, H/ws, W/ws, ws, ws, C
+        x = x.contiguous().view(B * ph * pw, ws * ws, C)  # (B*H/ws*W/ws), ws², C))
+
+        total_windows = x.size(0)
+        chunk_size = int(min(self.chunk_size, total_windows))
+        token_out = x.new_empty(total_windows, C)
+
+        use_ckpt = self.training and torch.is_grad_enabled()
+
+        def _run_blocks(t: torch.Tensor) -> torch.Tensor:
+            for blk in self.blocks:
+                t = blk(t)
+            return t
+
+        for i in range(0, total_windows, chunk_size):
+            chunk = x[i : i + chunk_size]  # (m, ws², C)
+            m = chunk.size(0)
+            cls = self.w_cls.expand(m, -1, -1)  # (m, 1, C)
+            chunk = torch.cat([cls, chunk], dim=1)  # (m, 1+ws², C)
+            chunk = chunk + self.local_pos_embed  # add local PE
+
+            if use_ckpt:
+                chunk = checkpoint(_run_blocks, chunk, use_reentrant=False)
+            else:
+                chunk = _run_blocks(chunk)
+
+            token_out[i : i + m] = chunk[:, 0]  # take CLS out
+
+        token = token_out.view(B, ph * pw, C)  # (B, (H/ws)*(W
+        return token, (ph, pw)
+
+
+# -------------------------------------------------------------------------------
+class GroupedQueryAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        num_kv_heads=None,
+        input_dims=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        qkv_bias=True,
+        qk_scale=None,
+        proj_bias=True,
+        use_qk_norm=True,
+        v_shortcut=False,
+        layer_scale_init_value=0.0,
+    ):
+        super().__init__()
+        # Core dims
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        assert self.num_heads % self.num_kv_heads == 0, (
+            "num_kv_heads must divide num_heads"
+        )
+        self.head_dim = embed_dims // num_heads
+        self.input_dims = input_dims or embed_dims
+        # Features
+        self.attn_drop = attn_drop
+        self.v_shortcut = v_shortcut
+        self.use_qk_norm = use_qk_norm
+
+        # Attention operation selection
+        if qk_scale is not None:
+            scale = qk_scale
+        else:
+            scale = self.head_dim**-0.5
+
+        assert qk_scale is None, "qk_scale is not supported"
+        self.attn_op = F.scaled_dot_product_attention
+
+        # Q/K/V projections
+        self.wq = nn.Linear(self.input_dims, embed_dims, bias=qkv_bias)
+        self.wk = nn.Linear(
+            self.input_dims, self.num_kv_heads * self.head_dim, bias=qkv_bias
+        )
+        self.wv = nn.Linear(
+            self.input_dims, self.num_kv_heads * self.head_dim, bias=qkv_bias
+        )
+
+        if self.use_qk_norm:
+            self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6)
+            self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6)
+
+        # Output projection + dropout
+        self.proj = nn.Linear(embed_dims, embed_dims, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        # Optional LayerScale
+        if layer_scale_init_value > 0:
+            self.gamma = LayerScale(embed_dims, scale=layer_scale_init_value)
+        else:
+            self.gamma = nn.Identity()
+
+    def apply_rope(
+        self, q: Tensor, k: Tensor, rope: Tensor | Tuple[Tensor, Tensor]
+    ) -> Tuple[Tensor, Tensor]:
+        # All operations will use the dtype of rope, the output is cast back to the dtype of q and k
+        q_dtype = q.dtype
+        k_dtype = k.dtype
+        sin, cos = rope
+        rope_dtype = sin.dtype
+        q = q.to(dtype=rope_dtype)
+        k = k.to(dtype=rope_dtype)
+        N = q.shape[-2]
+        prefix = N - sin.shape[-2]  ## extra tokens
+        assert prefix >= 0
+        q_prefix = q[:, :, :prefix, :]
+        q = self._rope_apply(q[:, :, prefix:, :], sin, cos)  # [B, head, hw, D//head]
+        q = torch.cat((q_prefix, q), dim=-2)  # [B, head, N, D//head]
+        k_prefix = k[:, :, :prefix, :]
+        k = self._rope_apply(k[:, :, prefix:, :], sin, cos)  # [B, head, hw, D//head]
+        k = torch.cat((k_prefix, k), dim=-2)  # [B, head, N, D//head]
+        q = q.to(dtype=q_dtype)
+        k = k.to(dtype=k_dtype)
+        return q, k
+
+    def _rope_rotate_half(self, x: Tensor) -> Tensor:
+        # x:   [ x0  x1  x2  x3  x4  x5]
+        # out: [-x3 -x4 -x5  x0  x1  x2]
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat([-x2, x1], dim=-1)
+
+    def _rope_apply(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+        # x:   [..., D], eg [x0,     x1,   x2,   x3,   x4,   x5]
+        # sin: [..., D], eg [sin0, sin1, sin2, sin0, sin1, sin2]
+        # cos: [..., D], eg [cos0, cos1, cos2, cos0, cos1, cos2]
+        return (x * cos) + (self._rope_rotate_half(x) * sin)
+
+    def forward(self, x, rope=None):
+        B, N, _ = x.shape
+        # Q: (B, N, num_heads, head_dim)
+        q = self.wq(x).view(B, N, self.num_heads, self.head_dim)
+        # K/V: (B, N, num_kv_heads, head_dim)
+        k = self.wk(x).view(B, N, self.num_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, N, self.num_kv_heads, self.head_dim)
+
+        # (B, heads, N, head_dim)
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        # Repeat KV heads if group ratio >1
+        if self.num_kv_heads != self.num_heads:
+            factor = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(factor, dim=1)
+            v = v.repeat_interleave(factor, dim=1)
+
+        if rope is not None:
+            q, k = self.apply_rope(q, k, rope)
+
+        # Scaled dot-product attention
+        attn_out = self.attn_op(
+            q, k, v, dropout_p=self.attn_drop if self.training else 0.0
+        )  # (B, num_heads, N, head_dim)
+
+        # Merge heads -> (B, N, embed_dims)
+        out = attn_out.permute(0, 2, 1, 3).reshape(B, N, self.embed_dims)
+
+        # Output projection + drop + layer scale
+        out = self.proj(out)
+        out = self.gamma(self.proj_drop(out))
+
+        # Optional V-shortcut (only when MQA)
+        if self.v_shortcut and self.num_kv_heads == 1:
+            raise NotImplementedError
+        return out
+
+
+# -------------------------------------------------------------------------------
+class TransformerEncoderLayer2(nn.Module):
+    def __init__(
+        self,
+        embed_dims,
+        num_heads,
+        num_kv_heads=None,
+        feedforward_channels=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        layer_scale_init_value=0.0,
+        use_qk_norm=True,
+        qkv_bias=True,
+    ):
+        super(TransformerEncoderLayer2, self).__init__()
+
+        self.embed_dims = embed_dims
+        self.ln1 = nn.RMSNorm(self.embed_dims, eps=1e-6)
+        self.attn = GroupedQueryAttention(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            attn_drop=attn_drop_rate,
+            proj_drop=drop_rate,
+            qkv_bias=qkv_bias,
+            layer_scale_init_value=layer_scale_init_value,
+            use_qk_norm=use_qk_norm,
+        )
+
+        self.ln2 = nn.RMSNorm(self.embed_dims, eps=1e-6)
+        self.ffn = SwiGLUFFN(
+            embed_dims=embed_dims,
+            feedforward_channels=feedforward_channels,
+        )
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+    @property
+    def norm2(self):
+        return self.ln2
+
+    def forward(self, x, rope=None):
+        x = x + self.attn(self.ln1(x), rope=rope)
+        x = self.ffn(self.ln2(x), identity=x)
+        return x
+
+
+##-----------------------------------
+class Sapiens2(nn.Module):
+    arch_zoo = {
+        **dict.fromkeys(
+            ["sapiens2_0.1b"],
+            {
+                "embed_dims": 768,
+                "num_layers": 12,
+                "num_heads": 12,
+                "feedforward_channels": 768 * 4,
+                "num_tokenizer_layers": 2,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_0.4b"],
+            {
+                "embed_dims": 1024,
+                "num_layers": 24,
+                "num_heads": 16,
+                "feedforward_channels": 1024 * 4,
+                "num_tokenizer_layers": 2,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_0.8b"],
+            {
+                "embed_dims": 1280,
+                "num_layers": 32,
+                "num_heads": 16,
+                "feedforward_channels": 1280 * 4,
+                "num_tokenizer_layers": 3,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_1b"],
+            {
+                "embed_dims": 1536,
+                "num_layers": 40,
+                "num_heads": 24,
+                "feedforward_channels": 1536 * 4,
+                "num_tokenizer_layers": 4,
+            },
+        ),
+        **dict.fromkeys(
+            ["sapiens2_5b"],
+            {
+                "embed_dims": 2432,
+                "num_layers": 56,
+                "num_heads": 32,
+                "feedforward_channels": 2432 * 4,
+                "num_tokenizer_layers": 6,
+            },
+        ),
+    }
+
+    num_extra_tokens = 1  # class token
+    OUT_TYPES = {"raw", "cls_token", "featmap"}
+
+    def __init__(
+        self,
+        arch="sapiens2_1b",
+        img_size=(1024, 768),
+        patch_size=16,
+        in_channels=3,
+        out_indices=-1,
+        drop_rate=0.0,
+        window_size=4,
+        use_tokenizer=False,  ## 4k resolution
+        use_qk_norm=True,
+        qkv_bias=True,
+        final_norm=True,
+        out_type="raw",
+        with_cls_token=True,
+        layer_scale_init_value=1e-4,  ## non zero init to activate layerscale
+        frozen_stages=-1,
+        patch_cfg=dict(),
+        layer_cfgs=dict(),
+        pos_embed_rope_base: float = 100.0,
+        pos_embed_rope_min_period: float | None = None,
+        pos_embed_rope_max_period: float | None = None,
+        pos_embed_rope_normalize_coords: Literal["min", "max", "separate"] = "separate",
+        pos_embed_rope_shift_coords: float | None = None,
+        pos_embed_rope_jitter_coords: float | None = None,
+        pos_embed_rope_rescale_coords: float | None = None,
+        pos_embed_rope_dtype: str = "bf16",
+        n_storage_tokens: int = 8,
+    ):
+        super().__init__()
+
+        arch = arch.lower()
+        assert arch in set(self.arch_zoo), (
+            f"Arch {arch} is not in default archs {set(self.arch_zoo)}"
+        )
+        self.arch_settings = self.arch_zoo[arch]
+
+        self.embed_dims = self.arch_settings["embed_dims"]
+        self.num_layers = self.arch_settings["num_layers"]
+        self.patch_size = patch_size
+
+        self.window_size = window_size
+        img_size = to_2tuple(img_size)
+        encoder_img_size = (
+            (img_size[0] // window_size, img_size[1] // window_size)
+            if use_tokenizer
+            else img_size
+        )
+        self.img_size = to_2tuple(encoder_img_size)
+
+        # Set patch embedding
+        _patch_cfg = dict(
+            in_channels=in_channels,
+            input_size=self.img_size,
+            embed_dims=self.embed_dims,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+        _patch_cfg.update(patch_cfg)
+        self.patch_embed = PatchEmbed(**_patch_cfg)
+        self.patch_resolution = self.patch_embed.init_out_size
+        num_patches = self.patch_resolution[0] * self.patch_resolution[1]
+
+        self.rope_embed = RopePositionEmbedding(
+            embed_dim=self.embed_dims,
+            num_heads=self.arch_settings["num_heads"],
+            base=pos_embed_rope_base,
+            min_period=pos_embed_rope_min_period,
+            max_period=pos_embed_rope_max_period,
+            normalize_coords=pos_embed_rope_normalize_coords,
+            shift_coords=pos_embed_rope_shift_coords,
+            jitter_coords=pos_embed_rope_jitter_coords,
+            rescale_coords=pos_embed_rope_rescale_coords,
+            dtype=torch.bfloat16 if pos_embed_rope_dtype == "bf16" else torch.float32,
+        )
+
+        # Set out type
+        if out_type not in self.OUT_TYPES:
+            raise ValueError(
+                f"Unsupported `out_type` {out_type}, please "
+                f"choose from {self.OUT_TYPES}"
+            )
+        self.out_type = out_type
+
+        if use_tokenizer == True:
+            self.tokenizer = Tokenizer(
+                embed_dims=self.embed_dims,
+                window_size=self.window_size,
+                num_heads=self.arch_settings["num_heads"],
+                num_tokenizer_layers=self.arch_settings["num_tokenizer_layers"],
+                qkv_bias=True,
+                use_qk_norm=False,
+            )
+        else:
+            self.tokenizer = None
+
+        # Set cls + storage tokens
+        self.with_cls_token = with_cls_token
+        if with_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims))
+        elif out_type != "cls_token":
+            self.cls_token = None
+            self.num_extra_tokens = 0
+        else:
+            raise ValueError('with_cls_token must be True when `out_type="cls_token"`.')
+
+        ## registers
+        self.n_storage_tokens = int(n_storage_tokens)
+        self.storage_tokens = (
+            nn.Parameter(torch.zeros(1, self.n_storage_tokens, self.embed_dims))
+            if self.n_storage_tokens > 0
+            else None
+        )
+        # how many non-patch tokens are at the front
+        self.num_extra_tokens = (
+            1 if self.cls_token is not None else 0
+        ) + self.n_storage_tokens
+
+        if isinstance(out_indices, int):
+            out_indices = [out_indices]
+        assert isinstance(out_indices, Sequence), (
+            f'"out_indices" must by a sequence or int, get {type(out_indices)} instead.'
+        )
+        for i, index in enumerate(out_indices):
+            if index < 0:
+                out_indices[i] = self.num_layers + index
+            assert 0 <= out_indices[i] <= self.num_layers, (
+                f"Invalid out_indices {index}"
+            )
+        self.out_indices = out_indices
+
+        self.blocks = nn.Sequential()
+        if isinstance(layer_cfgs, dict):
+            layer_cfgs = [layer_cfgs] * self.num_layers
+
+        mhsa_early, mhsa_late = 8, 8
+        for i in range(self.num_layers):
+            if i < mhsa_early or i >= self.num_layers - mhsa_late:
+                num_kv_heads = None  ## use MHSA
+            else:
+                num_kv_heads = self.arch_settings["num_heads"] // 2  # Use GQA
+
+            _layer_cfg = dict(
+                embed_dims=self.embed_dims,
+                num_heads=self.arch_settings["num_heads"],
+                num_kv_heads=num_kv_heads,
+                feedforward_channels=self.arch_settings["feedforward_channels"],
+                use_qk_norm=use_qk_norm,
+                layer_scale_init_value=layer_scale_init_value,
+                drop_rate=drop_rate,
+                qkv_bias=qkv_bias,
+            )
+            _layer_cfg.update(layer_cfgs[i])
+            self.blocks.append(TransformerEncoderLayer2(**_layer_cfg))
+
+        self.frozen_stages = frozen_stages
+
+        self.final_norm = final_norm
+        if final_norm:
+            self.ln1 = nn.RMSNorm(self.embed_dims, eps=1e-6)
+
+        # freeze stages only when self.frozen_stages > 0
+        if self.frozen_stages > 0:
+            self._freeze_stages()
+
+        ## load init weights
+        self.init_weights()
+
+        return
+
+    def init_weights(self):
+        # Initialize class token and storagr token embeddings
+        if self.with_cls_token:
+            trunc_normal_(self.cls_token, std=0.02)
+
+        if self.storage_tokens is not None:
+            trunc_normal_(self.storage_tokens, std=0.02)
+
+        # Apply custom initialization to all submodules
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # Use a truncated normal distribution for linear layer weights
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+        elif isinstance(m, (nn.LayerNorm, nn.RMSNorm)):
+            # Initialize normalization layers to act as an identity function
+            if hasattr(m, "bias") and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+            if hasattr(m, "weight") and m.weight is not None:
+                nn.init.constant_(m.weight, 1.0)
+
+        elif isinstance(m, nn.Conv2d):
+            # Initialize conv layer weights like linear layers
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def _freeze_stages(self):
+        ## freeze tokenizer
+        if self.frozen_stages >= 1 and self.tokenizer is not None:
+            self.tokenizer.eval()
+            for param in self.tokenizer.parameters():
+                param.requires_grad = False
+
+        # freeze patch embedding
+        self.patch_embed.eval()
+        for param in self.patch_embed.parameters():
+            param.requires_grad = False
+        # freeze cls_token
+        if self.cls_token is not None:
+            self.cls_token.requires_grad = False
+        if self.storage_tokens is not None:
+            self.storage_tokens.requires_grad = False
+        # freeze layers
+        for i in range(1, self.frozen_stages + 1):
+            m = self.blocks[i - 1]
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+
+        # freeze the last layer norm
+        if self.frozen_stages == len(self.blocks):
+            if self.final_norm:
+                self.ln1.eval()
+                for param in self.ln1.parameters():
+                    param.requires_grad = False
+
+    def forward(self, x):
+        B = x.shape[0]
+
+        x, patch_resolution = self.patch_embed(x)  # (B, 256*256, C)
+        if self.tokenizer is not None:
+            x, patch_resolution = self.tokenizer(x, patch_resolution)
+
+        # prepend [CLS] and storage tokens
+        prepend = []
+        if self.cls_token is not None:
+            prepend.append(self.cls_token.expand(B, -1, -1))
+        if self.storage_tokens is not None:
+            prepend.append(self.storage_tokens.expand(B, -1, -1))
+        if len(prepend) > 0:
+            x = torch.cat(prepend + [x], dim=1)
+
+        rope_sincos = self.rope_embed(H=patch_resolution[0], W=patch_resolution[1])
+        outs = []
+        for i, layer in enumerate(self.blocks):
+            x = layer(x, rope=rope_sincos)
+
+            if i == len(self.blocks) - 1 and self.final_norm:
+                x = self.ln1(x)
+
+            if i in self.out_indices:
+                outs.append(self._format_output(x, patch_resolution))
+
+        return tuple(outs)
+
+    def _format_output(self, x, hw):
+        if self.out_type == "raw":
+            return x
+        if self.out_type == "cls_token":
+            return x[:, 0]
+
+        patch_token = x[:, self.num_extra_tokens :]
+        if self.out_type == "featmap":
+            B = x.size(0)
+            # (B, N, C) -> (B, H, W, C) -> (B, C, H, W)
+            return patch_token.reshape(B, *hw, -1).permute(0, 3, 1, 2)
+
+    @property
+    def norm1(self):
+        return self.ln1
+
+
+# ----------------------------------------------------------------------------
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inplace: bool = False,
+        data_format: str = "channels_last",
+        scale: float = 1e-5,
+    ):
+        super().__init__()
+        assert data_format in (
+            "channels_last",
+            "channels_first",
+        ), "'data_format' could only be channels_last or channels_first."
+        self.inplace = inplace
+        self.data_format = data_format
+        self.weight = nn.Parameter(torch.ones(dim) * scale)
+
+    def forward(self, x) -> torch.Tensor:
+        if self.data_format == "channels_first":
+            shape = tuple((1, -1, *(1 for _ in range(x.dim() - 2))))
+        else:
+            shape = tuple((*(1 for _ in range(x.dim() - 1)), -1))
+        if self.inplace:
+            return x.mul_(self.weight.view(*shape))
+        else:
+            return x * self.weight.view(*shape)
+
+
+# ----------------------------------------------------------------------------
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        embed_dims=768,
+        kernel_size=16,
+        stride=16,
+        padding="corner",
+        dilation=1,
+        bias=True,
+        input_size=None,
+    ):
+        super().__init__()
+
+        self.embed_dims = embed_dims
+        if stride is None:
+            stride = kernel_size
+
+        kernel_size = to_2tuple(kernel_size)
+        stride = to_2tuple(stride)
+        dilation = to_2tuple(dilation)
+        padding = 0
+        padding = to_2tuple(padding)
+
+        self.projection = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=embed_dims,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+        if input_size:
+            input_size = to_2tuple(input_size)
+            self.init_input_size = input_size
+            h_out = (
+                input_size[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1
+            ) // stride[0] + 1
+            w_out = (
+                input_size[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1
+            ) // stride[1] + 1
+            self.init_out_size = (h_out, w_out)
+        else:
+            self.init_input_size = None
+            self.init_out_size = None
+
+    def forward(self, x):
+        x = self.projection(x)
+        out_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)
+        return x, out_size
+
+
+# ----------------------------------------------------------------------------
+class SwiGLUFFN(nn.Module):
+    """SwiGLU FFN layer.
+    https://github.com/facebookresearch/dinov2/blob/main/dinov2/layers/swiglu_ffn.py
+    """  # noqa
+
+    def __init__(
+        self,
+        embed_dims: int,
+        feedforward_channels: Optional[int] = None,
+        out_dims: Optional[int] = None,
+        layer_scale_init_value: float = 0.0,
+        bias: bool = True,
+        add_identity: bool = True,
+    ) -> None:
+        super().__init__()
+        self.embed_dims = embed_dims
+        self.out_dims = out_dims or embed_dims
+        hidden_dims = feedforward_channels or embed_dims
+
+        self.w12 = nn.Linear(self.embed_dims, 2 * hidden_dims, bias=bias)
+        self.w3 = nn.Linear(hidden_dims, self.out_dims, bias=bias)
+
+        if layer_scale_init_value > 0:
+            self.gamma2 = LayerScale(dim=embed_dims, scale=layer_scale_init_value)
+        else:
+            self.gamma2 = nn.Identity()
+
+        self.add_identity = add_identity
+
+    def forward(
+        self, x: torch.Tensor, identity: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        out = self.w3(hidden)
+        out = self.gamma2(out)
+
+        if self.out_dims != self.embed_dims or not self.add_identity:
+            # due to the dimension inconsistence or user setting
+            # not to apply residual operation
+            return out
+
+        if identity is None:
+            identity = x
+        return identity + out

sapiens/dense/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import pathlib
+import pkgutil
+
+from .. import __version__
+
+_src = pathlib.Path(__file__).with_name("src")
+__path__ = pkgutil.extend_path(__path__, __name__)  # allow namespace merge
+__path__.append(str(_src))
+del pathlib, pkgutil, _src
+
+
+# -----------------------------------------------------
+from importlib import import_module as _imp
+
+_pkg = _imp(__name__ + ".src")  # runs src/__init__.py

sapiens/dense/configs/albedo/render_people/sapiens2_0.4b_albedo_render_people-1024x768.py

@@ -0,0 +1,274 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 10
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.4b"
+embed_dim = 1024
+num_layers = 24
+num_heads = 16
+
+layer_decay_rate = 0.8
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.4b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="AlbedoVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="AlbedoRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="AlbedoResize", height=1024, width=768),
+    dict(type="RandomGaussianNoise", prob=0.2, var_range=(5.0, 20.0)),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="AlbedoResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="AlbedoPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="AlbedoResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+
+render_people_dataset = dict(
+    type="AlbedoRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="AlbedoRenderPeopleDataset",
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="AlbedoEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="AlbedoEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="AlbedoHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),
+            dict(type="AlbedoGradL1Loss", loss_weight=2.0),
+            # dict(type="AlbedoLowFreqL1Loss", down_sample=32, loss_weight=1.0),
+            dict(type="AlbedoChromaticityL1Loss", loss_weight=1.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=2.0, norm_type=2.0)

sapiens/dense/configs/albedo/render_people/sapiens2_0.8b_albedo_render_people-1024x768.py

@@ -0,0 +1,275 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+# val_every_iters = 2000
+val_every_iters = 10000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 10
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.8b"
+embed_dim = 1280
+num_layers = 32
+num_heads = 16
+
+layer_decay_rate = 0.85
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.8b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="AlbedoVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="AlbedoRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="AlbedoResize", height=1024, width=768),
+    dict(type="RandomGaussianNoise", prob=0.2, var_range=(5.0, 20.0)),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="AlbedoResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="AlbedoPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="AlbedoResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+
+render_people_dataset = dict(
+    type="AlbedoRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="AlbedoRenderPeopleDataset",
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="AlbedoEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="AlbedoEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="AlbedoHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),
+            dict(type="AlbedoGradL1Loss", loss_weight=2.0),
+            # dict(type="AlbedoLowFreqL1Loss", down_sample=32, loss_weight=1.0),
+            dict(type="AlbedoChromaticityL1Loss", loss_weight=1.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/albedo/render_people/sapiens2_1b_albedo_render_people-1024x768.py

@@ -0,0 +1,274 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 10
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_1b"
+embed_dim = 1536
+num_layers = 40
+num_heads = 24
+layer_decay_rate = 0.9
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_1b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="AlbedoVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="AlbedoRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="AlbedoResize", height=1024, width=768),
+    dict(type="RandomGaussianNoise", prob=0.2, var_range=(5.0, 20.0)),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="AlbedoResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="AlbedoPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="AlbedoResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+
+render_people_dataset = dict(
+    type="AlbedoRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="AlbedoRenderPeopleDataset",
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="AlbedoEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="AlbedoEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="AlbedoHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),
+            dict(type="AlbedoGradL1Loss", loss_weight=2.0),
+            # dict(type="AlbedoLowFreqL1Loss", down_sample=32, loss_weight=1.0),
+            dict(type="AlbedoChromaticityL1Loss", loss_weight=1.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/albedo/render_people/sapiens2_5b_albedo_render_people-1024x768.py

@@ -0,0 +1,280 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+# val_every_iters = 1000
+# val_every_iters = 20000
+val_every_iters = 40000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 10
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_5b"
+embed_dim = 2432
+num_layers = 56
+num_heads = 32
+layer_decay_rate = 0.94
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_5b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+        # parallelism_cfg=dict(
+        #     dp_shard_size=2,  # Fully Sharded Data Parallel degree
+        #     dp_replicate_size=1,  # Data Parallel degree
+        #     tp_size=1,  # Tensor Parallel degree
+        #     cp_size=4,  # Context Parallel degree
+        # ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="AlbedoVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="AlbedoRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="AlbedoRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="AlbedoResize", height=1024, width=768),
+    dict(type="RandomGaussianNoise", prob=0.2, var_range=(5.0, 20.0)),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="AlbedoResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="AlbedoPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="AlbedoResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="AlbedoPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+
+render_people_dataset = dict(
+    type="AlbedoRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    shuffle=False,
+    dataset=dict(
+        type="AlbedoRenderPeopleDataset",
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_albedo_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="AlbedoEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="AlbedoEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="AlbedoHead",
+        in_channels=embed_dim,
+        upsample_channels=[1536, 768, 512, 256],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),
+            dict(type="AlbedoGradL1Loss", loss_weight=2.0),
+            # dict(type="AlbedoLowFreqL1Loss", down_sample=32, loss_weight=1.0),
+            dict(type="AlbedoChromaticityL1Loss", loss_weight=1.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/normal/metasim_render_people/sapiens2_0.4b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.4b"
+embed_dim = 1024
+num_layers = 24
+num_heads = 16
+
+layer_decay_rate = 0.8
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.4b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=8,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=2.0, norm_type=2.0)

sapiens/dense/configs/normal/metasim_render_people/sapiens2_0.8b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 1e4
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.8b"
+embed_dim = 1280
+num_layers = 32
+num_heads = 16
+
+layer_decay_rate = 0.85
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.8b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=8,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/normal/metasim_render_people/sapiens2_1b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,306 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+# num_iters = 1e4  ## light finetune
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_1b"
+embed_dim = 1536
+num_layers = 40
+num_heads = 24
+layer_decay_rate = 0.9
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_1b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        # with_cls_token=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[768, 512, 256, 128],  ## 1K resolution
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/normal/metasim_render_people/sapiens2_5b_normal_metasim_render_people-1024x768.py

@@ -0,0 +1,312 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+# num_iters = 1e4  ## light finetune
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_5b"
+embed_dim = 2432
+num_layers = 56
+num_heads = 32
+layer_decay_rate = 0.94
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_5b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+        # parallelism_cfg=dict(
+        #     dp_shard_size=2,  # Fully Sharded Data Parallel degree
+        #     dp_replicate_size=1,  # Data Parallel degree
+        #     tp_size=1,  # Tensor Parallel degree
+        #     cp_size=4,  # Context Parallel degree
+        # ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="NormalVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(type="RandomDownUpSampleImage", scale_range=(0.1, 0.7), prob=0.2),
+    dict(
+        type="NormalRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="NormalRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="NormalResize", height=1024, width=768),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(type="RandomSolarize", prob=0.3, threshold=128),
+    dict(type="NormalGenerateTarget"),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="NormalResize", height=1024, width=768, test_mode=True),
+    dict(
+        type="NormalPackInputs",
+        test_mode=True,
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="NormalResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="NormalPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+        ),
+    ),
+]
+
+metasim_dataset = dict(
+    type="NormalMetaSimDataset",
+    airstore_template="airstore://codec_avatar_sapiens_metasim_v1_no_user_data",
+    json_path=f"{_DATA_ROOT}/seg/data/metasim/meta_data_v1.json",
+)
+
+render_people_dataset = dict(
+    type="NormalRenderPeopleBodyDataset",  ## body only
+    data_root=f"{_DATA_ROOT}/synthetic",
+    seg_data_root=f"{_DATA_ROOT}/RenderPeople/part_seg",
+)
+
+multihuman_render_people_dataset = dict(
+    type="NormalRenderPeopleMultihumanDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human",
+    normal_extension=".npz",
+    seg_data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_multi_human/part_seg",  ## supervise on face for multihuman
+)
+
+# train_datasets = 2 * [metasim_dataset] + [
+#     render_people_dataset,
+#     multihuman_render_people_dataset,
+# ]
+
+# train_datasets = [render_people_dataset]
+# train_datasets = [multihuman_render_people_dataset]
+train_datasets = [metasim_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="NormalRenderPeopleBodyDataset",  ## body only
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/metasim/evaluation",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="NormalEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="NormalEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="NormalHead",
+        in_channels=embed_dim,
+        upsample_channels=[1536, 768, 512, 256],  ## 1K resolution
+        conv_out_channels=[128, 64, 32],
+        conv_kernel_sizes=[3, 3, 3],
+        loss_decode=[
+            dict(
+                type="NormalCosineSimilarityLoss",
+                loss_weight=10.0,
+            ),
+            dict(type="L1Loss", loss_weight=1.0),
+            dict(type="NormalGradL1Loss", loss_weight=10.0),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    # lr=5e-4,
+    lr=1e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/pointmap/render_people/sapiens2_0.4b_pointmap_render_people-1024x768.py

@@ -0,0 +1,322 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4  ## 16 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.4b"
+embed_dim = 1024
+num_layers = 24
+num_heads = 16
+
+layer_decay_rate = 0.8
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.4b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+
+patch_size = 16
+num_tokens = (image_size[0] // patch_size) * (image_size[1] // patch_size)
+canonical_focal_length = 768.0
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+    ## Note: to merge sharded weight using FSDP
+    # accelerate merge-weights pytorch_model_fsdp_0/ .
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="PointmapVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(
+        type="PointmapRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="PointmapResize", height=1024, width=768),
+    ## target is same res as output, otherwise we get artifacts.
+    dict(
+        type="PointmapGenerateTarget",
+        canonical_focal_length=canonical_focal_length,
+        target_downsample_factor=1,
+    ),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "flip",
+            "flip_direction",
+            "original_K",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="PointmapResize", height=1024, width=768),
+    dict(type="PointmapGenerateTarget", canonical_focal_length=canonical_focal_length),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "padding_size",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="PointmapResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "padding_size",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+render_people_dataset = dict(
+    type="PointmapRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="PointmapRenderPeopleDataset",
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="PointmapEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="PointmapEstimator",
+    canonical_focal_length=canonical_focal_length,
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="PointmapHead",
+        in_channels=embed_dim,
+        upsample_channels=[1536, 768, 512, 256],
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        scale_conv_out_channels=(1536, 512, 128),
+        scale_conv_kernel_sizes=(1, 1, 1),
+        scale_final_layer=(
+            (num_tokens // ((2 * 2 * 2) * (2 * 2 * 2))) * 128,
+            512,
+            128,
+            1,
+        ),  ## scale regress
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),  ## on pointmap, XYZ
+            dict(
+                type="MultiscaleL1Loss",
+                loss_weight=1.0,
+                scale_factor=2,
+            ),
+            dict(type="SiLogLoss", loss_weight=1.0),  ## only applies silog loss
+            dict(
+                type="PointmapIntrinsicsConsistencyLoss",
+                loss_weight=1.0,
+            ),
+            dict(
+                type="PointmapShiftInvariantL1Loss",
+                loss_weight=1.0,
+            ),
+            dict(type="PointmapNormalLoss", loss_weight=2.0),
+            dict(
+                type="PointmapScaleL1Loss", loss_weight=4.0
+            ),  ## Canonical XYZ = scale * XYZ
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=2.0, norm_type=2.0)

sapiens/dense/configs/pointmap/render_people/sapiens2_0.8b_pointmap_render_people-1024x768.py

@@ -0,0 +1,325 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4  ## 16 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+# num_iters = 1e4  ## light finetune
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.8b"
+embed_dim = 1280
+num_layers = 32
+num_heads = 16
+
+layer_decay_rate = 0.85
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.8b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+
+patch_size = 16
+num_tokens = (image_size[0] // patch_size) * (image_size[1] // patch_size)
+canonical_focal_length = 768.0
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            # mixed_precision=dict(
+            #     param_dtype="bf16",
+            #     reduce_dtype="bf16",
+            # ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="PointmapVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(
+        type="PointmapRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="PointmapResize", height=1024, width=768),
+    ## target is same res as output, otherwise we get artifacts.
+    dict(
+        type="PointmapGenerateTarget",
+        canonical_focal_length=canonical_focal_length,
+        target_downsample_factor=1,
+    ),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "flip",
+            "flip_direction",
+            "original_K",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="PointmapResize", height=1024, width=768),
+    dict(type="PointmapGenerateTarget", canonical_focal_length=canonical_focal_length),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "padding_size",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="PointmapResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "padding_size",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+render_people_dataset = dict(
+    type="PointmapRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="PointmapRenderPeopleDataset",
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="PointmapEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="PointmapEstimator",
+    canonical_focal_length=canonical_focal_length,
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="PointmapHead",
+        in_channels=embed_dim,
+        upsample_channels=[1536, 768, 512, 256],
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        scale_conv_out_channels=(1536, 512, 128),
+        scale_conv_kernel_sizes=(1, 1, 1),
+        scale_final_layer=(
+            (num_tokens // ((2 * 2 * 2) * (2 * 2 * 2))) * 128,
+            512,
+            128,
+            1,
+        ),  ## scale regress
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),  ## on pointmap, XYZ
+            dict(
+                type="MultiscaleL1Loss",
+                loss_weight=1.0,
+                scale_factor=2,
+            ),
+            dict(type="SiLogLoss", loss_weight=1.0),  ## only applies silog loss
+            dict(
+                type="PointmapIntrinsicsConsistencyLoss",
+                loss_weight=1.0,
+            ),
+            dict(
+                type="PointmapShiftInvariantL1Loss",
+                loss_weight=1.0,
+            ),
+            dict(type="PointmapNormalLoss", loss_weight=2.0),
+            dict(
+                type="PointmapScaleL1Loss", loss_weight=4.0
+            ),  ## Canonical XYZ = scale * XYZ
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=2.0, norm_type=2.0)

sapiens/dense/configs/pointmap/render_people/sapiens2_1b_pointmap_render_people-1024x768.py

@@ -0,0 +1,319 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 3, global bs: 768. num samples: 1e6. 1e6/768 = 1302. 1 epoch = 1e3 iters.
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_1b"
+embed_dim = 1536
+num_layers = 40
+num_heads = 24
+layer_decay_rate = 0.9
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_1b_pretrain.safetensors"
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+
+patch_size = 16
+num_tokens = (image_size[0] // patch_size) * (image_size[1] // patch_size)
+canonical_focal_length = 768.0
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            # mixed_precision=dict(
+            #     param_dtype="bf16",
+            #     reduce_dtype="bf16",
+            # ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="PointmapVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(
+        type="PointmapRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="PointmapResize", height=1024, width=768),
+    ## target is same res as output, otherwise we get artifacts.
+    dict(
+        type="PointmapGenerateTarget",
+        canonical_focal_length=canonical_focal_length,
+        target_downsample_factor=1,
+    ),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "flip",
+            "flip_direction",
+            "original_K",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="PointmapResize", height=1024, width=768),
+    dict(type="PointmapGenerateTarget", canonical_focal_length=canonical_focal_length),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "padding_size",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="PointmapResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "padding_size",
+        ),
+    ),
+]
+
+render_people_dataset = dict(
+    type="PointmapRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="PointmapRenderPeopleDataset",
+        # num_samples=100,  ## debug: only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="PointmapEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="PointmapEstimator",
+    canonical_focal_length=canonical_focal_length,
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="PointmapHead",
+        in_channels=embed_dim,
+        upsample_channels=[1536, 768, 512, 256],
+        conv_out_channels=[64, 32, 16],
+        conv_kernel_sizes=[3, 3, 3],
+        scale_conv_out_channels=(1536, 512, 128),
+        scale_conv_kernel_sizes=(1, 1, 1),
+        scale_final_layer=(
+            (num_tokens // ((2 * 2 * 2) * (2 * 2 * 2))) * 128,
+            512,
+            128,
+            1,
+        ),  ## scale regress
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),  ## on pointmap, XYZ
+            dict(
+                type="MultiscaleL1Loss",
+                loss_weight=1.0,
+                scale_factor=2,
+            ),
+            dict(type="SiLogLoss", loss_weight=1.0),  ## only applies silog loss
+            dict(
+                type="PointmapIntrinsicsConsistencyLoss",
+                loss_weight=1.0,
+            ),
+            dict(
+                type="PointmapShiftInvariantL1Loss",
+                loss_weight=1.0,
+            ),
+            dict(type="PointmapNormalLoss", loss_weight=2.0),
+            dict(
+                type="PointmapScaleL1Loss", loss_weight=4.0
+            ),  ## Canonical XYZ = scale * XYZ
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/pointmap/render_people/sapiens2_5b_pointmap_render_people-1024x768.py

@@ -0,0 +1,329 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 4e4  ## 32 nodes, 8 gpus: 256 gpus. bs: 1, global bs: 256. num samples: 1e6. 1e6/256 = 3906. 1 epoch = 3906 iters.
+
+## debug
+# warmup_iters = 100
+# num_iters = 300
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 10
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_5b"
+embed_dim = 2432
+num_layers = 56
+num_heads = 32
+layer_decay_rate = 0.94
+
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_5b_pretrain.safetensors"
+
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+
+patch_size = 16
+num_tokens = (image_size[0] // patch_size) * (image_size[1] // patch_size)
+canonical_focal_length = 768.0
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            # mixed_precision=dict(
+            #     param_dtype="bf16",
+            #     reduce_dtype="bf16",
+            # ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="PointmapVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(type="PhotoMetricDistortion"),
+    dict(
+        type="PointmapRandomScale",
+        scale_min=0.5,
+        scale_max=2.0,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomCropContinuous",
+        ar_range=(0.5, 2.0),
+        area_range=(0.4, 1.0),
+        num_attempts=8,
+        prob=0.3,
+    ),
+    dict(
+        type="PointmapRandomFlip",
+        prob=0.3,
+    ),
+    dict(type="PointmapResize", height=1024, width=768),
+    ## target is same res as output, otherwise we get artifacts.
+    dict(
+        type="PointmapGenerateTarget",
+        canonical_focal_length=canonical_focal_length,
+        target_downsample_factor=1,
+    ),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "ori_shape",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "flip",
+            "flip_direction",
+            "original_K",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+val_pipeline = [
+    dict(type="PointmapResize", height=1024, width=768),
+    dict(type="PointmapGenerateTarget", canonical_focal_length=canonical_focal_length),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "padding_size",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+test_pipeline = [
+    dict(type="PointmapResizePadImage", height=1024, width=768, pad_val=0),
+    dict(
+        type="PointmapPackInputs",
+        meta_keys=(
+            "img_path",
+            "orig_img_height",
+            "orig_img_width",
+            "img_shape",
+            "pad_shape",
+            "scale",
+            "padding_size",
+            "K",
+            "M",
+        ),
+    ),
+]
+
+render_people_dataset = dict(
+    type="PointmapRenderPeopleDataset",
+    data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2",
+)
+
+train_datasets = [render_people_dataset]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=False,
+    dataset=dict(
+        type="PointmapRenderPeopleDataset",
+        # num_samples=100,  ## only use N samples for validation
+        test_mode=True,
+        data_root=f"{_DATA_ROOT}/seg/data/render_people/synthetic_v2_test",
+        pipeline=val_pipeline,
+    ),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(
+        type="PointmapEvaluator",
+    ),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="PointmapEstimator",
+    canonical_focal_length=canonical_focal_length,
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="PointmapHead",
+        in_channels=embed_dim,
+        # upsample_channels=[1536, 768, 512, 256],
+        # conv_out_channels=[64, 32, 16],
+        # conv_kernel_sizes=[3, 3, 3],
+        upsample_channels=[1536, 768, 768, 768],  ## 1K resolution
+        conv_out_channels=[128, 64, 32],
+        conv_kernel_sizes=[3, 3, 3],
+        scale_conv_out_channels=(1536, 512, 128),
+        scale_conv_kernel_sizes=(1, 1, 1),
+        scale_final_layer=(
+            (num_tokens // ((2 * 2 * 2) * (2 * 2 * 2))) * 128,
+            512,
+            128,
+            1,
+        ),  ## scale regress
+        loss_decode=[
+            dict(type="L1Loss", loss_weight=2.0),  ## on pointmap, XYZ
+            dict(
+                type="MultiscaleL1Loss",
+                loss_weight=1.0,
+                scale_factor=2,
+            ),
+            dict(type="SiLogLoss", loss_weight=1.0),  ## only applies silog loss
+            dict(
+                type="PointmapIntrinsicsConsistencyLoss",
+                loss_weight=1.0,
+            ),
+            dict(
+                type="PointmapShiftInvariantL1Loss",
+                loss_weight=1.0,
+            ),
+            dict(type="PointmapNormalLoss", loss_weight=2.0),
+            dict(
+                type="PointmapScaleL1Loss", loss_weight=4.0
+            ),  ## Canonical XYZ = scale * XYZ
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    # lr=5e-4,
+    lr=1e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/seg/shutterstock_goliath/sapiens2_0.4b_seg_shutterstock_goliath-1024x768.py

@@ -0,0 +1,364 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # # # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.4b"
+embed_dim = 1024
+num_layers = 24
+num_heads = 16
+
+layer_decay_rate = 0.8
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.4b_pretrain.safetensors"
+
+num_classes = 29  ## 29 classes
+CLASS_WEIGHT = [
+    0.1,
+    10,
+    10,
+    3,
+    2,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    1,
+    10,
+    10,
+    10,
+    10,
+    10,
+]  ## 29 classes
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+# use_fsdp = True
+use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="SegVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+    class_palette_type="dome29",
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="SegRandomBackground",
+        prob=0.8,
+        skip_key="is_itw",
+        background_images_root=f"{_DATA_ROOT}/BG-20k/train",
+    ),
+    dict(
+        type="SegRandomResize",
+        base_height=1024,
+        base_width=768,
+        ratio_range=(0.4, 2.0),
+        keep_ratio=True,
+    ),
+    dict(
+        type="SegRandomCrop",
+        crop_height=1024,
+        crop_width=768,
+        prob=0.3,
+        cat_max_ratio=0.75,
+    ),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(
+        type="SegRandomRotate", prob=0.5, degree=60, seg_pad_val=0
+    ),  ## the black pixels are set as background
+    dict(
+        type="SegRandomHorizontalFlip",
+        prob=0.5,
+        swap_seg_labels=[
+            (5, 14),
+            (6, 15),
+            (7, 16),
+            (8, 17),
+            (9, 18),
+            (10, 19),
+            (11, 20),
+            (12, 21),
+        ],
+    ),  ## for the 29 classes,
+    dict(type="PhotoMetricDistortion"),
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False),
+    dict(type="SegPackInputs"),
+]
+
+val_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False, test_mode=True),
+    dict(type="SegPackInputs", test_mode=True),
+]
+
+test_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False, test_mode=True),
+    dict(type="SegPackInputs", test_mode=True),
+]
+
+##------------------------------------------------------------------------
+dataset_dome_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/sociopticon_body_segmentation_33_train:2024092600.json",
+)
+
+dataset_shutterstock_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_train:2024121600.json",
+)
+
+dataset_ca3_wide_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_wide_angle_body_segmentation_33_train:2024091700.json",
+)
+
+dataset_caa_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/cca_segmentation_33_train:2024092400.json",
+)
+
+dataset_ca3_zoom_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_zoom_in_body_segmentation_50_train:2024091700.json",
+)
+
+dataset_lighticon_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/lighticon_lightful_body_segmentation_51_train:2025021900.json",
+)
+
+dataset_internal_train = dict(
+    type="SegInternalClass29Dataset",
+    # ann_file=f"{_DATA_ROOT}/annotations/stylized_sapiens/20250807/Internal_segmentation_32:2025080700.json",
+    ann_file=f"{_DATA_ROOT}/annotations/internal_dataset/20251103/internal_keypoint_344_segmentation_32_train:2025091500.json",
+)
+
+train_datasets = [
+    dataset_dome_train,
+    dataset_ca3_wide_train,
+    dataset_caa_train,
+    dataset_ca3_zoom_train,
+    dataset_lighticon_train,
+    dataset_internal_train,
+] + 2 * [dataset_shutterstock_train]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=8,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",  ## avoids fork error with airstore
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="SegShutterstockClass29Dataset",
+        ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_test:2024121600.json",
+        test_mode=True,
+        pipeline=val_pipeline,
+    ),
+    collate_fn=dict(type="eval_collate"),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(type="SegEvaluator", class_names="dome29", nan_to_num=0.0),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="SegEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="SegHead",
+        in_channels=embed_dim,
+        deconv_out_channels=(
+            512,
+            256,
+            128,
+            64,
+        ),  ## this will 2x at each step. so total is 16x. 1K output.
+        deconv_kernel_sizes=(4, 4, 4, 4),
+        conv_out_channels=(64, 64),
+        conv_kernel_sizes=(1, 1),
+        num_classes=num_classes,
+        loss_decode=[
+            dict(
+                type="CrossEntropyLoss",
+                loss_weight=1.0,
+                reduction="none",
+                class_weight=CLASS_WEIGHT,
+                ignore_index=255,
+            ),
+            dict(
+                type="DiceLoss",
+                loss_weight=1.0,
+                reduction="none",
+                activate=True,
+                use_sigmoid=False,
+                include_background=False,
+                ignore_index=255,
+            ),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,  ## use fused AdamW
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=2.0, norm_type=2.0)

sapiens/dense/configs/seg/shutterstock_goliath/sapiens2_0.8b_seg_shutterstock_goliath-1024x768.py

@@ -0,0 +1,368 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 3e4  ## bs: 5; 16 gpus
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # # # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_0.8b"
+embed_dim = 1280
+num_layers = 32
+num_heads = 16
+
+layer_decay_rate = 0.85
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_0.8b_pretrain.safetensors"
+
+num_classes = 29  ## 29 classes
+CLASS_WEIGHT = [
+    0.1,
+    10,
+    10,
+    3,
+    2,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    1,
+    10,
+    10,
+    10,
+    10,
+    10,
+]  ## 29 classes
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+# use_fsdp = True
+use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+    ## Note: to merge sharded weight using FSDP
+    # accelerate merge-weights pytorch_model_fsdp_0/ .
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="SegVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+    class_palette_type="dome29",
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="SegRandomBackground",
+        prob=0.8,
+        skip_key="is_itw",
+        background_images_root=f"{_DATA_ROOT}/BG-20k/train",
+    ),
+    dict(
+        type="SegRandomResize",
+        base_height=1024,
+        base_width=768,
+        ratio_range=(0.4, 2.0),
+        keep_ratio=True,
+    ),
+    dict(
+        type="SegRandomCrop",
+        crop_height=1024,
+        crop_width=768,
+        prob=0.3,
+        cat_max_ratio=0.75,
+    ),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(
+        type="SegRandomRotate", prob=0.5, degree=60, seg_pad_val=0
+    ),  ## the black pixels are set as background
+    dict(
+        type="SegRandomHorizontalFlip",
+        prob=0.5,
+        swap_seg_labels=[
+            (5, 14),
+            (6, 15),
+            (7, 16),
+            (8, 17),
+            (9, 18),
+            (10, 19),
+            (11, 20),
+            (12, 21),
+        ],
+    ),  ## for the 29 classes,
+    dict(type="PhotoMetricDistortion"),
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False),
+    dict(type="SegPackInputs"),
+]
+
+val_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False, test_mode=True),
+    dict(type="SegPackInputs", test_mode=True),
+]
+
+test_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False, test_mode=True),
+    dict(type="SegPackInputs", test_mode=True),
+]
+
+##------------------------------------------------------------------------
+dataset_dome_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/sociopticon_body_segmentation_33_train:2024092600.json",
+)
+
+dataset_shutterstock_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_train:2024121600.json",
+)
+
+dataset_ca3_wide_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_wide_angle_body_segmentation_33_train:2024091700.json",
+)
+
+dataset_caa_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/cca_segmentation_33_train:2024092400.json",
+)
+
+dataset_ca3_zoom_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_zoom_in_body_segmentation_50_train:2024091700.json",
+)
+
+dataset_lighticon_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/lighticon_lightful_body_segmentation_51_train:2025021900.json",
+)
+
+dataset_internal_train = dict(
+    type="SegInternalClass29Dataset",
+    # ann_file=f"{_DATA_ROOT}/annotations/stylized_sapiens/20250807/Internal_segmentation_32:2025080700.json",
+    ann_file=f"{_DATA_ROOT}/annotations/internal_dataset/20251103/internal_keypoint_344_segmentation_32_train:2025091500.json",
+)
+
+train_datasets = [
+    dataset_dome_train,
+    dataset_ca3_wide_train,
+    dataset_caa_train,
+    dataset_ca3_zoom_train,
+    dataset_lighticon_train,
+    dataset_internal_train,
+] + 2 * [dataset_shutterstock_train]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",  ## avoids fork error with airstore
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="SegShutterstockClass29Dataset",
+        # num_samples=40,  ## only use N samples for validation
+        ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_test:2024121600.json",
+        test_mode=True,
+        pipeline=val_pipeline,
+    ),
+    collate_fn=dict(type="eval_collate"),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(type="SegEvaluator", class_names="dome29", nan_to_num=0.0),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="SegEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="SegHead",
+        in_channels=embed_dim,
+        deconv_out_channels=(
+            512,
+            256,
+            128,
+            64,
+        ),  ## this will 2x at each step. so total is 16x. 1K output.
+        deconv_kernel_sizes=(4, 4, 4, 4),
+        conv_out_channels=(64, 64),
+        conv_kernel_sizes=(1, 1),
+        num_classes=num_classes,
+        loss_decode=[
+            dict(
+                type="CrossEntropyLoss",
+                loss_weight=1.0,
+                reduction="none",
+                class_weight=CLASS_WEIGHT,
+                ignore_index=255,
+            ),
+            dict(
+                type="DiceLoss",
+                loss_weight=1.0,
+                reduction="none",
+                activate=True,
+                use_sigmoid=False,
+                include_background=False,
+                ignore_index=255,
+            ),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/seg/shutterstock_goliath/sapiens2_1b_seg_shutterstock_goliath-1024x768.py

@@ -0,0 +1,366 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 2e4
+# num_iters = 4e4
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 1000
+val_every_iters = 1000
+
+# # # # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 2
+# save_every_iters = 1000
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_1b"
+embed_dim = 1536
+num_layers = 40
+num_heads = 24
+
+layer_decay_rate = 0.9
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_1b_pretrain.safetensors"
+
+num_classes = 29  ## 29 classes
+CLASS_WEIGHT = [
+    0.1,
+    10,
+    10,
+    3,
+    2,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    1,
+    10,
+    10,
+    10,
+    10,
+    10,
+]  ## 29 classes
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="SegVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+    class_palette_type="dome29",
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="SegRandomBackground",
+        prob=0.8,
+        skip_key="is_itw",
+        background_images_root=f"{_DATA_ROOT}/BG-20k/train",
+    ),
+    dict(
+        type="SegRandomResize",
+        base_height=1024,
+        base_width=768,
+        ratio_range=(0.4, 2.0),
+        keep_ratio=True,
+    ),
+    dict(
+        type="SegRandomCrop",
+        crop_height=1024,
+        crop_width=768,
+        prob=0.3,
+        cat_max_ratio=0.75,
+    ),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(
+        type="SegRandomRotate", prob=0.5, degree=60, seg_pad_val=0
+    ),  ## the black pixels are set as background
+    dict(
+        type="SegRandomHorizontalFlip",
+        prob=0.5,
+        swap_seg_labels=[
+            (5, 14),
+            (6, 15),
+            (7, 16),
+            (8, 17),
+            (9, 18),
+            (10, 19),
+            (11, 20),
+            (12, 21),
+        ],
+    ),  ## for the 29 classes,
+    dict(type="PhotoMetricDistortion"),
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False),
+    dict(type="SegPackInputs"),
+]
+
+val_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False, test_mode=True),
+    dict(type="SegPackInputs", test_mode=True),
+]
+
+test_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False, test_mode=True),
+    dict(type="SegPackInputs", test_mode=True),
+]
+
+##------------------------------------------------------------------------
+dataset_dome_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/sociopticon_body_segmentation_33_train:2024092600.json",
+)
+
+dataset_shutterstock_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_train:2024121600.json",
+)
+
+dataset_ca3_wide_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_wide_angle_body_segmentation_33_train:2024091700.json",
+)
+
+dataset_caa_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/cca_segmentation_33_train:2024092400.json",
+)
+
+dataset_ca3_zoom_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_zoom_in_body_segmentation_50_train:2024091700.json",
+)
+
+dataset_lighticon_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/lighticon_lightful_body_segmentation_51_train:2025021900.json",
+)
+
+dataset_internal_train = dict(
+    type="SegInternalClass29Dataset",
+    # ann_file=f"{_DATA_ROOT}/annotations/stylized_sapiens/20250807/Internal_segmentation_32:2025080700.json",
+    ann_file=f"{_DATA_ROOT}/annotations/internal_dataset/20251103/internal_keypoint_344_segmentation_32_train:2025091500.json",
+)
+
+train_datasets = [
+    dataset_dome_train,
+    dataset_ca3_wide_train,
+    dataset_caa_train,
+    dataset_ca3_zoom_train,
+    dataset_lighticon_train,
+    dataset_internal_train,
+] + 2 * [dataset_shutterstock_train]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",  ## avoids fork error with airstore
+    # num_workers=0, # debug
+    # persistent_workers=False, # debug
+    shuffle=False,
+    dataset=dict(
+        type="SegShutterstockClass29Dataset",
+        # num_samples=40,  ## only use N samples for validation
+        ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_test:2024121600.json",
+        test_mode=True,
+        pipeline=val_pipeline,
+    ),
+    collate_fn=dict(type="eval_collate"),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(type="SegEvaluator", class_names="dome29", nan_to_num=0.0),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="SegEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="SegHead",
+        in_channels=embed_dim,
+        deconv_out_channels=(
+            512,
+            256,
+            128,
+            64,
+        ),  ## this will 2x at each step. so total is 16x. 1K output.
+        deconv_kernel_sizes=(4, 4, 4, 4),
+        conv_out_channels=(64, 64),
+        conv_kernel_sizes=(1, 1),
+        num_classes=num_classes,
+        loss_decode=[
+            dict(
+                type="CrossEntropyLoss",
+                loss_weight=1.0,
+                reduction="none",
+                class_weight=CLASS_WEIGHT,
+                ignore_index=255,
+            ),
+            dict(
+                type="DiceLoss",
+                loss_weight=1.0,
+                reduction="none",
+                activate=True,
+                use_sigmoid=False,
+                include_background=False,
+                ignore_index=255,
+            ),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/configs/seg/shutterstock_goliath/sapiens2_5b_seg_shutterstock_goliath-1024x768.py

@@ -0,0 +1,365 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+_CHECKPOINT_ROOT = os.path.expanduser(
+    os.environ.get("SAPIENS_CHECKPOINT_ROOT", "~/sapiens2_host")
+)
+_DATA_ROOT = os.path.expanduser(os.environ.get("DATA_ROOT", "~/sapiens_data"))
+
+warmup_iters = 500
+num_iters = 5e4  ## for h200; bs is 4
+
+# ------------------------------------------------------------------------------
+vis_every_iters = 100
+log_every_iters = 10
+save_every_iters = 2000
+# val_every_iters = 2000
+val_every_iters = 10000
+
+# # # # debug
+# vis_every_iters = 1
+# log_every_iters = 1
+# val_every_iters = 10
+
+load_from = None
+resume = False
+
+# ------------------------------------------------------------------
+model_name = "sapiens2_5b"
+embed_dim = 2432
+num_layers = 56
+num_heads = 32
+
+layer_decay_rate = 0.94
+pretrained_checkpoint = f"{_CHECKPOINT_ROOT}/pretrain/sapiens2_5b_pretrain.safetensors"
+
+num_classes = 29  ## 29 classes
+CLASS_WEIGHT = [
+    0.1,
+    10,
+    10,
+    3,
+    2,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    4,
+    4,
+    2,
+    2,
+    6,
+    10,
+    3,
+    3,
+    1,
+    1,
+    10,
+    10,
+    10,
+    10,
+    10,
+]  ## 29 classes
+
+##-----------------------------------------------------------------
+image_size = (1024, 768)  ## height x width
+patch_size = 16
+
+# ------------------------------------------------------------------
+use_fsdp = True
+# use_fsdp = False
+
+use_compile = True
+# use_compile = False
+
+## DDP config
+if use_fsdp is False:
+    accelerator_cfg = dict(
+        type="DDP",
+        log_with="tensorboard",
+        # find_unused_parameters=True,
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        # mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,  ## schedule independent of n_gpus
+    )
+
+else:
+    accelerator_cfg = dict(
+        type="FSDP",
+        log_with="tensorboard",
+        gradient_accumulation_steps=1,  # only accumulation=1 is supported. Otherwise, the LR scheduler will be off.
+        max_interval=num_iters,
+        mixed_precision="bf16",  # Options: ‘no’,‘fp16’,‘bf16’ or ‘fp8’.
+        step_scheduler_with_optimizer=False,
+        fsdp_cfg=dict(
+            fsdp_version=2,  # DTensor-based engine
+            state_dict_type="SHARDED_STATE_DICT",  # SHARDED_STATE_DICT | FULL_STATE_DICT
+            # state_dict_type="FULL_STATE_DICT",  # TODO: resume from this is not working
+            mixed_precision=dict(
+                param_dtype="bf16",
+                reduce_dtype="bf16",
+            ),
+            cpu_ram_efficient_loading=False,
+        ),
+    )
+
+    ## Note: to merge sharded weight using FSDP
+    # accelerate merge-weights pytorch_model_fsdp_0/ .
+
+if use_compile:
+    accelerator_cfg["compile_cfg"] = dict(
+        backend="inductor",
+        mode="default",  # Options: "default", "reduce-overhead", "max-autotune"
+        fullgraph=False,
+        dynamic=False,
+    )
+
+# ------------------------------------------------------------------
+randomness = dict(seed=0, deterministic=False, diff_rank_seed=True)
+logger = dict(
+    type="Logger",
+    log_interval=log_every_iters,
+)
+checkpoint = dict(
+    type="Checkpointer",
+    save_interval=save_every_iters,
+)
+
+visualizer = dict(
+    type="SegVisualizer",
+    vis_interval=vis_every_iters,
+    vis_max_samples=4,
+    vis_image_width=384,
+    vis_image_height=512,
+    class_palette_type="dome29",
+)
+
+
+##-----------------------------------------------------------------
+train_pipeline = [
+    dict(
+        type="SegRandomBackground",
+        prob=0.8,
+        skip_key="is_itw",
+        background_images_root=f"{_DATA_ROOT}/BG-20k/train",
+    ),
+    dict(
+        type="SegRandomResize",
+        base_height=1024,
+        base_width=768,
+        ratio_range=(0.4, 2.0),
+        keep_ratio=True,
+    ),
+    dict(
+        type="SegRandomCrop",
+        crop_height=1024,
+        crop_width=768,
+        prob=0.3,
+        cat_max_ratio=0.75,
+    ),
+    dict(
+        type="RandomGaussianBlur", prob=0.3, kernel_size=(3, 3), sigma_range=(0.1, 2.0)
+    ),
+    dict(type="RandomGaussianNoise", prob=0.3, var_range=(5.0, 20.0)),
+    dict(
+        type="SegRandomRotate", prob=0.5, degree=60, seg_pad_val=0
+    ),  ## the black pixels are set as background
+    dict(
+        type="SegRandomHorizontalFlip",
+        prob=0.5,
+        swap_seg_labels=[
+            (5, 14),
+            (6, 15),
+            (7, 16),
+            (8, 17),
+            (9, 18),
+            (10, 19),
+            (11, 20),
+            (12, 21),
+        ],
+    ),  ## for the 29 classes,
+    dict(type="PhotoMetricDistortion"),
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False),
+    dict(type="SegPackInputs"),
+]
+
+val_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False),
+    dict(type="SegPackInputs"),
+]
+
+test_pipeline = [
+    dict(type="SegResize", height=1024, width=768, keep_ratio=False),
+    dict(type="SegPackInputs"),
+]
+
+##------------------------------------------------------------------------
+dataset_dome_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/sociopticon_body_segmentation_33_train:2024092600.json",
+)
+
+dataset_shutterstock_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_train:2024121600.json",
+)
+
+dataset_ca3_wide_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_wide_angle_body_segmentation_33_train:2024091700.json",
+)
+
+dataset_caa_train = dict(
+    type="SegDomeClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/cca_segmentation_33_train:2024092400.json",
+)
+
+dataset_ca3_zoom_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/ca3_zoom_in_body_segmentation_50_train:2024091700.json",
+)
+
+dataset_lighticon_train = dict(
+    type="SegShutterstockClass29Dataset",
+    ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/lighticon_lightful_body_segmentation_51_train:2025021900.json",
+)
+
+dataset_internal_train = dict(
+    type="SegInternalClass29Dataset",
+    # ann_file=f"{_DATA_ROOT}/annotations/stylized_sapiens/20250807/Internal_segmentation_32:2025080700.json",
+    ann_file=f"{_DATA_ROOT}/annotations/internal_dataset/20251103/internal_keypoint_344_segmentation_32_train:2025091500.json",
+)
+
+train_datasets = [
+    dataset_dome_train,
+    dataset_ca3_wide_train,
+    dataset_caa_train,
+    dataset_ca3_zoom_train,
+    dataset_lighticon_train,
+    dataset_internal_train,
+] + 2 * [dataset_shutterstock_train]
+
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    shuffle=True,
+    dataset=dict(
+        type="CombinedDataset", datasets=train_datasets, pipeline=train_pipeline
+    ),
+)
+
+val_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    multiprocessing_context="spawn",  ## avoids fork error with airstore
+    shuffle=False,
+    dataset=dict(
+        type="SegShutterstockClass29Dataset",
+        ann_file=f"{_DATA_ROOT}/annotations/ingestion_90942/itw_shutterstock_body_segmentation_51_test:2024121600.json",
+        test_mode=True,
+        pipeline=val_pipeline,
+    ),
+    collate_fn=dict(type="eval_collate"),
+)
+
+val_cfg = dict(
+    val_interval=val_every_iters,
+    evaluator=dict(type="SegEvaluator", class_names="dome29", nan_to_num=0.0),
+)
+
+data_preprocessor = dict(
+    type="ImagePreprocessor",
+    mean=[123.675, 116.28, 103.53],
+    std=[58.395, 57.12, 57.375],
+    bgr_to_rgb=True,  ## convert from bgr to rgb for pretrained models
+)
+
+##-----------------------------------------------------------------
+model = dict(
+    type="SegEstimator",
+    backbone=dict(
+        type="Sapiens2",
+        arch=model_name,
+        img_size=image_size,
+        patch_size=patch_size,
+        final_norm=True,
+        use_tokenizer=False,
+        with_cls_token=True,
+        out_type="featmap",
+        init_cfg=dict(type="Pretrained", checkpoint=pretrained_checkpoint),
+    ),
+    decode_head=dict(
+        type="SegHead",
+        in_channels=embed_dim,
+        deconv_out_channels=(
+            512,
+            256,
+            128,
+            64,
+        ),  ## this will 2x at each step. so total is 16x. 1K output.
+        deconv_kernel_sizes=(4, 4, 4, 4),
+        conv_out_channels=(64, 64),
+        conv_kernel_sizes=(1, 1),
+        num_classes=num_classes,
+        loss_decode=[
+            dict(
+                type="CrossEntropyLoss",
+                loss_weight=1.0,
+                reduction="none",
+                class_weight=CLASS_WEIGHT,
+                ignore_index=255,
+            ),
+            dict(
+                type="DiceLoss",
+                loss_weight=1.0,
+                reduction="none",
+                activate=True,
+                use_sigmoid=False,
+                include_background=False,
+                ignore_index=255,
+            ),
+        ],
+    ),
+)
+
+
+##-----------------------------------------------------------------
+optimizer = dict(
+    type="AdamW",
+    lr=5e-4,
+    betas=(0.9, 0.999),
+    weight_decay=0.1,
+    paramwise_cfg=dict(
+        num_layers=num_layers,
+        layer_decay_rate=layer_decay_rate,
+    ),
+    fused=True,
+)
+
+scheduler = dict(
+    type="SequentialLR",
+    milestones=[warmup_iters],
+    schedulers=[
+        dict(type="LinearLR", start_factor=1e-3, total_iters=warmup_iters),
+        dict(
+            type="PolynomialLR",
+            total_iters=num_iters - warmup_iters,
+            power=1.0,
+        ),
+    ],
+)
+
+clip_grad = dict(mode="norm", max_norm=4.0, norm_type=2.0)

sapiens/dense/scripts/albedo/train/sapiens2_0.4b/node.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+
+cd "$(dirname "$(realpath "$0")")/../../../.." || exit
+
+#-------------------------------------------------------------------------------
+DEVICES=0,1,2,3,4,5,6,7
+# DEVICES=0
+
+#-------------------------------------------------------------------------------
+TASK="albedo"
+DATASET="render_people"
+MODEL="sapiens2_0.4b_${TASK}_${DATASET}-1024x768"
+
+CONFIG_FILE="configs/${TASK}/$DATASET/${MODEL}.py"
+TRAIN_BATCH_SIZE_PER_GPU=20
+
+#-------------------------------------------------------------------------------
+# mode='debug'
+mode='multi-gpu'
+
+#-------------------------------------------------------------------------------
+OUTPUT_DIR="Outputs/${TASK}/train/${MODEL}/node"
+OUTPUT_DIR="$(echo "${OUTPUT_DIR}/$(date +"%m-%d-%Y_%H:%M:%S")")"
+
+#-------------------------------------------------------------------------------
+OPTIONS="train_dataloader.batch_size=$TRAIN_BATCH_SIZE_PER_GPU"
+OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+CMD_RESUME="${RESUME_FROM:+--resume $RESUME_FROM}"
+
+export TF_CPP_MIN_LOG_LEVEL=2
+PORT=$(( ((RANDOM<<15)|RANDOM) % 63001 + 2000 ))
+
+#-------------------------------------------------------------------------------
+if [ "$mode" = "debug" ]; then
+    export TORCH_DISTRIBUTED_DEBUG=DETAIL
+    TRAIN_BATCH_SIZE_PER_GPU=1
+    OPTIONS="train_dataloader.batch_size=${TRAIN_BATCH_SIZE_PER_GPU} train_dataloader.num_workers=0 train_dataloader.persistent_workers=False"
+    OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} python tools/train.py ${CONFIG_FILE} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME}
+
+elif [ "$mode" = "multi-gpu" ]; then
+    NUM_GPUS=$(echo $DEVICES | tr -s ',' ' ' | wc -w)
+
+    LOG_FILE="${OUTPUT_DIR}/log.txt"
+    mkdir -p ${OUTPUT_DIR}
+    touch ${LOG_FILE}
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} PORT=${PORT} 'tools/dist_train.sh' ${CONFIG_FILE} \
+        ${NUM_GPUS} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME} \
+        | tee ${LOG_FILE}
+fi

sapiens/dense/scripts/albedo/train/sapiens2_0.8b/node.sh

@@ -0,0 +1,59 @@
+#!/bin/bash
+
+cd "$(dirname "$(realpath "$0")")/../../../.." || exit
+
+#-------------------------------------------------------------------------------
+DEVICES=0,1,2,3,4,5,6,7
+# DEVICES=0
+
+#-------------------------------------------------------------------------------
+TASK="albedo"
+DATASET="render_people"
+MODEL="sapiens2_0.8b_${TASK}_${DATASET}-1024x768"
+
+CONFIG_FILE="configs/${TASK}/$DATASET/${MODEL}.py"
+TRAIN_BATCH_SIZE_PER_GPU=12
+LOAD_FROM=''
+
+#-------------------------------------------------------------------------------
+# mode='debug'
+mode='multi-gpu'
+
+#-------------------------------------------------------------------------------
+OUTPUT_DIR="Outputs/${TASK}/train/${MODEL}/node"
+OUTPUT_DIR="$(echo "${OUTPUT_DIR}/$(date +"%m-%d-%Y_%H:%M:%S")")"
+
+#-------------------------------------------------------------------------------
+OPTIONS="train_dataloader.batch_size=$TRAIN_BATCH_SIZE_PER_GPU"
+OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+CMD_RESUME="${RESUME_FROM:+--resume $RESUME_FROM}"
+
+export TF_CPP_MIN_LOG_LEVEL=2
+PORT=$(( ((RANDOM<<15)|RANDOM) % 63001 + 2000 ))
+
+#-------------------------------------------------------------------------------
+if [ "$mode" = "debug" ]; then
+    export TORCH_DISTRIBUTED_DEBUG=DETAIL
+    TRAIN_BATCH_SIZE_PER_GPU=1
+    OPTIONS="train_dataloader.batch_size=${TRAIN_BATCH_SIZE_PER_GPU} train_dataloader.num_workers=0 train_dataloader.persistent_workers=False"
+    OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} python tools/train.py ${CONFIG_FILE} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME}
+
+elif [ "$mode" = "multi-gpu" ]; then
+    NUM_GPUS=$(echo $DEVICES | tr -s ',' ' ' | wc -w)
+
+    LOG_FILE="${OUTPUT_DIR}/log.txt"
+    mkdir -p ${OUTPUT_DIR}
+    touch ${LOG_FILE}
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} PORT=${PORT} 'tools/dist_train.sh' ${CONFIG_FILE} \
+        ${NUM_GPUS} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME} \
+        | tee ${LOG_FILE}
+fi

sapiens/dense/scripts/albedo/train/sapiens2_1b/node.sh

@@ -0,0 +1,59 @@
+#!/bin/bash
+
+cd "$(dirname "$(realpath "$0")")/../../../.." || exit
+
+#-------------------------------------------------------------------------------
+DEVICES=0,1,2,3,4,5,6,7
+# DEVICES=0
+
+#-------------------------------------------------------------------------------
+TASK="albedo"
+DATASET="render_people"
+MODEL="sapiens2_1b_${TASK}_${DATASET}-1024x768"
+
+CONFIG_FILE="configs/${TASK}/$DATASET/${MODEL}.py"
+TRAIN_BATCH_SIZE_PER_GPU=7
+
+
+#-------------------------------------------------------------------------------
+# mode='debug'
+mode='multi-gpu'
+
+#-------------------------------------------------------------------------------
+OUTPUT_DIR="Outputs/${TASK}/train/${MODEL}/node"
+OUTPUT_DIR="$(echo "${OUTPUT_DIR}/$(date +"%m-%d-%Y_%H:%M:%S")")"
+
+#-------------------------------------------------------------------------------
+OPTIONS="train_dataloader.batch_size=$TRAIN_BATCH_SIZE_PER_GPU"
+OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+CMD_RESUME="${RESUME_FROM:+--resume $RESUME_FROM}"
+
+export TF_CPP_MIN_LOG_LEVEL=2
+PORT=$(( ((RANDOM<<15)|RANDOM) % 63001 + 2000 ))
+
+#-------------------------------------------------------------------------------
+if [ "$mode" = "debug" ]; then
+    export TORCH_DISTRIBUTED_DEBUG=DETAIL
+    TRAIN_BATCH_SIZE_PER_GPU=1
+    OPTIONS="train_dataloader.batch_size=${TRAIN_BATCH_SIZE_PER_GPU} train_dataloader.num_workers=0 train_dataloader.persistent_workers=False"
+    OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} python tools/train.py ${CONFIG_FILE} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME}
+
+elif [ "$mode" = "multi-gpu" ]; then
+    NUM_GPUS=$(echo $DEVICES | tr -s ',' ' ' | wc -w)
+
+    LOG_FILE="${OUTPUT_DIR}/log.txt"
+    mkdir -p ${OUTPUT_DIR}
+    touch ${LOG_FILE}
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} PORT=${PORT} 'tools/dist_train.sh' ${CONFIG_FILE} \
+        ${NUM_GPUS} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME} \
+        | tee ${LOG_FILE}
+fi

sapiens/dense/scripts/albedo/train/sapiens2_5b/node.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+
+cd "$(dirname "$(realpath "$0")")/../../../.." || exit
+
+#-------------------------------------------------------------------------------
+DEVICES=0,1,2,3,4,5,6,7
+# DEVICES=0
+
+#-------------------------------------------------------------------------------
+TASK="albedo"
+DATASET="render_people"
+MODEL="sapiens2_5b_${TASK}_${DATASET}-1024x768"
+
+CONFIG_FILE="configs/${TASK}/$DATASET/${MODEL}.py"
+TRAIN_BATCH_SIZE_PER_GPU=3
+
+# LOAD_FROM=""
+
+#-------------------------------------------------------------------------------
+# mode='debug'
+mode='multi-gpu'
+
+#-------------------------------------------------------------------------------
+OUTPUT_DIR="Outputs/${TASK}/train/${MODEL}/node"
+OUTPUT_DIR="$(echo "${OUTPUT_DIR}/$(date +"%m-%d-%Y_%H:%M:%S")")"
+
+#-------------------------------------------------------------------------------
+OPTIONS="train_dataloader.batch_size=$TRAIN_BATCH_SIZE_PER_GPU"
+OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+CMD_RESUME="${RESUME_FROM:+--resume $RESUME_FROM}"
+
+export TF_CPP_MIN_LOG_LEVEL=2
+PORT=$(( ((RANDOM<<15)|RANDOM) % 63001 + 2000 ))
+
+#-------------------------------------------------------------------------------
+if [ "$mode" = "debug" ]; then
+    export TORCH_DISTRIBUTED_DEBUG=DETAIL
+    TRAIN_BATCH_SIZE_PER_GPU=1
+    OPTIONS="train_dataloader.batch_size=${TRAIN_BATCH_SIZE_PER_GPU} train_dataloader.num_workers=0 train_dataloader.persistent_workers=False"
+    OPTIONS="${OPTIONS}${LOAD_FROM:+ load_from=$LOAD_FROM}"
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} python tools/train.py ${CONFIG_FILE} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME}
+
+elif [ "$mode" = "multi-gpu" ]; then
+    NUM_GPUS=$(echo $DEVICES | tr -s ',' ' ' | wc -w)
+
+    LOG_FILE="${OUTPUT_DIR}/log.txt"
+    mkdir -p ${OUTPUT_DIR}
+    touch ${LOG_FILE}
+
+    CUDA_VISIBLE_DEVICES=${DEVICES} PORT=${PORT} 'tools/dist_train.sh' ${CONFIG_FILE} \
+        ${NUM_GPUS} \
+        --work-dir ${OUTPUT_DIR} \
+        --cfg-options ${OPTIONS} \
+        ${CMD_RESUME} \
+        | tee ${LOG_FILE}
+fi