File size: 2,348 Bytes
b63b631 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
# Run Stable Diffusion to generate images from text prompts.
# This script uses the Modal framework to run Stable Diffusion in a cloud environment.
# It requires the `modal` package and a Hugging Face token stored in a Modal secret.
# Make sure to set up the Modal environment and install the necessary dependencies.
# Usage: modal run stablediffusion.py
from icecream import ic
import io
import os
import modal
app = modal.App()
@app.function(
image=modal.Image.debian_slim().pip_install(
"icecream",
"torch",
"diffusers[torch]",
"transformers",
"ftfy"
),
secrets=[modal.Secret.from_name("huggingface-secret")],
gpu="any",
)
def run_stable_diffusion(prompt: str):
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
use_auth_token=os.environ["HF_TOKEN"],
).to("cuda")
image = pipe(prompt, num_inference_steps=10).images[0] # type: ignore
buf = io.BytesIO()
image.save(buf, format="PNG")
img_bytes = buf.getvalue()
return img_bytes
@app.local_entrypoint()
def main():
prompt = "Wu-Tang Clan climbing Mount Everest"
# prompt = "A robot dog walking down a vineyard" # Example prompt
# out_path = "/tmp/output.png"
out_path = "stablediffusion_output.png"
# ic(os.getcwd())
# img_bytes = b"<image_bytes>" # Placeholder for the actual image bytes
print("DEBUG: Starting Stable Diffusion with prompt:", prompt)
img_bytes = run_stable_diffusion.remote(prompt=prompt)
print("DEBUG: Writing img_bytes length:", len(img_bytes))
with open(out_path, "wb") as f:
f.write(img_bytes)
print("DEBUG: Image saved to ", out_path)
if __name__ == "__main__":
main()
print("Image saved to /tmp/output.png")
print("Run `modal deploy` to deploy this app.")
print("Run `modal serve` to serve this app locally.")
print("Run `modal run` to run this app in the cloud.")
print("Run `modal logs` to view the logs of this app.")
print("Run `modal shell` to open a shell in the cloud environment.")
print("Run `modal run --help` to see all available options.")
print("Run `modal deploy --help` to see all available options for deployment.") |