forked from damian0815/compel
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathcompel-demo.py
24 lines (19 loc) · 992 Bytes
/
compel-demo.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
import torch
from compel import Compel, ReturnedEmbeddingsType
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from torch import Generator
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')