Upload convert_da3.py with huggingface_hub

convert_da3.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+"""Repackage an original Depth Anything 3 checkpoint into ComfyUI's native layout.
+
+Usage:
+    python scripts/convert_da3.py \
+        --input da3_original.safetensors \
+        --output models/diffusion_models/da3_comfy.safetensors
+
+This applies, *offline*, the exact transform that ComfyUI used to do at load
+time for DA3:
+
+  * remap the DINOv2 backbone keys ``backbone.pretrained.*`` (upstream DA3)
+    to the ``Dinov2Model`` runtime layout (``backbone.embeddings.*``,
+    ``backbone.encoder.layer.*``, ``backbone.layernorm.*``);
+  * split each fused ``attn.qkv`` projection into separate query/key/value
+    linears;
+  * drop the unused Gaussian-splat head weights (``gs_head.*``, ``gs_adapter.*``).
+
+The head (``head.*``), camera encoder/decoder (``cam_enc.*``, ``cam_dec.*``)
+and any other keys are passed through unchanged. After conversion the file
+loads directly via ComfyUI auto-detection with no in-code remap.
+"""
+
+import argparse
+import glob
+import os
+
+import torch
+
+from safetensors.torch import load_file, save_file
+
+DROP_PREFIXES = ("gs_head.", "gs_adapter.")
+
+
+def remap_backbone_keys(state_dict, prefix="backbone."):
+    """Map ``backbone.pretrained.*`` (upstream DA3) keys to ``Dinov2Model`` layout."""
+    pre = prefix + "pretrained."
+    src_keys = [k for k in state_dict.keys() if k.startswith(pre)]
+    if not src_keys:
+        return state_dict
+
+    static_renames = {
+        pre + "patch_embed.proj.weight":  prefix + "embeddings.patch_embeddings.projection.weight",
+        pre + "patch_embed.proj.bias":    prefix + "embeddings.patch_embeddings.projection.bias",
+        pre + "pos_embed":                prefix + "embeddings.position_embeddings",
+        pre + "cls_token":                prefix + "embeddings.cls_token",
+        pre + "camera_token":             prefix + "embeddings.camera_token",
+        pre + "norm.weight":              prefix + "layernorm.weight",
+        pre + "norm.bias":                prefix + "layernorm.bias",
+    }
+    for src, dst in static_renames.items():
+        if src in state_dict:
+            state_dict[dst] = state_dict.pop(src)
+
+    block_pre = pre + "blocks."
+    block_keys = [k for k in state_dict.keys() if k.startswith(block_pre)]
+    for k in block_keys:
+        rest = k[len(block_pre):]                 # e.g. "5.attn.qkv.weight"
+        idx_str, _, sub = rest.partition(".")
+        target_block = "{}encoder.layer.{}.".format(prefix, idx_str)
+
+        # Fused QKV -> split query/key/value linears.
+        if sub == "attn.qkv.weight":
+            qkv = state_dict.pop(k)
+            c = qkv.shape[0] // 3
+            state_dict[target_block + "attention.attention.query.weight"] = qkv[:c].clone()
+            state_dict[target_block + "attention.attention.key.weight"]   = qkv[c:2 * c].clone()
+            state_dict[target_block + "attention.attention.value.weight"] = qkv[2 * c:].clone()
+            continue
+        if sub == "attn.qkv.bias":
+            qkv = state_dict.pop(k)
+            c = qkv.shape[0] // 3
+            state_dict[target_block + "attention.attention.query.bias"] = qkv[:c].clone()
+            state_dict[target_block + "attention.attention.key.bias"]   = qkv[c:2 * c].clone()
+            state_dict[target_block + "attention.attention.value.bias"] = qkv[2 * c:].clone()
+            continue
+
+        # Sub-key remap (suffix preserved).
+        if sub.startswith("attn.proj."):
+            tail = sub[len("attn.proj."):]
+            new = "attention.output.dense." + tail
+        elif sub.startswith("attn.q_norm."):
+            new = "attention.q_norm." + sub[len("attn.q_norm."):]
+        elif sub.startswith("attn.k_norm."):
+            new = "attention.k_norm." + sub[len("attn.k_norm."):]
+        elif sub == "ls1.gamma":
+            new = "layer_scale1.lambda1"
+        elif sub == "ls2.gamma":
+            new = "layer_scale2.lambda1"
+        elif sub.startswith("mlp.w12."):
+            new = "mlp.weights_in." + sub[len("mlp.w12."):]
+        elif sub.startswith("mlp.w3."):
+            new = "mlp.weights_out." + sub[len("mlp.w3."):]
+        elif sub.startswith(("norm1.", "norm2.", "mlp.fc1.", "mlp.fc2.")):
+            new = sub
+        else:
+            # Unrecognised key -- leave as-is so a later load can complain.
+            raise ValueError("Unrecognised DA3 backbone key: {}".format(k))
+
+        state_dict[target_block + new] = state_dict.pop(k)
+
+    return state_dict
+
+
+def drop_unused(state_dict):
+    for k in list(state_dict.keys()):
+        if k.startswith(DROP_PREFIXES):
+            state_dict.pop(k)
+    return state_dict
+
+
+def load_state_dict(path):
+    if os.path.isdir(path):
+        sd = {}
+        files = sorted(glob.glob(os.path.join(path, "*.safetensors")))
+        if not files:
+            raise FileNotFoundError("No .safetensors files in {}".format(path))
+        for f in files:
+            sd.update(load_file(f))
+        return sd
+
+    if path.endswith(".safetensors"):
+        return load_file(path)
+
+    sd = torch.load(path, map_location="cpu", weights_only=False)
+    # Unwrap common nesting (e.g. {"model": ...} / {"state_dict": ...}).
+    for wrap in ("state_dict", "model", "module"):
+        if isinstance(sd, dict) and wrap in sd and isinstance(sd[wrap], dict):
+            sd = sd[wrap]
+            break
+    return sd
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Repackage an original Depth Anything 3 checkpoint into ComfyUI's native layout"
+    )
+    parser.add_argument("--input", type=str, required=True,
+                        help="Path to original DA3 .safetensors / .pt / .pth file or directory")
+    parser.add_argument("--output", type=str, required=True,
+                        help="Output .safetensors file path")
+    args = parser.parse_args()
+
+    print("Loading: {}".format(args.input))
+    sd = load_state_dict(args.input)
+    print("  Loaded {} keys".format(len(sd)))
+
+    # Original DA3 checkpoints store everything under a "model." string prefix
+    # (e.g. "model.backbone.pretrained.*"). Strip it so the remap works on bare
+    # "backbone.*" keys, then re-add it at the end: ComfyUI's loader resolves the
+    # diffusion-model prefix to "model." for DA3, so the saved file must keep it.
+    if any(k.startswith("model.") for k in sd):
+        print('  Stripping "model." prefix for processing')
+        sd = {(k[len("model."):] if k.startswith("model.") else k): v for k, v in sd.items()}
+
+    if any(k.startswith("backbone.pretrained.") for k in sd):
+        print("  Remapping backbone (backbone.pretrained.* -> Dinov2Model layout)...")
+        sd = remap_backbone_keys(sd, prefix="backbone.")
+    elif any(k.startswith("backbone.embeddings.") for k in sd):
+        print("  Backbone already in ComfyUI layout, skipping remap.")
+    else:
+        raise ValueError("Input does not look like a DA3 checkpoint (no backbone.* keys found)")
+
+    n_before = len(sd)
+    sd = drop_unused(sd)
+    dropped = n_before - len(sd)
+    if dropped:
+        print("  Dropped {} unused Gaussian-head keys".format(dropped))
+
+    # Re-add the "model." prefix expected by ComfyUI's diffusion-model loader.
+    sd = {"model." + k: v for k, v in sd.items()}
+
+    # safetensors requires contiguous tensors; the qkv split slices are cloned
+    # above but enforce contiguity defensively for all tensors.
+    sd = {k: v.contiguous() for k, v in sd.items()}
+
+    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
+    save_file(sd, args.output)
+    print("  Saved {} keys to {}".format(len(sd), args.output))
+
+
+if __name__ == "__main__":
+    main()