push_to_hf.py

import json
import shutil
from pathlib import Path
from tempfile import TemporaryDirectory

import click
import torch
from huggingface_hub import HfApi
from safetensors.torch import save_file

from msma import EDMScorer, ScoreFlow, build_model_from_pickle


@click.command
@click.option(
    "--basedir",
    help="Directory holding the model weights and logs",
    type=str,
    required=True,
)
@click.option(
    "--preset", help="Preset of the score model used", type=str, required=True
)
def main(basedir, preset):
    basedir = Path(basedir)
    modeldir = basedir / preset

    net = build_model_from_pickle(preset)
    with open(modeldir / "config.json", "rb") as f:
        model_params = json.load(f)

    model = ScoreFlow(
        net,
        device="cpu",
        **model_params["PatchFlow"],
    )
    model.flow.load_state_dict(torch.load(modeldir / "flow.pt"))

    api = HfApi()

    # Use your own repo
    repo_name = "ahsanMah/localizing-edm"

    # Create repo if not existing yet and get the associated repo_id
    repo_id = api.create_repo(repo_name, exist_ok=True).repo_id

    # Save all files in a temporary directory and push them in a single commit
    with TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)

        # Save weights
        save_file(model.state_dict(), tmpdir / "model.safetensors")

        # save config
        (tmpdir / "config.json").write_text(
            json.dumps(model.config, sort_keys=True, indent=4)
        )

        # save gmm and cached score likelihoods
        shutil.copyfile(modeldir / "gmm.pkl", tmpdir / "gmm.pkl")
        shutil.copyfile(modeldir / "refscores.npz", tmpdir / "refscores.npz")

        # Generate model card
        # card = generate_model_card(model)
        # (tmpdir / "README.md").write_text(card)

        # Save logs
        shutil.copytree(modeldir / "logs", tmpdir / "logs")

        # Push to hub
        api.upload_folder(repo_id=repo_id, path_in_repo=preset, folder_path=tmpdir)


if __name__ == "__main__":
    main()