feat(checkpoints): support #fname query in HTTPStorage

Author: gadicc

api/convert_to_diffusers.py

@@ -22,14 +22,15 @@ def main(
     checkpoint_url: str,
     checkpoint_config_url: str,
     checkpoint_args: dict = {},
+    path=None,
 ):
-    fname = CHECKPOINT_DIR + "/" + checkpoint_url.split("/").pop()
+    if not path:
+        fname = checkpoint_url.split("/").pop()
+        path = os.path.join(CHECKPOINT_DIR, fname)
 
     if checkpoint_config_url and checkpoint_config_url != "":
         storage = Storage(checkpoint_config_url)
-        configPath = (
-            CHECKPOINT_DIR + "/" + checkpoint_url.split("/").pop() + "_config.yaml"
-        )
+        configPath = CHECKPOINT_DIR + "/" + path + "_config.yaml"
         print(f"Downloading {checkpoint_config_url} to {configPath}...")
         storage.download_file(configPath)
 
@@ -47,7 +48,7 @@ def main(
     #     "./diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py"
     # )
 
-    print("Converting " + fname + " to diffusers model " + model_id + "...", flush=True)
+    print("Converting " + path + " to diffusers model " + model_id + "...", flush=True)
 
     # These are now in main requirements.txt.
     # subprocess.run(
@@ -112,11 +113,11 @@ def main(
     # our defaults
     args.update(
         {
-            "checkpoint_path": fname,
+            "checkpoint_path": path,
             "original_config_file": configPath if checkpoint_config_url else None,
             "device": device_id,
             "extract_ema": True,
-            "from_safetensors": "safetensor" in fname.lower(),
+            "from_safetensors": "safetensor" in path.lower(),
         }
     )
 

api/download.py

@@ -80,11 +80,12 @@ async def download_model(
             await asyncio.to_thread(storage.download_and_extract, model_file, model_dir)
         else:
             if checkpoint_url:
-                download_checkpoint(checkpoint_url)
+                path = download_checkpoint(checkpoint_url)
                 convert_to_diffusers(
                     model_id=model_id,
                     checkpoint_url=checkpoint_url,
                     checkpoint_config_url=checkpoint_config_url,
+                    path=path,
                 )
             else:
                 print("Does not exist, let's try find it on huggingface")
@@ -138,11 +139,12 @@ async def download_model(
 
     else:
         if checkpoint_url:
-            download_checkpoint(checkpoint_url)
+            path = download_checkpoint(checkpoint_url)
             convert_to_diffusers(
                 model_id=model_id,
                 checkpoint_url=checkpoint_url,
                 checkpoint_config_url=checkpoint_config_url,
+                path=path,
             )
         else:
             # do a dry run of loading the huggingface model, which will download weights at build time

api/download_checkpoint.py

@@ -8,10 +8,19 @@
 def main(checkpoint_url: str):
     if not os.path.isdir(CHECKPOINT_DIR):
         os.makedirs(CHECKPOINT_DIR)
-    fname = CHECKPOINT_DIR + "/" + checkpoint_url.split("/").pop()
-    if not os.path.isfile(fname):
-        storage = Storage(checkpoint_url)
-        storage.download_file(fname)
+
+    storage = Storage(checkpoint_url)
+    storage_query_fname = storage.query.get("fname")
+    if storage_query_fname:
+        fname = storage_query_fname[0]
+    else:
+        fname = checkpoint_url.split("/").pop()
+    path = os.path.join(CHECKPOINT_DIR, fname)
+
+    if not os.path.isfile(path):
+        storage.download_file(path)
+
+    return path
 
 
 if __name__ == "__main__":