tomesd_token_merging_stable_diffusion_inference_speedup.py

1import torch, tomesd
2from diffusers import StableDiffusionPipeline
3
4# Load the model
5pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
6pipe.to("cuda")
7
8# Apply ToMe to the pipeline
9tomesd.apply_patch(pipe, ratio=0.5)
10
11# Generate an image as usual
12prompt = "a photo of an astronaut riding a horse on mars"
13image = pipe(prompt).images[0]
14
15image.save("astronaut.png")