ssd_1b.py

# -*- coding: utf-8 -*-
"""SSD 1B.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1U2ImlQGa0gnZXhVXm_3JH9v9IclYvkuB
"""

!pip install git+https://github.com/huggingface/diffusers
!pip install transformers accelerate safetensors

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained(
    "segmind/SSD-1B",
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16"
)
pipe.to("cuda")

# Commented out IPython magic to ensure Python compatibility.
# %time
from transformers.pipelines.image_to_text import Image
from IPython.display import display

prompt = "cute cat sitting on the chair"
neg_prompt = "ugly, blurry, poor quality"

image = pipe(prompt=prompt, negative_prompt=neg_prompt).images[0]

display(image)