Run Stable Diffusion 2.0 on Modal

modal_stable_diffusion_2.py

import io
import sys

import modal

stub = modal.Stub(
    image=modal.Image.debian_slim()
    .apt_install(["git"])
    .pip_install(
        [
            "git+https://github.com/huggingface/diffusers.git",
            "transformers",
            "accelerate",
            "scipy",
        ]
    )
)


@stub.function(gpu=True)
def run(prompt):
    import torch
    from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler

    model_id = "stabilityai/stable-diffusion-2"
    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16
    )
    pipe = pipe.to("cuda")

    image = pipe(prompt, height=768, width=768).images[0]
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    return buf.getvalue()


if __name__ == "__main__":
    with stub.run():
        png_data = run(sys.argv[1])
        with open("output.png", "wb") as f:
            f.write(png_data)

Note – this code is not very optimized. There are a few simple optimizations you can do though:

1. Use an A100 GPU – just do gpu=modal.gpu.A100()
2. Use a shared volume to store the weights – see this example: https://github.com/modal-labs/modal-examples/blob/main/06_gpu/stable_diffusion_slackbot.py
3. If you deploy this function, use an __enter__ method similar to this example: https://github.com/modal-labs/modal-examples/blob/main/misc/batch_inference_using_huggingface.py

There are a few other optimizations too. We will publish some more info about those hopefully soon!