-
Notifications
You must be signed in to change notification settings - Fork 8
/
t2i.py
48 lines (40 loc) · 1.39 KB
/
t2i.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
37
38
39
40
41
42
43
44
45
46
47
48
import torch
from diffusers import DiffusionPipeline
from utils import (
attn_maps,
cross_attn_init,
register_cross_attention_hook,
set_layer_with_name_and_path,
save_by_timesteps_and_path,
save_by_timesteps
)
##### 1. Init modules #####
cross_attn_init()
###########################
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda:0")
##### 2. Replace modules and Register hook #####
pipe.unet = set_layer_with_name_and_path(pipe.unet)
pipe.unet = register_cross_attention_hook(pipe.unet)
################################################
height = 512
width = 768
prompt = "A portrait photo of a kangaroo wearing an orange hoodie and blue sunglasses standing on the grass in front of the Sydney Opera House holding a sign on the chest that says 'SDXL'!."
image = pipe(
prompt,
height=height,
width=width,
num_inference_steps=15,
).images[0]
image.save('test.png')
##### 3. Process and Save attention map #####
print('resizing and saving ...')
##### 3-1. save by timesteps and path (2~3 minutes) #####
save_by_timesteps_and_path(pipe.tokenizer, prompt, height, width)
#########################################################
##### 3-2. save by timesteps (1~2 minutes) #####
# save_by_timesteps(pipe.tokenizer, prompt, height, width)
################################################