Skip to content

Instantly share code, notes, and snippets.

@victormurcia
Created January 12, 2023 01:22
Show Gist options
  • Save victormurcia/d594b64f987153ebaf305368d19f2652 to your computer and use it in GitHub Desktop.
Save victormurcia/d594b64f987153ebaf305368d19f2652 to your computer and use it in GitHub Desktop.
code to setup Stable Diffusion pipeline
import mediapy as media
import torch
from diffusers import StableDiffusionPipeline
device = "cuda"
if model_id.startswith("stabilityai/"):
model_revision = "fp16"
else:
model_revision = None
if scheduler is None:
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16,
revision=model_revision,
)
else:
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
scheduler=scheduler,
torch_dtype=torch.float16,
revision=model_revision,
)
pipe = pipe.to(device)
pipe.enable_xformers_memory_efficient_attention()
if model_id.endswith('-base'):
image_length = 512
else:
image_length = 768
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment