Views
No views yet

pip install -U diffusersDiffusionPipeline to run the model1import torch
2from diffusers import DiffusionPipeline
3
4# Load the diffusion pipeline from a pretrained model, using bfloat16 for tensor types.
5pipe = DiffusionPipeline.from_pretrained(
6 "shuttleai/shuttle-3-diffusion", torch_dtype=torch.bfloat16
7).to("cuda")
8
9# Uncomment the following line to save VRAM by offloading the model to CPU if needed.
10# pipe.enable_model_cpu_offload()
11
12# Uncomment the lines below to enable torch.compile for potential performance boosts on compatible GPUs.
13# Note that this can increase loading times considerably.
14# pipe.transformer.to(memory_format=torch.channels_last)
15# pipe.transformer = torch.compile(
16# pipe.transformer, mode="max-autotune", fullgraph=True
17# )
18
19# Set your prompt for image generation.
20prompt = "A cat holding a sign that says hello world"
21
22# Generate the image using the diffusion pipeline.
23image = pipe(
24 prompt,
25 height=1024,
26 width=1024,
27 guidance_scale=3.5,
28 num_inference_steps=4,
29 max_sequence_length=256,
30 # Uncomment the line below to use a manual seed for reproducible results.
31 # generator=torch.Generator("cpu").manual_seed(0)
32).images[0]
33
34# Save the generated image.
35image.save("shuttle.png")