-
Notifications
You must be signed in to change notification settings - Fork 9
/
pnp.py
151 lines (120 loc) · 6.17 KB
/
pnp.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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import glob
import os
from pathlib import Path
import torch
import torch.nn as nn
import torchvision.transforms as T
import argparse
from PIL import Image
import yaml
from tqdm import tqdm
from transformers import logging
from diffusers import DDIMScheduler, StableDiffusionPipeline
from pnp_utils import *
# suppress partial model loading warning
logging.set_verbosity_error()
class PNP(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.device = config["device"]
sd_version = config["sd_version"]
if sd_version == '2.1':
model_key = "stabilityai/stable-diffusion-2-1-base"
elif sd_version == '2.0':
model_key = "stabilityai/stable-diffusion-2-base"
elif sd_version == '1.5':
model_key = "runwayml/stable-diffusion-v1-5"
else:
raise ValueError(f'Stable-diffusion version {sd_version} not supported.')
# Create SD models
print('Loading SD model')
pipe = StableDiffusionPipeline.from_pretrained(model_key, torch_dtype=torch.float16).to("cuda")
pipe.enable_xformers_memory_efficient_attention()
self.vae = pipe.vae
self.tokenizer = pipe.tokenizer
self.text_encoder = pipe.text_encoder
self.unet = pipe.unet
self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
self.scheduler.set_timesteps(config["n_timesteps"], device=self.device)
print('SD model loaded')
# load image
self.image, self.eps = self.get_data()
self.text_embeds = self.get_text_embeds(config["prompt"], config["negative_prompt"])
self.pnp_guidance_embeds = self.get_text_embeds("", "").chunk(2)[0]
@torch.no_grad()
def get_text_embeds(self, prompt, negative_prompt, batch_size=1):
# Tokenize text and get embeddings
text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
truncation=True, return_tensors='pt')
text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
# Do the same for unconditional embeddings
uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
return_tensors='pt')
uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
# Cat for final embeddings
text_embeddings = torch.cat([uncond_embeddings] * batch_size + [text_embeddings] * batch_size)
return text_embeddings
@torch.no_grad()
def decode_latent(self, latent):
with torch.autocast(device_type='cuda', dtype=torch.float32):
latent = 1 / 0.18215 * latent
img = self.vae.decode(latent).sample
img = (img / 2 + 0.5).clamp(0, 1)
return img
@torch.autocast(device_type='cuda', dtype=torch.float32)
def get_data(self):
# load image
image = Image.open(self.config["image_path"]).convert('RGB')
image = image.resize((512, 512), resample=Image.Resampling.LANCZOS)
image = T.ToTensor()(image).to(self.device)
# get noise
latents_path = os.path.join(self.config["latents_path"], os.path.splitext(os.path.basename(self.config["image_path"]))[0], f'noisy_latents_{self.scheduler.timesteps[0]}.pt')
noisy_latent = torch.load(latents_path).to(self.device)
return image, noisy_latent
@torch.no_grad()
def denoise_step(self, x, t):
# register the time step and features in pnp injection modules
source_latents = load_source_latents_t(t, os.path.join(self.config["latents_path"], os.path.splitext(os.path.basename(self.config["image_path"]))[0]))
latent_model_input = torch.cat([source_latents] + ([x] * 2))
register_time(self, t.item())
# compute text embeddings
text_embed_input = torch.cat([self.pnp_guidance_embeds, self.text_embeds], dim=0)
# apply the denoising network
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input)['sample']
# perform guidance
_, noise_pred_uncond, noise_pred_cond = noise_pred.chunk(3)
noise_pred = noise_pred_uncond + self.config["guidance_scale"] * (noise_pred_cond - noise_pred_uncond)
# compute the denoising step with the reference model
denoised_latent = self.scheduler.step(noise_pred, t, x)['prev_sample']
return denoised_latent
def init_pnp(self, conv_injection_t, qk_injection_t):
self.qk_injection_timesteps = self.scheduler.timesteps[:qk_injection_t] if qk_injection_t >= 0 else []
self.conv_injection_timesteps = self.scheduler.timesteps[:conv_injection_t] if conv_injection_t >= 0 else []
register_attention_control_efficient(self, self.qk_injection_timesteps)
register_conv_control_efficient(self, self.conv_injection_timesteps)
def run_pnp(self):
pnp_f_t = int(self.config["n_timesteps"] * self.config["pnp_f_t"])
pnp_attn_t = int(self.config["n_timesteps"] * self.config["pnp_attn_t"])
self.init_pnp(conv_injection_t=pnp_f_t, qk_injection_t=pnp_attn_t)
edited_img = self.sample_loop(self.eps)
def sample_loop(self, x):
with torch.autocast(device_type='cuda', dtype=torch.float32):
for i, t in enumerate(tqdm(self.scheduler.timesteps, desc="Sampling")):
x = self.denoise_step(x, t)
decoded_latent = self.decode_latent(x)
T.ToPILImage()(decoded_latent[0]).save(f'{self.config["output_path"]}/output-{self.config["prompt"]}.png')
return decoded_latent
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='pnp-configs/config-horse.yaml')
opt = parser.parse_args()
with open(opt.config_path, "r") as f:
config = yaml.safe_load(f)
os.makedirs(config["output_path"], exist_ok=True)
with open(os.path.join(config["output_path"], "config.yaml"), "w") as f:
yaml.dump(config, f)
seed_everything(config["seed"])
print(config)
pnp = PNP(config)
pnp.run_pnp()