schedulers: local_files_only speedsup from_config from 800ms to 1ms

gadicc · gadicc · commit a1614df1b442 · 2022-11-08T13:01:24.000Z
diff --git a/Dockerfile b/Dockerfile
@@ -94,11 +94,6 @@ ADD convert-to-diffusers.py .
 RUN python3 convert-to-diffusers.py
 # RUN rm -rf checkpoints
 
-# Loading a new scheduler for the first time takes an extra ~800ms,
-# so set this to your most common one.
-ARG DEFAULT_SCHEDULER="LMSDiscreteScheduler"
-ENV DEFAULT_SCHEDULER=${DEFAULT_SCHEDULER}
-
 # Add your model weight files 
 # (in this case we have a python script)
 ADD getScheduler.py .
diff --git a/download.py b/download.py
@@ -3,7 +3,6 @@
 
 import os
 from loadModel import loadModel, MODEL_IDS
-from getScheduler import getScheduler, SCHEDULERS, DEFAULT_SCHEDULER
 
 MODEL_ID = os.environ.get("MODEL_ID")
 
@@ -17,8 +16,6 @@ def download_model():
     else:
         loadModel(MODEL_ID, False)
 
-    getScheduler(MODEL_ID, DEFAULT_SCHEDULER)
-
 
 if __name__ == "__main__":
     download_model()
diff --git a/getScheduler.py b/getScheduler.py
@@ -4,7 +4,6 @@
 from diffusers import schedulers as _schedulers
 
 HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
-DEFAULT_SCHEDULER = os.getenv("DEFAULT_SCHEDULER")
 
 SCHEDULERS = [
     "LMSDiscreteScheduler",
@@ -14,6 +13,9 @@
     "EulerDiscreteScheduler",
 ]
 
+DEFAULT_SCHEDULER = os.getenv("DEFAULT_SCHEDULER", SCHEDULERS[0])
+
+
 """
 # This was a nice idea but until we have default init vars for all schedulers
 # via from_config(), it's a no go.  In any case, loading a scheduler takes time
@@ -31,15 +33,18 @@
 """
 
 
-def initScheduler(MODEL_ID: str, scheduler_id: str):
+def initScheduler(MODEL_ID: str, scheduler_id: str, download=False):
     print(f"Initializing {scheduler_id} for {MODEL_ID}...")
     start = time.time()
     scheduler = getattr(_schedulers, scheduler_id)
     if scheduler == None:
         return None
 
     inittedScheduler = scheduler.from_config(
-        MODEL_ID, subfolder="scheduler", use_auth_token=HF_AUTH_TOKEN
+        MODEL_ID,
+        subfolder="scheduler",
+        use_auth_token=HF_AUTH_TOKEN,
+        local_files_only=not download,
     )
     diff = round((time.time() - start) * 1000)
     print(f"Initialized {scheduler_id} for {MODEL_ID} in {diff}ms")
@@ -50,7 +55,7 @@ def initScheduler(MODEL_ID: str, scheduler_id: str):
 schedulers = {}
 
 
-def getScheduler(MODEL_ID: str, scheduler_id: str):
+def getScheduler(MODEL_ID: str, scheduler_id: str, download=False):
     schedulersByModel = schedulers.get(MODEL_ID, None)
     if schedulersByModel == None:
         schedulersByModel = {}
@@ -73,7 +78,7 @@ def getScheduler(MODEL_ID: str, scheduler_id: str):
 
     scheduler = schedulersByModel.get(scheduler_id, None)
     if scheduler == None:
-        scheduler = initScheduler(MODEL_ID, scheduler_id)
+        scheduler = initScheduler(MODEL_ID, scheduler_id, download)
         schedulersByModel.update({scheduler_id: scheduler})
 
     return scheduler
diff --git a/loadModel.py b/loadModel.py
@@ -1,7 +1,7 @@
 import torch
 import os
 from diffusers import pipelines as _pipelines, StableDiffusionPipeline
-from getScheduler import getScheduler, SCHEDULERS, DEFAULT_SCHEDULER
+from getScheduler import getScheduler, DEFAULT_SCHEDULER
 
 HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
 PIPELINE = os.getenv("PIPELINE")
@@ -22,8 +22,7 @@ def loadModel(model_id: str, load=True):
         StableDiffusionPipeline if PIPELINE == "ALL" else getattr(_pipelines, PIPELINE)
     )
 
-    print("DEFAULT SCHEDULER=" + DEFAULT_SCHEDULER)
-    scheduler = getScheduler(model_id, DEFAULT_SCHEDULER)
+    scheduler = getScheduler(model_id, DEFAULT_SCHEDULER, not load)
 
     model = pipeline.from_pretrained(
         model_id,
