Skip to content

Instantly share code, notes, and snippets.

@erikbern

erikbern/modal_stable_diffusion_2.py

Last active July 29, 2023 20:10
Show Gist options
  • Save erikbern/03a056b79e18692eb9b848f84eefc57b to your computer and use it in GitHub Desktop.
Save erikbern/03a056b79e18692eb9b848f84eefc57b to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. erikbern revised this gist Nov 25, 2022. 1 changed file with 0 additions and 2 deletions.
    2 changes: 0 additions & 2 deletions modal_stable_diffusion_2.py
    View file
    Original file line number Diff line number Diff line change
    @@ -23,8 +23,6 @@ def run(prompt):
    from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler

    model_id = "stabilityai/stable-diffusion-2"

    # Use the Euler scheduler here instead
    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
    pipe = StableDiffusionPipeline.from_pretrained(
    model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16
  2. erikbern created this gist Nov 25, 2022.
    44 changes: 44 additions & 0 deletions modal_stable_diffusion_2.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,44 @@
    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"

    # Use the Euler scheduler here instead
    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)