switch to safetensors

app.py

@@ -13,6 +13,7 @@ import spaces
 from chord import ChordModel
 from chord.module import make
 from chord.util import get_positions, rgb_to_srgb
+from chord.io import load_torch_file
 
 EXAMPLES_USECASE_1 = [
     [f"examples/generated/{f}"] 
@@ -29,14 +30,14 @@ EXAMPLES_USECASE_3 = [
 
 MODEL_OBJ = None
 login(token=os.environ["HF_TOKEN"])
-MODEL_CKPT_PATH = hf_hub_download(repo_id="Ubisoft/ubisoft-laforge-chord", filename="chord_v1.ckpt")
+MODEL_CKPT_PATH = hf_hub_download(repo_id="Ubisoft/ubisoft-laforge-chord", filename="chord_v1.safetensors")
 def load_model(ckpt_path):
     print("Loading model from:", ckpt_path)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     config = OmegaConf.load("config/chord.yaml")
     model = ChordModel(config)
-    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-    model.load_state_dict(ckpt["state_dict"])
+    state_dict = load_torch_file(ckpt_path)
+    model.load_state_dict(state_dict)
     model.eval()
     model.to(device)
     return model

chord/io.py

@@ -3,6 +3,7 @@ import imageio.v3 as imageio
 import numpy as np
 import warnings
 import os
+import safetensors
 
 import torchvision.transforms.functional as F
 
@@ -77,4 +78,23 @@ def save_maps(path: str, maps: dict):
         os.makedirs(path)
     for name, image in maps.items():
         out_img = create_img(image)
-        out_img.save(os.path.join(path, name+".png"))
+        out_img.save(os.path.join(path, name+".png"))
+
+def load_torch_file(ckpt, device=None):
+    if device is None:
+        device = torch.device("cpu")
+    if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"):
+        with safetensors.safe_open(ckpt, framework="pt", device=device.type) as f:
+            state_dict = {}
+            for k in f.keys():
+                tensor = f.get_tensor(k)
+                state_dict[k] = tensor
+    else:
+        torch_args = {}
+        ckpt = torch.load(ckpt, map_location=device, weights_only=True, **torch_args)
+
+        if "state_dict" in ckpt:
+            state_dict = ckpt["state_dict"]
+        else:
+            state_dict = ckpt
+    return state_dict