Merge pull request #27 from Stability-AI/os_hf_hub

Add huggingface hub support

Author: zqevans

run_gradio.py

@@ -1,3 +1,4 @@
+from stable_audio_tools import get_pretrained_model
 from stable_audio_tools.interface.gradio import create_ui
 import json 
 
@@ -6,21 +7,16 @@
 def main(args):
     torch.manual_seed(42)
 
-    print(f"Loading model config from {args.model_config}")
-    # Load config from json file
-    with open(args.model_config) as f:
-        model_config = json.load(f)
-
-
-    interface = create_ui(model_config, args.ckpt_path, args.pretransform_ckpt_path)
+    interface = create_ui(model_config_path = args.model_config, ckpt_path=args.ckpt_path, pretrained_name=args.pretrained_name, pretransform_ckpt_path=args.pretransform_ckpt_path)
     interface.queue()
     interface.launch(share=True, auth=(args.username, args.password) if args.username is not None else None)
 
 if __name__ == "__main__":
     import argparse
     parser = argparse.ArgumentParser(description='Run gradio interface')
-    parser.add_argument('--model-config', type=str, help='Path to model config', required=True)
-    parser.add_argument('--ckpt-path', type=str, help='Path to model checkpoint', required=True)
+    parser.add_argument('--pretrained-name', type=str, help='Name of pretrained model', required=False)
+    parser.add_argument('--model-config', type=str, help='Path to model config', required=False)
+    parser.add_argument('--ckpt-path', type=str, help='Path to model checkpoint', required=False)
     parser.add_argument('--pretransform-ckpt-path', type=str, help='Optional to model pretransform checkpoint', required=False)
     parser.add_argument('--username', type=str, help='Gradio username', required=False)
     parser.add_argument('--password', type=str, help='Gradio password', required=False)

setup.py

@@ -18,6 +18,7 @@
         'ema-pytorch==0.2.3',
         'encodec==0.1.1',
         'gradio==3.42.0',
+        'huggingface_hub',
         'importlib-resources==5.12.0',
         'k-diffusion==0.1.1',
         'laion-clap==1.1.4',

stable_audio_tools/__init__.py

@@ -0,0 +1,2 @@
+from .models.factory import create_model_from_config, create_model_from_config_path
+from .models.pretrained import get_pretrained_model

stable_audio_tools/models/__init__.py

@@ -1 +1 @@
-from .factory import create_model_from_config
+from .factory import create_model_from_config, create_model_from_config_path

stable_audio_tools/models/factory.py

@@ -1,3 +1,5 @@
+import json
+
 def create_model_from_config(model_config):
     model_type = model_config.get('model_type', None)
 
@@ -20,7 +22,13 @@ def create_model_from_config(model_config):
         return create_musicgen_from_config(model_config)
     else:
         raise NotImplementedError(f'Unknown model type: {model_type}')
+
+def create_model_from_config_path(model_config_path):
+    with open(model_config_path) as f:
+        model_config = json.load(f)
 
+    return create_model_from_config(model_config)
+
 def create_pretransform_from_config(pretransform_config, sample_rate):
     pretransform_type = pretransform_config.get('type', None)
 

stable_audio_tools/models/pretrained.py

@@ -0,0 +1,25 @@
+import json
+
+from .factory import create_model_from_config
+from .utils import load_ckpt_state_dict
+
+from huggingface_hub import hf_hub_download
+
+def get_pretrained_model(name: str):
+    
+    model_config_path = hf_hub_download(name, filename="model_config.json", repo_type='model')
+
+    with open(model_config_path) as f:
+        model_config = json.load(f)
+
+    model = create_model_from_config(model_config)
+
+    # Try to download the model.safetensors file first, if it doesn't exist, download the model.ckpt file
+    try:
+        model_ckpt_path = hf_hub_download(name, filename="model.safetensors", repo_type='model')
+    except Exception as e:
+        model_ckpt_path = hf_hub_download(name, filename="model.ckpt", repo_type='model')
+
+    model.load_state_dict(load_ckpt_state_dict(model_ckpt_path))
+
+    return model, model_config