Update deployment code to use explicit downloader (#54)

diffusion/inference/inference_model.py

@@ -5,7 +5,7 @@
 
 import base64
 import io
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 import torch
 from composer.utils.file_helpers import get_file
@@ -17,6 +17,15 @@
 LOCAL_CHECKPOINT_PATH = '/tmp/model.pt'
 
 
+def download_checkpoint(chkpt_path: str):
+    """Downloads the Stable Diffusion checkpoint to the local filesystem.
+
+    Args:
+        chkpt_path (str): The path to the local folder, URL or object score that contains the checkpoint.
+    """
+    get_file(path=chkpt_path, destination=LOCAL_CHECKPOINT_PATH)
+
+
 class StableDiffusionInference():
     """Inference endpoint class for Stable Diffusion.
 
@@ -26,13 +35,11 @@ class StableDiffusionInference():
             Default: ``None``.
     """
 
-    def __init__(self, chkpt_path: Optional[str] = None):
-        pretrained_flag = chkpt_path is None
+    def __init__(self, pretrained: bool = False):
         self.device = torch.cuda.current_device()
 
-        model = stable_diffusion_2(pretrained=pretrained_flag, encode_latents_in_fp16=True, fsdp=False)
-        if not pretrained_flag:
-            get_file(path=chkpt_path, destination=LOCAL_CHECKPOINT_PATH)
+        model = stable_diffusion_2(pretrained=pretrained, encode_latents_in_fp16=True, fsdp=False)
+        if not pretrained:
             state_dict = torch.load(LOCAL_CHECKPOINT_PATH)
             for key in list(state_dict['state']['model'].keys()):
                 if 'val_metrics.' in key:

diffusion/inference/mosaic_inference.yaml

@@ -10,7 +10,12 @@ integrations:
     git_branch: main
     pip_install: .[all]
 model:
+  downloader: diffusion.inference.inference_model.download_checkpoint
+  download_parameters:
+    chkpt_path:  # Path to download the checkpoint to evaluate
   model_handler: diffusion.inference.inference_model.StableDiffusionInference
+  model_parameters:
+    pretrained: false
 command: |
   export PYTHONPATH=$PYTHONPATH:/code/diffusion
   rm /usr/lib/python3/dist-packages/packaging-23.1.dist-info/REQUESTED