Add trusses for Stable Diffusion models (#233)

* add trusses for SD text2image, image2image, and inpainting

* add shared external package for b64 utils + readme

Author: spal1

examples/shared/base64_utils.py

@@ -0,0 +1,17 @@
+import base64
+from io import BytesIO
+
+from PIL import Image
+
+BASE64_PREAMBLE = "data:image/png;base64,"
+
+
+def pil_to_b64(pil_img):
+    buffered = BytesIO()
+    pil_img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue())
+    return BASE64_PREAMBLE + str(img_str)[2:-1]
+
+
+def b64_to_pil(b64_str):
+    return Image.open(BytesIO(base64.b64decode(b64_str.replace(BASE64_PREAMBLE, ""))))

examples/stable-diffusion-1-5-img2img/config.yaml

@@ -0,0 +1,33 @@
+bundled_packages_dir: packages
+data_dir: data
+description: null
+environment_variables: {}
+examples_filename: examples.yaml
+external_package_dirs:
+- ../shared/
+input_type: Any
+live_reload: false
+model_class_filename: model.py
+model_class_name: Model
+model_framework: custom
+model_metadata: {}
+model_module_dir: model
+model_name: Stable Diffusion 1.5 Img2Img
+model_type: custom
+python_version: py39
+requirements:
+- diffusers
+- transformers
+- accelerate
+- scipy
+- safetensors
+- xformers
+- triton
+resources:
+  accelerator: A10G
+  cpu: "8"
+  memory: 30Gi
+  use_gpu: true
+secrets: {}
+spec_version: '2.0'
+system_packages: []

examples/stable-diffusion-1-5-img2img/examples.yaml

@@ -0,0 +1,35 @@
+from dataclasses import asdict
+from typing import Dict
+
+import torch
+from base64_utils import b64_to_pil, pil_to_b64
+from diffusers import StableDiffusionImg2ImgPipeline
+
+STABLE_DIFFUSION_MODEL_ID = "runwayml/stable-diffusion-v1-5"
+
+
+class Model:
+    def __init__(self, **kwargs) -> None:
+        self._model = None
+
+    def load(self):
+        self._model = StableDiffusionImg2ImgPipeline.from_pretrained(
+            STABLE_DIFFUSION_MODEL_ID,
+            torch_dtype=torch.float16,
+        )
+        self._model = self._model.to("cuda")
+
+    def preprocess(self, request: Dict) -> Dict:
+        # Convert from base64
+        if "image" in request:
+            request["image"] = b64_to_pil(request["image"]).convert("RGB")
+        return request
+
+    def postprocess(self, request: Dict) -> Dict:
+        # Convert to base64
+        request.images = [pil_to_b64(img) for img in request.images]
+        return asdict(request)
+
+    def predict(self, request: Dict):
+        response = self._model(**request)
+        return response

examples/stable-diffusion-1-5-img2img/model/__init__.py

@@ -0,0 +1,33 @@
+bundled_packages_dir: packages
+data_dir: data
+description: null
+environment_variables: {}
+examples_filename: examples.yaml
+external_package_dirs:
+- ../shared/
+input_type: Any
+live_reload: false
+model_class_filename: model.py
+model_class_name: Model
+model_framework: custom
+model_metadata: {}
+model_module_dir: model
+model_name: Stable Diffusion 1.5
+model_type: custom
+python_version: py39
+requirements:
+- diffusers
+- transformers
+- accelerate
+- scipy
+- safetensors
+- xformers
+- triton
+resources:
+  accelerator: A10G
+  cpu: "8"
+  memory: 30Gi
+  use_gpu: true
+secrets: {}
+spec_version: '2.0'
+system_packages: []

examples/stable-diffusion-1-5-img2img/model/model.py

@@ -0,0 +1,51 @@
+# Adding *custom weights* to Stable Diffusion
+If you have custom weights for Stable Diffusion and want to build a truss for it, its as simple as doing the following!
+
+## Add weights to the data directory
+In your truss, add your weights to the `data` directory, and make sure its specified in your `config.yaml` file here:
+```
+...
+data_dir: data
+...
+```
+Make sure you have the appropriate subfolders, your `data` directory structure should look something like this:
+```
+your_truss/
+    ...
+    data/
+        feature_extractor/
+        safety_checker/
+        scheduler/
+        text_encoder/
+        tokenizer/
+        unet/
+        vae/
+        model_index.json
+    ...
+```
+
+## Reference the data directory in the `models.py` file
+We can [utilize the functionality](https://huggingface.co/docs/diffusers/main/en/api/diffusion_pipeline#diffusers.DiffusionPipeline.from_pretrained) to pass in a path to the `from_pretrained` method off of the `StableDiffusionPipeline`.
+
+You'll need to reference this data directory in your `model.py` file, in this case `_data_dir` is the name of your `data` directory:
+
+```
+class Model:
+    def __init__(self, **kwargs) -> None:
+        self._data_dir = kwargs["data_dir"]
+        self._model = None
+
+    def load(self):
+        self._model = StableDiffusionPipeline.from_pretrained(
+            self._data_dir,
+            torch_dtype=torch.float16,
+        )
+        self._model = self._model.to("cuda")
+
+    ...
+```
+
+And this is all you need to load your custom weights and create a truss for your custom Stable Diffusion model!
+
+## Truss in the process...
+[Follow the documentation here](https://truss.baseten.co/develop/localhost) to build and run your truss!

examples/stable-diffusion-1-5/config.yaml

@@ -0,0 +1,3 @@
+- input:
+    prompt: a photo of an astronaut riding a horse on mars
+  name: astronaut example

examples/stable-diffusion-1-5/data/README.md

@@ -0,0 +1,35 @@
+from dataclasses import asdict
+from typing import Dict
+
+import torch
+from base64_utils import pil_to_b64
+from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline
+
+STABLE_DIFFUSION_MODEL_ID = "runwayml/stable-diffusion-v1-5"
+
+
+class Model:
+    def __init__(self, **kwargs) -> None:
+        self._model = None
+
+    def load(self):
+        scheduler = EulerDiscreteScheduler.from_pretrained(
+            STABLE_DIFFUSION_MODEL_ID,
+            subfolder="scheduler",
+        )
+        self._model = StableDiffusionPipeline.from_pretrained(
+            STABLE_DIFFUSION_MODEL_ID,
+            scheduler=scheduler,
+            torch_dtype=torch.float16,
+        )
+        self._model.unet.set_use_memory_efficient_attention_xformers(True)
+        self._model = self._model.to("cuda")
+
+    def postprocess(self, request: Dict) -> Dict:
+        # Convert to base64
+        request.images = [pil_to_b64(img) for img in request.images]
+        return asdict(request)
+
+    def predict(self, request: Dict):
+        response = self._model(**request)
+        return response

examples/stable-diffusion-1-5/examples.yaml

@@ -0,0 +1,33 @@
+bundled_packages_dir: packages
+data_dir: data
+description: null
+environment_variables: {}
+examples_filename: examples.yaml
+external_package_dirs:
+- ../shared/
+input_type: Any
+live_reload: false
+model_class_filename: model.py
+model_class_name: Model
+model_framework: custom
+model_metadata: {}
+model_module_dir: model
+model_name: Stable Diffusion Inpainting
+model_type: custom
+python_version: py39
+requirements:
+- diffusers
+- transformers
+- accelerate
+- scipy
+- safetensors
+- xformers
+- triton
+resources:
+  accelerator: A10G
+  cpu: "8"
+  memory: 30Gi
+  use_gpu: true
+secrets: {}
+spec_version: '2.0'
+system_packages: []

examples/stable-diffusion-1-5/model/__init__.py

@@ -0,0 +1,36 @@
+from dataclasses import asdict
+from typing import Dict
+
+import torch
+from base64_utils import b64_to_pil, pil_to_b64
+from diffusers import StableDiffusionInpaintPipeline
+
+STABLE_DIFFUSION_MODEL_ID = "runwayml/stable-diffusion-inpainting"
+
+
+class Model:
+    def __init__(self, **kwargs) -> None:
+        self._model = None
+
+    def load(self):
+        self._model = StableDiffusionInpaintPipeline.from_pretrained(
+            STABLE_DIFFUSION_MODEL_ID, revision="fp16", torch_dtype=torch.float16
+        )
+        self._model = self._model.to("cuda")
+
+    def preprocess(self, request: Dict) -> Dict:
+        # Convert from base64
+        if "image" in request:
+            request["image"] = b64_to_pil(request["image"])
+        if "mask_image" in request:
+            request["mask_image"] = b64_to_pil(request["mask_image"])
+        return request
+
+    def postprocess(self, request: Dict) -> Dict:
+        # Convert to base64
+        request.images = [pil_to_b64(img) for img in request.images]
+        return asdict(request)
+
+    def predict(self, request: Dict):
+        response = self._model(**request)
+        return response