ahsanMah commited on
Commit
1b96548
·
1 Parent(s): cb54c64

added utility to push to hub

Browse files
Files changed (1) hide show
  1. push_to_hf.py +42 -0
push_to_hf.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import shutil
2
+ from pathlib import Path
3
+ from tempfile import TemporaryDirectory
4
+
5
+ import torch
6
+ from huggingface_hub import HfApi
7
+ from safetensors.torch import save_file
8
+
9
+ from msma import ScoreFlow
10
+
11
+ basedir = Path("models/condgauss")
12
+ preset = "edm2-img64-s-fid"
13
+ modeldir = basedir / preset
14
+
15
+ model = ScoreFlow(preset)
16
+ model.flow.load_state_dict(torch.load(modeldir / "flow.pt"))
17
+
18
+ api = HfApi()
19
+ repo_name = "ahsanMah/localizing-edm"
20
+
21
+ # Create repo if not existing yet and get the associated repo_id
22
+ repo_id = api.create_repo(repo_name, exist_ok=True).repo_id
23
+
24
+ # Save all files in a temporary directory and push them in a single commit
25
+ with TemporaryDirectory() as tmpdir:
26
+ tmpdir = Path(tmpdir)
27
+
28
+ # Save weights
29
+ save_file(model.state_dict(), tmpdir / "model.safetensors")
30
+
31
+ # Generate model card
32
+ # card = generate_model_card(model)
33
+ # (tmpdir / "README.md").write_text(card)
34
+
35
+ # Save logs
36
+ shutil.copytree(modeldir / "logs", tmpdir / "logs")
37
+ # Save figures
38
+ # Save evaluation metrics
39
+ # ...
40
+
41
+ # Push to hub
42
+ api.upload_folder(repo_id=repo_id, folder_path=tmpdir)