Spaces:
vilarin
/
Running on Zero

vilarin commited on
Commit
0cffd40
1 Parent(s): a68c5d3

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +229 -0
app.py ADDED
@@ -0,0 +1,229 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import UNet2DConditionModel, AutoencoderKL, DDIMScheduler, AutoencoderTiny
2
+ from transformers import AutoTokenizer, CLIPTextModel, CLIPTextModelWithProjection
3
+ from accelerate import Accelerator
4
+ from huggingface_hub import hf_hub_download
5
+ import spaces
6
+ import gradio as gr
7
+ import numpy as np
8
+ import torch
9
+ import time
10
+ import PIL
11
+
12
+ base = "stabilityai/stable-diffusion-xl-base-1.0"
13
+ repo_id = "tianweiy/DMD2"
14
+ subfolder = "model/sdxl/sdxl_cond999_8node_lr5e-7_denoising4step_diffusion1000_gan5e-3_guidance8_noinit_noode_backsim_scratch_checkpoint_model_019000"
15
+ filename = "pytorch_model.bin"
16
+
17
+
18
+ class ModelWrapper:
19
+ def __init__(self, model_id, checkpoint_path, precision, image_resolution, latent_resolution, num_train_timesteps, conditioning_timestep, num_step, revision, accelerator):
20
+ super().__init__()
21
+ torch.set_grad_enabled(False)
22
+
23
+ self.DTYPE = getattr(torch, precision)
24
+ self.device = accelerator.device
25
+
26
+ self.tokenizer_one = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
27
+ self.tokenizer_two = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
28
+
29
+ self.text_encoder = SDXLTextEncoder(model_id, revision, accelerator, dtype=self.DTYPE)
30
+
31
+ self.vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").float().to(self.device)
32
+ self.vae_dtype = torch.float32
33
+
34
+ self.tiny_vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=self.DTYPE).to(self.device)
35
+ self.tiny_vae_dtype = self.DTYPE
36
+
37
+ self.model = self.create_generator(model_id, checkpoint_path).to(dtype=self.DTYPE).to(self.device)
38
+
39
+ self.accelerator = accelerator
40
+ self.image_resolution = image_resolution
41
+ self.latent_resolution = latent_resolution
42
+ self.num_train_timesteps = num_train_timesteps
43
+ self.vae_downsample_ratio = image_resolution // latent_resolution
44
+ self.conditioning_timestep = conditioning_timestep
45
+
46
+ self.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
47
+ self.alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
48
+ self.num_step = num_step
49
+
50
+ def create_generator(self, model_id, checkpoint_path):
51
+ generator = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(self.DTYPE)
52
+ state_dict = torch.load(checkpoint_path, map_location="cpu")
53
+ generator.load_state_dict(state_dict, strict=True)
54
+ generator.requires_grad_(False)
55
+ return generator
56
+
57
+ def build_condition_input(self, height, width):
58
+ original_size = (height, width)
59
+ target_size = (height, width)
60
+ crop_top_left = (0, 0)
61
+
62
+ add_time_ids = list(original_size + crop_top_left + target_size)
63
+ add_time_ids = torch.tensor([add_time_ids], device=self.device, dtype=self.DTYPE)
64
+ return add_time_ids
65
+
66
+ def _encode_prompt(self, prompt):
67
+ text_input_ids_one = self.tokenizer_one([prompt], padding="max_length", max_length=self.tokenizer_one.model_max_length, truncation=True, return_tensors="pt").input_ids
68
+ text_input_ids_two = self.tokenizer_two([prompt], padding="max_length", max_length=self.tokenizer_two.model_max_length, truncation=True, return_tensors="pt").input_ids
69
+
70
+ prompt_dict = {
71
+ 'text_input_ids_one': text_input_ids_one.unsqueeze(0).to(self.device),
72
+ 'text_input_ids_two': text_input_ids_two.unsqueeze(0).to(self.device)
73
+ }
74
+ return prompt_dict
75
+
76
+ @staticmethod
77
+ def _get_time():
78
+ torch.cuda.synchronize()
79
+ return time.time()
80
+
81
+ def sample(self, noise, unet_added_conditions, prompt_embed, fast_vae_decode):
82
+ alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
83
+
84
+ if self.num_step == 1:
85
+ all_timesteps = [self.conditioning_timestep]
86
+ step_interval = 0
87
+ elif self.num_step == 4:
88
+ all_timesteps = [999, 749, 499, 249]
89
+ step_interval = 250
90
+ else:
91
+ raise NotImplementedError()
92
+
93
+ DTYPE = prompt_embed.dtype
94
+
95
+ for constant in all_timesteps:
96
+ current_timesteps = torch.ones(len(prompt_embed), device=self.device, dtype=torch.long) * constant
97
+ eval_images = self.model(noise, current_timesteps, prompt_embed, added_cond_kwargs=unet_added_conditions).sample
98
+
99
+ eval_images = get_x0_from_noise(noise, eval_images, alphas_cumprod, current_timesteps).to(self.DTYPE)
100
+
101
+ next_timestep = current_timesteps - step_interval
102
+ noise = self.scheduler.add_noise(eval_images, torch.randn_like(eval_images), next_timestep).to(DTYPE)
103
+
104
+ if fast_vae_decode:
105
+ eval_images = self.tiny_vae.decode(eval_images.to(self.tiny_vae_dtype) / self.tiny_vae.config.scaling_factor, return_dict=False)[0]
106
+ else:
107
+ eval_images = self.vae.decode(eval_images.to(self.vae_dtype) / self.vae.config.scaling_factor, return_dict=False)[0]
108
+ eval_images = ((eval_images + 1.0) * 127.5).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1)
109
+ return eval_images
110
+
111
+ @spaces.GPU(enable_queue=True)
112
+ @torch.no_grad()
113
+ def inference(self, prompt, seed, height, width, num_images, fast_vae_decode):
114
+ print("Running model inference...")
115
+
116
+ if seed == -1:
117
+ seed = np.random.randint(0, 1000000)
118
+
119
+ generator = torch.manual_seed(seed)
120
+
121
+ add_time_ids = self.build_condition_input(height, width).repeat(num_images, 1)
122
+
123
+ noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device=self.device, dtype=self.DTYPE)
124
+
125
+ prompt_inputs = self._encode_prompt(prompt)
126
+
127
+ start_time = self._get_time()
128
+
129
+ prompt_embeds, pooled_prompt_embeds = self.text_encoder(prompt_inputs)
130
+
131
+ batch_prompt_embeds, batch_pooled_prompt_embeds = (
132
+ prompt_embeds.repeat(num_images, 1, 1),
133
+ pooled_prompt_embeds.repeat(num_images, 1, 1)
134
+ )
135
+
136
+ unet_added_conditions = {
137
+ "time_ids": add_time_ids,
138
+ "text_embeds": batch_pooled_prompt_embeds.squeeze(1)
139
+ }
140
+
141
+ eval_images = self.sample(noise=noise, unet_added_conditions=unet_added_conditions, prompt_embed=batch_prompt_embeds, fast_vae_decode=fast_vae_decode)
142
+
143
+ end_time = self._get_time()
144
+
145
+ output_image_list = []
146
+ for image in eval_images:
147
+ output_image_list.append(PIL.Image.fromarray(image.cpu().numpy()))
148
+
149
+ return output_image_list, f"Run successfully in {(end_time-start_time):.2f} seconds"
150
+
151
+ def get_x0_from_noise(sample, model_output, alphas_cumprod, timestep):
152
+ alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
153
+ beta_prod_t = 1 - alpha_prod_t
154
+
155
+ pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
156
+ return pred_original_sample
157
+
158
+ class SDXLTextEncoder(torch.nn.Module):
159
+ def __init__(self, model_id, revision, accelerator, dtype=torch.float32):
160
+ super().__init__()
161
+
162
+ self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(accelerator.device).to(dtype=dtype)
163
+ self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_id, subfolder="text_encoder_2", revision=revision).to(accelerator.device).to(dtype=dtype)
164
+
165
+ self.accelerator = accelerator
166
+
167
+ def forward(self, batch):
168
+ text_input_ids_one = batch['text_input_ids_one'].to(self.accelerator.device).squeeze(1)
169
+ text_input_ids_two = batch['text_input_ids_two'].to(self.accelerator.device).squeeze(1)
170
+ prompt_embeds_list = []
171
+
172
+ for text_input_ids, text_encoder in zip([text_input_ids_one, text_input_ids_two], [self.text_encoder_one, self.text_encoder_two]):
173
+ prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device), output_hidden_states=True)
174
+
175
+ pooled_prompt_embeds = prompt_embeds[0]
176
+
177
+ prompt_embeds = prompt_embeds.hidden_states[-2]
178
+ bs_embed, seq_len, _ = prompt_embeds.shape
179
+ prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
180
+ prompt_embeds_list.append(prompt_embeds)
181
+
182
+ prompt_embeds = torch.cat(prompt_embeds_list, dim=-1)
183
+ pooled_prompt_embeds = pooled_prompt_embeds.view(len(text_input_ids_one), -1)
184
+
185
+ return prompt_embeds, pooled_prompt_embeds
186
+
187
+ def create_demo():
188
+ TITLE = "# DMD2-SDXL Demo"
189
+ model_id = "stabilityai/stable-diffusion-xl-base-1.0"
190
+ checkpoint_path = hf_hub_download(repo_id=repo_id, subfolder=subfolder,filename=filename)
191
+ precision = "float16"
192
+ image_resolution = 1024
193
+ latent_resolution = 128
194
+ num_train_timesteps = 1000
195
+ conditioning_timestep = 999
196
+ num_step = 4
197
+ revision = None
198
+
199
+ torch.backends.cuda.matmul.allow_tf32 = True
200
+ torch.backends.cudnn.allow_tf32 = True
201
+
202
+ accelerator = Accelerator()
203
+
204
+ model = ModelWrapper(model_id, checkpoint_path, precision, image_resolution, latent_resolution, num_train_timesteps, conditioning_timestep, num_step, revision, accelerator)
205
+
206
+ with gr.Blocks() as demo:
207
+ gr.Markdown(TITLE)
208
+ with gr.Row():
209
+ with gr.Column():
210
+ prompt = gr.Text(value="An oil painting of two rabbits in the style of American Gothic, wearing the same clothes as in the original.", label="Prompt")
211
+ run_button = gr.Button("Run")
212
+ with gr.Accordion(label="Advanced options", open=True):
213
+ seed = gr.Slider(label="Seed", minimum=-1, maximum=1000000, step=1, value=0)
214
+ num_images = gr.Slider(label="Number of generated images", minimum=1, maximum=16, step=1, value=16)
215
+ fast_vae_decode = gr.Checkbox(label="Use Tiny VAE for faster decoding", value=True)
216
+ height = gr.Slider(label="Image Height", minimum=512, maximum=1536, step=64, value=1024)
217
+ width = gr.Slider(label="Image Width", minimum=512, maximum=1536, step=64, value=1024)
218
+ with gr.Column():
219
+ result = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=1024)
220
+ error_message = gr.Text(label="Job Status")
221
+
222
+ inputs = [prompt, seed, height, width, num_images, fast_vae_decode]
223
+ run_button.click(fn=model.inference, inputs=inputs, outputs=[result, error_message], concurrency_limit=1)
224
+ return demo
225
+
226
+ if __name__ == "__main__":
227
+ demo = create_demo()
228
+ demo.queue(api_open=False)
229
+ demo.launch(show_error=True, share=True)