refactor loadModel. use model_inputs to check {input,mask}_image

gadicc · gadicc · commit d2d5ad1b60a6 · 2022-09-21T09:15:10.000+03:00
diff --git a/Dockerfile b/Dockerfile
@@ -25,6 +25,7 @@ ADD DOWNLOAD_VARS.py .
 
 # Add your model weight files 
 # (in this case we have a python script)
+ADD loadModel.py .
 ADD download.py .
 RUN python3 download.py
 
diff --git a/app.py b/app.py
@@ -9,9 +9,9 @@
 )
 import base64
 from io import BytesIO
-import os
 import PIL
 import json
+from loadModel import loadModel
 
 from APP_VARS import MODEL_ID
 
@@ -54,8 +54,6 @@ def init():
     global schedulers
     global dummy_safety_checker
 
-    HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
-
     schedulers = {
         "LMS": LMSDiscreteScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
@@ -77,15 +75,7 @@ def init():
         last_model_id = None
         return
 
-    print("Loading model " + MODEL_ID)
-
-    model = _pipelines.StableDiffusionPipeline.from_pretrained(
-        MODEL_ID,
-        revision="fp16",
-        torch_dtype=torch.float16,
-        use_auth_token=HF_AUTH_TOKEN,
-    ).to("cuda")
-
+    model = loadModel(MODEL_ID)
     pipelines = createPipelinesFromModel(MODEL_ID)
 
 
@@ -120,17 +110,9 @@ def inference(all_inputs: dict) -> dict:
         return {"$error": "UPGRADE CLIENT - no model_inputs specified"}
 
     if MODEL_ID == "ALL":
-        HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
         model_id = call_inputs.get("MODEL_ID")
         if last_model_id != model_id:
-            print("Loading model " + model_id)
-            model = _pipelines.StableDiffusionPipeline.from_pretrained(
-                model_id,
-                revision="fp16",
-                torch_dtype=torch.float16,
-                use_auth_token=HF_AUTH_TOKEN,
-            ).to("cuda")
-
+            model = loadModel(model_id)
             pipelines = createPipelinesFromModel(model_id)
             last_model_id = model_id
 
@@ -155,18 +137,11 @@ def inference(all_inputs: dict) -> dict:
     # seed = model_inputs.get("seed", None)
     #   strength = model_inputs.get("strength", 0.75)
 
-    if call_inputs.get("PIPELINE") in [
-        "StableDiffusionImg2ImgPipeline",
-        "StableDiffusionInpaintPipeline",
-    ]:
-        model_inputs.update(
-            {"init_image": decodeBase64Image(model_inputs.get("init_image"))}
-        )
+    if "init_image" in model_inputs:
+        model_inputs["init_image"] = decodeBase64Image(model_inputs.get("init_image"))
 
-    if all_inputs.get("PIPELINE") == "StableDiffusionInpaintPipeline":
-        model_inputs.update(
-            {"mask_image": decodeBase64Image(model_inputs.get("mask_image"))}
-        )
+    if "mask_image" in model_inputs:
+        model_inputs["mask_image"] = decodeBase64Image(model_inputs.get("mask_image"))
 
     seed = model_inputs.get("seed", None)
     if seed == None:
diff --git a/download.py b/download.py
@@ -1,35 +1,18 @@
 # In this file, we define download_model
 # It runs during container build time to get model weights built into the container
 
-from diffusers import StableDiffusionPipeline
-import torch
-import os
 from DOWNLOAD_VARS import MODEL_ID
-
-MODEL_IDS = ["CompVis/stable-diffusion-v1-4", "hakurei/waifu-diffusion"]
+from loadModel import loadModel, MODEL_IDS
 
 
 def download_model():
     # do a dry run of loading the huggingface model, which will download weights at build time
-    # Set auth token which is required to download stable diffusion model weights
-    HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
-
-    # Bad for production serverless, great for local dev & preview deploys
+    # For local dev & preview deploys, download all the models (terrible for serverless deploys)
     if MODEL_ID == "ALL":
         for MODEL_I in MODEL_IDS:
-            StableDiffusionPipeline.from_pretrained(
-                MODEL_I,
-                revision="fp16",
-                torch_dtype=torch.float16,
-                use_auth_token=HF_AUTH_TOKEN,
-            )
+            loadModel(MODEL_I, False)
     else:
-        model = StableDiffusionPipeline.from_pretrained(
-            MODEL_ID,
-            revision="fp16",
-            torch_dtype=torch.float16,
-            use_auth_token=HF_AUTH_TOKEN,
-        )
+        loadModel(MODEL_ID, False)
 
 
 if __name__ == "__main__":
diff --git a/loadModel.py b/loadModel.py
@@ -0,0 +1,19 @@
+import torch
+import os
+from diffusers import StableDiffusionPipeline
+
+HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
+MODEL_IDS = ["CompVis/stable-diffusion-v1-4", "hakurei/waifu-diffusion"]
+
+
+def loadModel(model_id: str, load=True):
+    print(("Loading" if load else "Downloading") + " model: " + model_id)
+
+    model = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        revision="fp16",
+        torch_dtype=torch.float16,
+        use_auth_token=HF_AUTH_TOKEN,
+    )
+
+    return model.to("cuda") if load else None
