tipsv2/update_repo.py

"""Merge backbone vision encoder weights into DPT repo safetensors files.

For each DPT repo, downloads the current model.safetensors and the corresponding
backbone's model.safetensors, prepends the backbone's vision_encoder.* weights,
and pushes the merged file back to the DPT repo.
"""

import json
from huggingface_hub import (
    delete_file,
    hf_hub_download,
    list_repo_files,
    upload_file,
    upload_folder,
)
from safetensors.torch import load_file, save_file
from transformers import (
    AutoConfig,
    Tipsv2Config,
    Tipsv2DptConfig,
    Tipsv2DptImageProcessor,
    Tipsv2DptForDensePrediction,
    Tipsv2ImageProcessor,
    Tipsv2Model,
    Tipsv2Processor,
    Tipsv2Tokenizer,
)
import tempfile
from pathlib import Path

# (dpt_repo, backbone_repo, vision_fn). vision_fn keys into VISION_PRESETS and is carried here
# rather than read from the backbone config, because once a config has been converted to the
# transformers format it no longer stores the raw vision_fn key.
REPOS = [
    ("guarin/tipsv2-b14-dpt",       "guarin/tipsv2-b14",       "vit_base"),
    ("guarin/tipsv2-l14-dpt",       "guarin/tipsv2-l14",       "vit_large"),
    ("guarin/tipsv2-so400m14-dpt",  "guarin/tipsv2-so400m14",  "vit_so400m"),
    ("guarin/tipsv2-g14-dpt",       "guarin/tipsv2-g14",       "vit_giant2"),
]

# Remote code files (saved next to this script) pushed to each DPT repo so the
# trust_remote_code implementation matches the new config.
REMOTE_CODE_FILES = [
    "configuration_dpt.py",
    "modeling_dpt.py",
    "dpt_head.py",
    "image_encoder.py",
]

# Remote code files (saved next to this script) pushed to each backbone repo so the
# trust_remote_code implementation matches the new config.
BACKBONE_REMOTE_CODE_FILES = [
    "configuration_tips.py",
]

# trust_remote_code auto_map carried over from the backbone repos' original config so the pushed
# config keeps pointing at the custom configuration/modeling classes.
BACKBONE_AUTO_MAP = {
    "AutoConfig": "configuration_tips.TIPSv2Config",
    "AutoModel": "modeling_tips.TIPSv2Model",
}

# Per-architecture vision encoder sizes, keyed by the backbone's vision_fn. block_indices select
# which backbone stages feed the DPT neck.
VISION_PRESETS = {
    "vit_base":   {"num_hidden_layers": 12, "num_attention_heads": 12, "mlp_ratio": 4,             "block_indices": [2, 5, 8, 11]},
    "vit_large":  {"num_hidden_layers": 24, "num_attention_heads": 16, "mlp_ratio": 4,             "block_indices": [5, 11, 17, 23]},
    "vit_so400m": {"num_hidden_layers": 27, "num_attention_heads": 16, "mlp_ratio": 4304 / 1152,   "block_indices": [6, 13, 20, 26]},
    "vit_giant2": {"num_hidden_layers": 40, "num_attention_heads": 24, "mlp_ratio": 4,             "block_indices": [9, 19, 29, 39]},
}


def merge_weights(dpt_repo: str, backbone_repo: str, vision_fn: str) -> None:
    print(f"\n=== {dpt_repo} ===")

    # print(f"  Downloading {dpt_repo}/model.safetensors ...")
    # dpt_path = hf_hub_download(dpt_repo, "model.safetensors")

    # print(f"  Downloading {backbone_repo}/model.safetensors ...")
    # backbone_path = hf_hub_download(backbone_repo, "model.safetensors")

    # print("  Loading weights ...")
    # dpt_weights = load_file(dpt_path)
    # backbone_weights = load_file(backbone_path)

    # vision_weights = {
    #     key: tensor
    #     for key, tensor in backbone_weights.items()
    #     if key.startswith("vision_encoder.")
    # }
    # print(f"  Vision encoder keys: {len(vision_weights)}")
    # print(f"  DPT head keys:       {len(dpt_weights)}")

    # merged = {**vision_weights, **dpt_weights}
    # print(f"  Merged total keys:   {len(merged)}")

    # with tempfile.TemporaryDirectory() as tmp_dir:
    #     out_path = Path(tmp_dir) / "model.safetensors"
    #     save_file(merged, out_path)
    #     print(f"  Pushing merged weights to {dpt_repo} ...")
    #     upload_file(
    #         path_or_fileobj=str(out_path),
    #         path_in_repo="model.safetensors",
    #         repo_id=dpt_repo,
    #         repo_type="model",
    #         commit_message="Add vision encoder weights from backbone",
    #     )

    push_remote_code(dpt_repo)
    # push_model_config(dpt_repo, backbone_repo, vision_fn)
    # push_preprocessor_config(dpt_repo)
    # push_backbone_remote_code(backbone_repo)
    # push_backbone_model_config(backbone_repo, vision_fn)
    # push_backbone_preprocessor_config(backbone_repo)
    print(f"  Done.")


def push_remote_code(dpt_repo: str) -> None:
    this_dir = Path(__file__).parent
    for file_name in REMOTE_CODE_FILES:
        print(f"  Pushing {file_name} to {dpt_repo} ...")
        upload_file(
            path_or_fileobj=str(this_dir / file_name),
            path_in_repo=file_name,
            repo_id=dpt_repo,
            repo_type="model",
            commit_message=f"Update {file_name} for new config",
        )


def push_backbone_remote_code(backbone_repo: str) -> None:
    this_dir = Path(__file__).parent
    for file_name in BACKBONE_REMOTE_CODE_FILES:
        print(f"  Pushing {file_name} to {backbone_repo} ...")
        upload_file(
            path_or_fileobj=str(this_dir / file_name),
            path_in_repo=file_name,
            repo_id=backbone_repo,
            repo_type="model",
            commit_message=f"Update {file_name} for new config",
        )


def load_ade20k_labels() -> dict[int, str]:
    """Load the canonical ADE20K (150 classes) id -> label name mapping."""
    labels_path = hf_hub_download(
        "huggingface/label-files", "ade20k-id2label.json", repo_type="dataset"
    )
    id2label = json.loads(Path(labels_path).read_text())
    return {int(class_id): label_name for class_id, label_name in id2label.items()}


def sort_json_file(path: Path) -> None:
    """Rewrite a JSON file in place with all keys (recursively) in alphabetical order."""
    data = json.loads(path.read_text())
    path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n")


def build_vision_config_dict(backbone_repo: str, vision_fn: str) -> tuple[dict, dict]:
    """Rebuild the vision encoder config from a backbone repo, keyed by its architecture preset.

    Returns the raw backbone config dict and the corresponding Tipsv2VisionConfig dict. The layer
    geometry comes from VISION_PRESETS[vision_fn] and is always applied. The remaining values
    (hidden size, register tokens, image/patch size, ...) are only available from a raw
    pre-conversion backbone config; an already-converted config carries them correctly so the
    loaded values are kept.
    """
    backbone_dict = json.loads(Path(hf_hub_download(backbone_repo, "config.json")).read_text())
    preset = VISION_PRESETS[vision_fn]

    vision_config = AutoConfig.from_pretrained(backbone_repo).vision_config.to_dict()
    vision_config["num_hidden_layers"] = preset["num_hidden_layers"]
    vision_config["num_attention_heads"] = preset["num_attention_heads"]
    vision_config["mlp_ratio"] = preset["mlp_ratio"]

    if "vision_fn" in backbone_dict:
        vision_config["hidden_size"] = backbone_dict["embed_dim"]
        vision_config["use_swiglu_ffn"] = backbone_dict["ffn_layer"] in ("swiglu", "swiglufused")
        vision_config["num_register_tokens"] = backbone_dict["num_register_tokens"]
        vision_config["layerscale_value"] = backbone_dict.get("init_values", 1.0)
        vision_config["image_size"] = backbone_dict["img_size"]
        vision_config["patch_size"] = backbone_dict["patch_size"]

    vision_config.pop("out_features")
    vision_config.pop("out_indices")
    vision_config.pop("stage_names")

    return backbone_dict, vision_config


def push_backbone_model_config(backbone_repo: str, vision_fn: str) -> None:
    print(f"  Building config for {backbone_repo} ...")
    backbone_dict, vision_config = build_vision_config_dict(backbone_repo, vision_fn)

    loaded_config = AutoConfig.from_pretrained(backbone_repo)
    text_config = loaded_config.text_config.to_dict()
    temperature_init_value = loaded_config.temperature_init_value

    # Already-converted configs carry these directly; only raw configs need the overrides.
    if "vision_fn" in backbone_dict:
        text_config["vocab_size"] = backbone_dict["vocab_size"]
        text_config["hidden_size"] = backbone_dict["text_hidden_size"]
        text_config["intermediate_size"] = backbone_dict["text_mlp_dim"]
        text_config["num_attention_heads"] = backbone_dict["text_num_heads"]
        text_config["num_hidden_layers"] = backbone_dict["text_num_layers"]
        text_config["max_position_embeddings"] = backbone_dict["max_len"]
        temperature_init_value = backbone_dict["temperature"]

    config = Tipsv2Config(
        vision_config=vision_config,
        text_config=text_config,
        temperature_init_value=temperature_init_value,
    )
    config.architectures = ["TIPSv2Model"]

    print(f"  Initializing model ...")
    model = Tipsv2Model(config)

    with tempfile.TemporaryDirectory() as tmp_dir:
        model.config.save_pretrained(tmp_dir)
        out_path = Path(tmp_dir) / "config.json"

        # Preserve the trust_remote_code auto_map from the original config so the repo keeps
        # pointing at its custom configuration/modeling classes.
        saved_config = json.loads(out_path.read_text())
        saved_config["auto_map"] = BACKBONE_AUTO_MAP
        out_path.write_text(json.dumps(saved_config, indent=2, sort_keys=True) + "\n")

        print(f"  Pushing config.json to {backbone_repo} ...")
        upload_file(
            path_or_fileobj=str(out_path),
            path_in_repo="config.json",
            repo_id=backbone_repo,
            repo_type="model",
            commit_message="Add model config",
        )


def push_model_config(dpt_repo: str, backbone_repo: str, vision_fn: str) -> None:
    print(f"  Building config for {dpt_repo} ...")
    config_path = hf_hub_download(dpt_repo, "config.json")
    config_dict = json.loads(Path(config_path).read_text())

    # The model is trained on ADE20K, so the placeholder "LABEL_0, LABEL_1, ..." names are
    # replaced with the real 150 ADE20K class names.
    id2label = load_ade20k_labels()
    config_dict["id2label"] = {
        str(class_id): label_name for class_id, label_name in id2label.items()
    }
    config_dict["label2id"] = {label_name: class_id for class_id, label_name in id2label.items()}

    # The stored config does not carry a backbone_config, so the vision encoder config is rebuilt
    # from the backbone repo. block_indices select which backbone stages feed the DPT neck (the +1
    # offset accounts for the embedding output that precedes the transformer stages).
    backbone_dict, backbone_config = build_vision_config_dict(backbone_repo, vision_fn)
    preset = VISION_PRESETS[vision_fn]
    num_hidden_layers = preset["num_hidden_layers"]
    out_indices = [block_index + 1 for block_index in preset["block_indices"]]

    stage_names = ["stem"] + [f"stage{index}" for index in range(1, num_hidden_layers + 1)]
    backbone_config["stage_names"] = stage_names
    backbone_config["out_indices"] = out_indices
    backbone_config["out_features"] = [stage_names[out_index] for out_index in out_indices]
    backbone_config["apply_layernorm"] = True
    backbone_config["reshape_hidden_states"] = False
    config_dict["backbone_config"] = backbone_config

    config = Tipsv2DptConfig.from_dict(config_dict)

    print(f"  Initializing model ...")
    model = Tipsv2DptForDensePrediction(config)

    with tempfile.TemporaryDirectory() as tmp_dir:
        model.config.save_pretrained(tmp_dir)
        out_path = Path(tmp_dir) / "config.json"
        sort_json_file(out_path)
        print(f"  Pushing config.json to {dpt_repo} ...")
        upload_file(
            path_or_fileobj=str(out_path),
            path_in_repo="config.json",
            repo_id=dpt_repo,
            repo_type="model",
            commit_message="Add model config",
        )


def push_preprocessor_config(dpt_repo: str) -> None:
    print(f"  Creating preprocessor_config.json ...")
    image_processor = Tipsv2DptImageProcessor()

    with tempfile.TemporaryDirectory() as tmp_dir:
        image_processor.save_pretrained(tmp_dir)
        config_path = Path(tmp_dir) / "preprocessor_config.json"
        sort_json_file(config_path)
        print(f"  Pushing preprocessor_config.json to {dpt_repo} ...")
        upload_file(
            path_or_fileobj=str(config_path),
            path_in_repo="preprocessor_config.json",
            repo_id=dpt_repo,
            repo_type="model",
            commit_message="Add preprocessor_config.json",
        )


def push_backbone_preprocessor_config(backbone_repo: str) -> None:
    print(f"  Creating processor config for {backbone_repo} ...")

    print(f"  Downloading {backbone_repo}/tokenizer.model ...")
    vocab_path = hf_hub_download(backbone_repo, "tokenizer.model")

    image_processor = Tipsv2ImageProcessor()
    tokenizer = Tipsv2Tokenizer(vocab_file=vocab_path)
    processor = Tipsv2Processor(image_processor=image_processor, tokenizer=tokenizer)

    with tempfile.TemporaryDirectory() as tmp_dir:
        processor.save_pretrained(tmp_dir)
        # tokenizer.json must not be published, so it is removed before upload.
        tokenizer_json_path = Path(tmp_dir) / "tokenizer.json"
        if tokenizer_json_path.exists():
            tokenizer_json_path.unlink()
        for json_path in Path(tmp_dir).glob("*.json"):
            sort_json_file(json_path)
        generated_files = sorted(path.name for path in Path(tmp_dir).iterdir())
        print(f"  Generated config files: {generated_files}")
        print(f"  Pushing processor config files to {backbone_repo} ...")
        upload_folder(
            folder_path=tmp_dir,
            repo_id=backbone_repo,
            repo_type="model",
            commit_message="Add processor config files",
        )

    repo_files = list_repo_files(backbone_repo, repo_type="model")

    if "preprocessor_config.json" in repo_files:
        print(f"  Deleting obsolete preprocessor_config.json from {backbone_repo} ...")
        delete_file(
            path_in_repo="preprocessor_config.json",
            repo_id=backbone_repo,
            repo_type="model",
            commit_message="Remove obsolete preprocessor_config.json (merged into processor_config.json)",
        )

    if "tokenizer.json" in repo_files:
        print(f"  Deleting tokenizer.json from {backbone_repo} ...")
        delete_file(
            path_in_repo="tokenizer.json",
            repo_id=backbone_repo,
            repo_type="model",
            commit_message="Remove tokenizer.json",
        )


def main() -> None:
    for dpt_repo, backbone_repo, vision_fn in REPOS:
        merge_weights(dpt_repo, backbone_repo, vision_fn)


if __name__ == "__main__":
    main()