How to run Stable Diffusion 3.5 Medium on Modal

All posts

Back

Model Library

December 16, 2024•5 minute read

Yiren Lu@YirenLu

Solutions Engineer

Introduction to Stable Diffusion 3.5 Medium

Stable Diffusion 3.5 Medium is a powerful image generation model with 2.5 billion parameters that is designed to run “out of the box” on consumer hardware, striking a balance between quality and ease of customization. It is trained on 1024x1024 images, and it most suitable for generating images with those resolutions.

To run the following code, you will need to:

Create an account at modal.com
Run pip install modal to install the modal Python package
Run modal setup to authenticate (if this doesn’t work, try python -m modal setup)
Copy the code below into a file called app.py
Run modal run app.py

Please note that this code does not come with a UI. For a more detailed example of how to run Stable Diffusion 3.5 Medium as a CLI, API, and UI, refer here. (You can modify the code in that example to run the medium version instead of the large version.)

import io
import random
from pathlib import Path

import modal

app = modal.App("stable-diffusion-medium-model-library")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "accelerate==0.33.0",
        "diffusers==0.31.0",
        "fastapi[standard]==0.115.4",
        "huggingface-hub[hf_transfer]==0.25.2",
        "sentencepiece==0.2.0",
        "torch==2.5.1",
        "torchvision==0.20.1",
        "transformers~=4.44.0",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

with image.imports():
    import diffusers
    import torch


@app.cls(
    image=image,
    gpu="H100",
    timeout=600,
)
class Inference:
    @modal.build()
    @modal.enter()
    def initialize(self):
        self.pipe = diffusers.StableDiffusion3Pipeline.from_pretrained(
            "adamo1139/stable-diffusion-3.5-medium-ungated",
            revision="1cb9becf522803da9d25c2713d1a82b41bc6198d",
            torch_dtype=torch.bfloat16,
        )

    @modal.enter()
    def move_to_gpu(self):
        self.pipe.to("cuda")

    @modal.method()
    def run(self, prompt: str, batch_size: int = 4, seed: int = None) -> list[bytes]:
        seed = seed if seed is not None else random.randint(0, 2**32 - 1)
        print("seeding RNG with", seed)
        torch.manual_seed(seed)

        images = self.pipe(
            prompt,
            num_images_per_prompt=batch_size,
            num_inference_steps=4,
            guidance_scale=0.0,
            max_sequence_length=512,
        ).images

        image_output = []
        for image in images:
            with io.BytesIO() as buf:
                image.save(buf, format="PNG")
                image_output.append(buf.getvalue())
        torch.cuda.empty_cache()
        return image_output


@app.local_entrypoint()
def main(prompt: str = "A princess riding on a pony"):
    output_dir = Path("/tmp/stable-diffusion")
    output_dir.mkdir(exist_ok=True)

    images = Inference().run.remote(prompt, batch_size=1)

    for i, image_bytes in enumerate(images):
        output_path = output_dir / f"output_{i:02d}.png"
        output_path.write_bytes(image_bytes)
        print(f"Saved {output_path}")

Introduction to Stable Diffusion 3.5 Medium

Example code for running the Stable Diffusion 3.5 Medium image generation model on Modal

Ship your first app in minutes.