Support `diffusers` and `stable-diffusion` (pretty much everything)

convert.py

@@ -2,18 +2,18 @@ import argparse
 import json
 import os
 import shutil
-from tempfile import TemporaryDirectory
 from collections import defaultdict
 from inspect import signature
-from typing import Optional, List
+from tempfile import TemporaryDirectory
+from typing import Dict, List, Optional, Set
 
 import torch
 
-from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, get_repo_discussions
+from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download
 from huggingface_hub.file_download import repo_folder_name
+from safetensors.torch import load_file, save_file
 from transformers import AutoConfig
 from transformers.pipelines.base import infer_framework_load_model
-from safetensors.torch import save_file
 
 
 class AlreadyExists(Exception):
@@ -30,15 +30,18 @@ def shared_pointers(tensors):
             failing.append(names)
     return failing
 
+
 def check_file_size(sf_filename: str, pt_filename: str):
     sf_size = os.stat(sf_filename).st_size
     pt_size = os.stat(pt_filename).st_size
 
     if (sf_size - pt_size) / pt_size > 0.01:
-        raise RuntimeError(f"""The file size different is more than 1%:
+        raise RuntimeError(
+            f"""The file size different is more than 1%:
          - {sf_filename}: {sf_size}
          - {pt_filename}: {pt_size}
-         """)
+         """
+        )
 
 
 def rename(pt_filename: str) -> str:
@@ -53,15 +56,14 @@ def convert_multi(model_id: str, folder: str) -> List["CommitOperationAdd"]:
         data = json.load(f)
 
     filenames = set(data["weight_map"].values())
+    local_filenames = []
     for filename in filenames:
-        cached_filename = hf_hub_download(repo_id=model_id, filename=filename)
-        loaded = torch.load(cached_filename)
-        sf_filename = rename(filename)
+        pt_filename = hf_hub_download(repo_id=model_id, filename=filename)
 
-        local = os.path.join(folder, sf_filename)
-        save_file(loaded, local, metadata={"format": "pt"})
-        check_file_size(local, cached_filename)
-        local_filenames.append(local)
+        sf_filename = rename(pt_filename)
+        sf_filename = os.path.join(folder, sf_filename)
+        convert_file(pt_filename, sf_filename)
+        local_filenames.append(sf_filename)
 
     index = os.path.join(folder, "model.safetensors.index.json")
     with open(index, "w") as f:
@@ -71,17 +73,28 @@ def convert_multi(model_id: str, folder: str) -> List["CommitOperationAdd"]:
         json.dump(newdata, f)
     local_filenames.append(index)
 
-    operations = [CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames]
+    operations = [
+        CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames
+    ]
 
     return operations
 
 
 def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:
-    sf_filename = "model.safetensors"
-    filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
-    loaded = torch.load(filename)
+    pt_filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
+
+    sf_name = "model.safetensors"
+    sf_filename = os.path.join(folder, sf_name)
+    convert_file(pt_filename, sf_filename)
+    operations = [CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)]
+    return operations
 
-    local = os.path.join(folder, sf_filename)
+
+def convert_file(
+    pt_filename: str,
+    sf_filename: str,
+):
+    loaded = torch.load(pt_filename)
     shared = shared_pointers(loaded)
     for shared_weights in shared:
         for name in shared_weights[1:]:
@@ -90,23 +103,45 @@ def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:
     # For tensors to be contiguous
     loaded = {k: v.contiguous() for k, v in loaded.items()}
 
-    save_file(loaded, local, metadata={"format": "pt"})
+    dirname = sf_filename.rsplit(os.path.sep, 1)[0]
+    os.makedirs(dirname, exist_ok=True)
+    save_file(loaded, sf_filename, metadata={"format": "pt"})
+    check_file_size(sf_filename, pt_filename)
+    reloaded = load_file(sf_filename)
+    for k in loaded:
+        pt_tensor = loaded[k]
+        sf_tensor = reloaded[k]
+        if not torch.equal(pt_tensor, sf_tensor):
+            raise RuntimeError(f"The output tensors do not match for key {k}")
 
-    check_file_size(local, filename)
 
-    operations = [CommitOperationAdd(path_in_repo=sf_filename, path_or_fileobj=local)]
-    return operations
+def create_diff(pt_infos: Dict[str, List[str]], sf_infos: Dict[str, List[str]]) -> str:
+    errors = []
+    for key in ["missing_keys", "mismatched_keys", "unexpected_keys"]:
+        pt_set = set(pt_infos[key])
+        sf_set = set(sf_infos[key])
+
+        pt_only = pt_set - sf_set
+        sf_only = sf_set - pt_set
+
+        if pt_only:
+            errors.append(f"{key} : PT warnings contain {pt_only} which are not present in SF warnings")
+        if sf_only:
+            errors.append(f"{key} : SF warnings contain {sf_only} which are not present in PT warnings")
+    return "\n".join(errors)
+
 
 def check_final_model(model_id: str, folder: str):
     config = hf_hub_download(repo_id=model_id, filename="config.json")
     shutil.copy(config, os.path.join(folder, "config.json"))
     config = AutoConfig.from_pretrained(folder)
 
-    _, pt_model = infer_framework_load_model(model_id, config)
-    _, sf_model = infer_framework_load_model(folder, config)
+    _, (pt_model, pt_infos) = infer_framework_load_model(model_id, config, output_loading_info=True)
+    _, (sf_model, sf_infos) = infer_framework_load_model(folder, config, output_loading_info=True)
 
-    pt_model = pt_model
-    sf_model = sf_model
+    if pt_infos != sf_infos:
+        error_string = create_diff(pt_infos, sf_infos)
+        raise ValueError(f"Different infos when reloading the model: {error_string}")
 
     pt_params = pt_model.state_dict()
     sf_params = sf_model.state_dict()
@@ -134,7 +169,6 @@ def check_final_model(model_id: str, folder: str):
     if "image" in sig.parameters:
         kwargs["image"] = pixel_values
 
-
     if torch.cuda.is_available():
         pt_model = pt_model.cuda()
         sf_model = sf_model.cuda()
@@ -146,6 +180,7 @@ def check_final_model(model_id: str, folder: str):
     torch.testing.assert_close(sf_logits, pt_logits)
     print(f"Model {model_id} is ok !")
 
+
 def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
     try:
         discussions = api.get_repo_discussions(repo_id=model_id)
@@ -156,7 +191,22 @@ def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discuss
             return discussion
 
 
-def convert(api: "HfApi", model_id: str, force: bool=False) -> Optional["CommitInfo"]:
+def convert_generic(model_id: str, folder: str, filenames: Set[str]) -> List["CommitOperationAdd"]:
+    operations = []
+
+    extensions = set([".bin", ".ckpt"])
+    for filename in filenames:
+        prefix, ext = os.path.splitext(filename)
+        if ext in extensions:
+            pt_filename = hf_hub_download(model_id, filename=filename)
+            sf_in_repo = f"{filename}.safetensors"
+            sf_filename = os.path.join(folder, sf_in_repo)
+            convert_file(pt_filename, sf_filename)
+            operations.append(CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename))
+    return operations
+
+
+def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:
     pr_title = "Adding `safetensors` variant of this model"
     info = api.model_info(model_id)
     filenames = set(s.rfilename for s in info.siblings)
@@ -174,21 +224,27 @@ def convert(api: "HfApi", model_id: str, force: bool=False) -> Optional["CommitI
                 url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
                 new_pr = pr
                 raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}")
-            elif "pytorch_model.bin" in filenames:
-                operations = convert_single(model_id, folder)
-            elif "pytorch_model.bin.index.json" in filenames:
-                operations = convert_multi(model_id, folder)
+            elif info.library_name == "transformers":
+                if "pytorch_model.bin" in filenames:
+                    operations = convert_single(model_id, folder)
+                elif "pytorch_model.bin.index.json" in filenames:
+                    operations = convert_multi(model_id, folder)
+                else:
+                    raise RuntimeError(f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert")
+                check_final_model(model_id, folder)
             else:
-                raise RuntimeError(f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert")
+                operations = convert_generic(model_id, folder, filenames)
 
             if operations:
-                check_final_model(model_id, folder)
                 new_pr = api.create_commit(
                     repo_id=model_id,
                     operations=operations,
                     commit_message=pr_title,
                     create_pr=True,
                 )
+                print(f"Pr created at {new_pr.pr_url}")
+            else:
+                print("No files to convert")
         finally:
             shutil.rmtree(folder)
         return new_pr