-
Notifications
You must be signed in to change notification settings - Fork 47
/
compel-demo-sdxl.py
36 lines (28 loc) · 1.57 KB
/
compel-demo-sdxl.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import torch
from compel import Compel, ReturnedEmbeddingsType
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from torch import Generator
device='cuda'
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
variant="fp16",
use_safetensors=True,
torch_dtype=torch.float16).to(device)
compel = Compel(tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2] ,
text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
requires_pooled=[False, True])
device = "mps"
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
# dpm++
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config,
algorithm_type="dpmsolver++")
prompts = ["a cat playing with a ball++ in the forest", "a cat playing with a ball in the forest"]
compel = Compel(tokenizer=pipeline.tokenizer, text_encoder=pipeline.text_encoder)
prompt = "a cat playing with a ball++ in the forest"
conditioning, pooled = compel(prompt)
print(conditioning.shape, pooled.shape)
prompt_embeds = compel(prompts)
images = pipeline(prompt_embeds=prompt_embeds, num_inference_steps=10, width=256, height=256).images
print(images)
images[0].save('/tmp/img0.jpg')
images[1].save('/tmp/img1.jpg')