Copied pipe.unet.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last) Now compile and perform inference:   Copied # Compile the UNet and VAE.
pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

# First call to `pipe` is slow,