Download weights first on Spaces

app.py

@@ -41,7 +41,7 @@ from app_scribble import create_demo as create_demo_scribble
 from app_scribble_interactive import \
     create_demo as create_demo_scribble_interactive
 from app_seg import create_demo as create_demo_seg
-from model import Model
+from model import Model, download_all_controlnet_weights
 
 DESCRIPTION = '# [ControlNet](https://github.com/lllyasviel/ControlNet)'
 
@@ -55,6 +55,9 @@ if SPACE_ID is not None:
 MAX_IMAGES = int(os.getenv('MAX_IMAGES', '3'))
 DEFAULT_NUM_IMAGES = min(MAX_IMAGES, int(os.getenv('DEFAULT_NUM_IMAGES', '1')))
 
+if os.getenv('SYSTEM') == 'spaces':
+    download_all_controlnet_weights()
+
 DEFAULT_MODEL_ID = os.getenv('DEFAULT_MODEL_ID',
                              'runwayml/stable-diffusion-v1-5')
 model = Model(base_model_id=DEFAULT_MODEL_ID, task_name='canny')

model.py

@@ -38,6 +38,11 @@ CONTROLNET_MODEL_IDS = {
 }
 
 
+def download_all_controlnet_weights():
+    for model_id in CONTROLNET_MODEL_IDS.values():
+        ControlNetModel.from_pretrained(model_id)
+
+
 class Model:
     def __init__(self,
                  base_model_id: str = 'runwayml/stable-diffusion-v1-5',