torch

>>> pipe = VersatileDiffusionTextToImagePipeline.from_pretrained(
...     "shi-labs/versatile-diffusion", torch_dtype=torch.float16
... )
>>> pipe.remove_unused_weights()
>>> pipe = pipe.to("cuda")

>>> generator = torch.Generator(device="cuda").manual_seed(0)
>>> image = pipe("an astronaut riding on a horse on mars", generator=generator).images[0]
>>> image.save("./astronaut.png")

enable_sequential_cpu_offload

<
source
>
(
gpu_id = 0
