utility script for pushing to HF hub

push_to_hf.py

@@ -1,42 +1,67 @@
+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 ScoreFlow
+from msma import EDMScorer, ScoreFlow, build_model_from_pickle
 
-basedir = Path("models/condgauss")
-preset = "edm2-img64-s-fid"
-modeldir = basedir / preset
 
-model = ScoreFlow(preset)
-model.flow.load_state_dict(torch.load(modeldir / "flow.pt"))
+@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
 
-api = HfApi()
-repo_name = "ahsanMah/localizing-edm"
+    net = build_model_from_pickle(preset)
+    model = ScoreFlow(
+        net,
+        num_flows=8,
+    )
+    model.flow.load_state_dict(torch.load(modeldir / "flow.pt"))
 
-# Create repo if not existing yet and get the associated repo_id
-repo_id = api.create_repo(repo_name, exist_ok=True).repo_id
+    api = HfApi()
+    repo_name = "ahsanMah/localizing-edm"
 
-# Save all files in a temporary directory and push them in a single commit
-with TemporaryDirectory() as tmpdir:
-    tmpdir = Path(tmpdir)
+    # 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 weights
-    save_file(model.state_dict(), tmpdir / "model.safetensors")
+    # Save all files in a temporary directory and push them in a single commit
+    with TemporaryDirectory() as tmpdir:
+        tmpdir = Path(tmpdir)
 
-    # Generate model card
-    # card = generate_model_card(model)
-    # (tmpdir / "README.md").write_text(card)
+        # Save weights
+        save_file(model.state_dict(), tmpdir / "model.safetensors")
 
-    # Save logs
-    shutil.copytree(modeldir / "logs", tmpdir / "logs")
-    # Save figures
-    # Save evaluation metrics
-    # ...
+        # save config
+        (tmpdir / "config.json").write_text(json.dumps(model.config, sort_keys=True, indent=4))
+        
+        # TODO: save gmm and cached score norms
 
-    # Push to hub
-    api.upload_folder(repo_id=repo_id, folder_path=tmpdir)
+
+        # 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()