sdpipeline.py

import torch
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt
import time

begin = time.time()

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",

    # use_safetensors=True,
    # local_files_only=True,
    # device_map="auto",
    # safety_checker=None
    cache_dir="models",
)

print(f"Load             : {time.time()-begin}")
pipe = pipe.to("cuda")
print(f"Load+cuda        : {time.time()-begin}")

prompt = "native american chief"
image = pipe(prompt, num_inference_steps=25).images[0]

# If the image is a torch tensor, convert it to numpy array
if torch.is_tensor(image):
    image = image.cpu().numpy()

# Display the image
plt.imshow(image)
plt.axis('off')  # To turn off axis numbers
plt.show()