naver-ai/DenseDiffusion
38
1import gradio as gr2import numpy as np3import torch4import requests 5import random6import os7import sys8import pickle9from PIL import Image10 11from tqdm.auto import tqdm12from datetime import datetime13 14import diffusers15from diffusers import DDIMScheduler16from transformers import CLIPTextModel, CLIPTokenizer17import torch.nn.functional as F18 19from utils import preprocess_mask, process_sketch, process_prompts, process_example20 21 22#################################################23#################################################24canvas_html = "<div id='canvas-root' style='max-width:400px; margin: 0 auto'></div>"25load_js = """26async () => {27const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.js"28fetch(url)29 .then(res => res.text())30 .then(text => {31 const script = document.createElement('script');32 script.type = "module"33 script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));34 document.head.appendChild(script);35 });36}37"""38 39get_js_colors = """40async (canvasData) => {41 const canvasEl = document.getElementById("canvas-root");42 return [canvasEl._data]43}44"""45 46css = '''47#color-bg{display:flex;justify-content: center;align-items: center;}48.color-bg-item{width: 100%; height: 32px}49#main_button{width:100%}50<style>51'''52 53 54#################################################55#################################################56global sreg, creg, sizereg, COUNT, creg_maps, sreg_maps, pipe, text_cond57 58sreg = 059creg = 060sizereg = 061COUNT = 062reg_sizes = {}63creg_maps = {}64sreg_maps = {}65text_cond = 066device="cuda"67MAX_COLORS = 1268 69pipe = diffusers.StableDiffusionPipeline.from_pretrained(70 "runwayml/stable-diffusion-v1-5",71 variant="fp16").to(device)72 73pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)74pipe.scheduler.set_timesteps(50)75timesteps = pipe.scheduler.timesteps76sp_sz = pipe.unet.sample_size77 78with open('./valset.pkl', 'rb') as f:79 val_prompt = pickle.load(f)80 81 82#################################################83#################################################84def mod_forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None):85 86 residual = hidden_states87 88 if self.spatial_norm is not None:89 hidden_states = self.spatial_norm(hidden_states, temb)90 91 input_ndim = hidden_states.ndim92 93 if input_ndim == 4:94 batch_size, channel, height, width = hidden_states.shape95 hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)96 97 batch_size, sequence_length, _ = (hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape)98 attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)99 100 if self.group_norm is not None:101 hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)102 103 query = self.to_q(hidden_states)104 105 global sreg, creg, COUNT, creg_maps, sreg_maps, reg_sizes, text_cond106 107 sa_ = True if encoder_hidden_states is None else False108 encoder_hidden_states = text_cond if encoder_hidden_states is not None else hidden_states109 110 if self.norm_cross:111 encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)112 113 key = self.to_k(encoder_hidden_states)114 value = self.to_v(encoder_hidden_states)115 116 query = self.head_to_batch_dim(query)117 key = self.head_to_batch_dim(key)118 value = self.head_to_batch_dim(value)119 120 if COUNT/32 < 50*0.3:121 122 dtype = query.dtype123 if self.upcast_attention:124 query = query.float()125 key = key.float()126 127 sim = torch.baddbmm(torch.empty(query.shape[0], query.shape[1], key.shape[1], 128 dtype=query.dtype, device=query.device),129 query, key.transpose(-1, -2), beta=0, alpha=self.scale)130 131 treg = torch.pow(timesteps[COUNT//32]/1000, 5)132 133 ## reg at self-attn134 if sa_:135 min_value = sim[int(sim.size(0)/2):].min(-1)[0].unsqueeze(-1)136 max_value = sim[int(sim.size(0)/2):].max(-1)[0].unsqueeze(-1) 137 mask = sreg_maps[sim.size(1)].repeat(self.heads,1,1)138 size_reg = reg_sizes[sim.size(1)].repeat(self.heads,1,1)139 140 sim[int(sim.size(0)/2):] += (mask>0)*size_reg*sreg*treg*(max_value-sim[int(sim.size(0)/2):])141 sim[int(sim.size(0)/2):] -= ~(mask>0)*size_reg*sreg*treg*(sim[int(sim.size(0)/2):]-min_value)142 143 ## reg at cross-attn144 else:145 min_value = sim[int(sim.size(0)/2):].min(-1)[0].unsqueeze(-1)146 max_value = sim[int(sim.size(0)/2):].max(-1)[0].unsqueeze(-1) 147 mask = creg_maps[sim.size(1)].repeat(self.heads,1,1)148 size_reg = reg_sizes[sim.size(1)].repeat(self.heads,1,1)149 150 sim[int(sim.size(0)/2):] += (mask>0)*size_reg*creg*treg*(max_value-sim[int(sim.size(0)/2):])151 sim[int(sim.size(0)/2):] -= ~(mask>0)*size_reg*creg*treg*(sim[int(sim.size(0)/2):]-min_value)152 153 attention_probs = sim.softmax(dim=-1)154 attention_probs = attention_probs.to(dtype)155 156 else:157 attention_probs = self.get_attention_scores(query, key, attention_mask)158 159 COUNT += 1160 161 hidden_states = torch.bmm(attention_probs, value)162 hidden_states = self.batch_to_head_dim(hidden_states)163 164 # linear proj165 hidden_states = self.to_out[0](hidden_states)166 # dropout167 hidden_states = self.to_out[1](hidden_states)168 169 if input_ndim == 4:170 hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)171 172 if self.residual_connection:173 hidden_states = hidden_states + residual174 175 hidden_states = hidden_states / self.rescale_output_factor176 177 return hidden_states178 179for _module in pipe.unet.modules():180 if _module.__class__.__name__ == "Attention":181 _module.__class__.__call__ = mod_forward182 183 184#################################################185#################################################186def process_generation(binary_matrixes, seed, creg_, sreg_, sizereg_, bsz, master_prompt, *prompts):187 188 global creg, sreg, sizereg189 creg, sreg, sizereg = creg_, sreg_, sizereg_190 191 clipped_prompts = prompts[:len(binary_matrixes)]192 prompts = [master_prompt] + list(clipped_prompts)193 layouts = torch.cat([preprocess_mask(mask_, sp_sz, sp_sz, device) for mask_ in binary_matrixes])194 195 text_input = pipe.tokenizer(prompts, padding="max_length", return_length=True, return_overflowing_tokens=False, 196 max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt")197 cond_embeddings = pipe.text_encoder(text_input.input_ids.to(device))[0]198 199 uncond_input = pipe.tokenizer([""]*bsz, padding="max_length", max_length=pipe.tokenizer.model_max_length,200 truncation=True, return_tensors="pt")201 uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(device))[0]202 203 204 ###########################205 ###### prep for sreg ###### 206 ###########################207 global sreg_maps, reg_sizes208 sreg_maps = {}209 reg_sizes = {}210 211 for r in range(4):212 res = int(sp_sz/np.power(2,r))213 layouts_s = F.interpolate(layouts,(res, res),mode='nearest')214 layouts_s = (layouts_s.view(layouts_s.size(0),1,-1)*layouts_s.view(layouts_s.size(0),-1,1)).sum(0).unsqueeze(0).repeat(bsz,1,1)215 reg_sizes[np.power(res, 2)] = 1-sizereg*layouts_s.sum(-1, keepdim=True)/(np.power(res, 2))216 sreg_maps[np.power(res, 2)] = layouts_s217 218 219 ###########################220 ###### prep for creg ######221 ###########################222 pww_maps = torch.zeros(1,77,sp_sz,sp_sz).to(device)223 for i in range(1,len(prompts)):224 wlen = text_input['length'][i] - 2225 widx = text_input['input_ids'][i][1:1+wlen]226 for j in range(77):227 try:228 if (text_input['input_ids'][0][j:j+wlen] == widx).sum() == wlen:229 pww_maps[:,j:j+wlen,:,:] = layouts[i-1:i]230 cond_embeddings[0][j:j+wlen] = cond_embeddings[i][1:1+wlen]231 break232 except:233 raise gr.Error("Please check whether every segment prompt is included in the full text !")234 return235 236 global creg_maps237 creg_maps = {}238 for r in range(4):239 res = int(sp_sz/np.power(2,r))240 layout_c = F.interpolate(pww_maps,(res,res),mode='nearest').view(1,77,-1).permute(0,2,1).repeat(bsz,1,1)241 creg_maps[np.power(res, 2)] = layout_c242 243 244 ########################### 245 #### prep for text_emb ####246 ###########################247 global text_cond248 text_cond = torch.cat([uncond_embeddings, cond_embeddings[:1].repeat(bsz,1,1)]) 249 250 global COUNT251 COUNT = 0252 253 if seed == -1:254 latents = torch.randn(bsz,4,sp_sz,sp_sz).to(device)255 else:256 latents = torch.randn(bsz,4,sp_sz,sp_sz, generator=torch.Generator().manual_seed(seed)).to(device)257 258 image = pipe(prompts[:1]*bsz, latents=latents).images259 260 return(image)261 262 263#################################################264#################################################265### define the interface266with gr.Blocks(css=css) as demo:267 binary_matrixes = gr.State([])268 color_layout = gr.State([])269 gr.Markdown('''## DenseDiffusion: Dense Text-to-Image Generation with Attention Modulation''')270 gr.Markdown('''271 #### 😺 Instruction to generate images 😺 <br>272 (1) Create the image layout. <br>273 (2) Label each segment with a text prompt. <br>274 (3) Adjust the full text. The default full text is automatically concatenated from each segment's text. The default one works well, but refineing the full text will further improve the result. <br>275 (4) Check the generated images, and tune the hyperparameters if needed. <br>276 - w<sup>c</sup> : The degree of attention modulation at cross-attention layers. <br>277 - w<sup>s</sup> : The degree of attention modulation at self-attention layers. <br>278 ''')279 280 with gr.Row():281 with gr.Box(elem_id="main-image"):282 canvas_data = gr.JSON(value={}, visible=False)283 canvas = gr.HTML(canvas_html)284 button_run = gr.Button("(1) I've finished my sketch ! 😺", elem_id="main_button", interactive=True)285 286 prompts = []287 colors = []288 color_row = [None] * MAX_COLORS289 with gr.Column(visible=False) as post_sketch:290 for n in range(MAX_COLORS):291 if n == 0 :292 with gr.Row(visible=False) as color_row[n]:293 colors.append(gr.Image(shape=(100, 100), label="background", type="pil", image_mode="RGB", width=100, height=100))294 prompts.append(gr.Textbox(label="Prompt for the background (white region)", value=""))295 else:296 with gr.Row(visible=False) as color_row[n]:297 colors.append(gr.Image(shape=(100, 100), label="segment "+str(n), type="pil", image_mode="RGB", width=100, height=100))298 prompts.append(gr.Textbox(label="Prompt for the segment "+str(n)))299 300 get_genprompt_run = gr.Button("(2) I've finished segment labeling ! 😺", elem_id="prompt_button", interactive=True)301 302 with gr.Column(visible=False) as gen_prompt_vis:303 general_prompt = gr.Textbox(value='', label="(3) Textual Description for the entire image", interactive=True)304 with gr.Accordion("(4) Tune the hyperparameters", open=False):305 creg_ = gr.Slider(label=" w\u1D9C (The degree of attention modulation at cross-attention layers) ", minimum=0, maximum=2., value=1.0, step=0.1)306 sreg_ = gr.Slider(label=" w \u02E2 (The degree of attention modulation at self-attention layers) ", minimum=0, maximum=2., value=0.3, step=0.1)307 sizereg_ = gr.Slider(label="The degree of mask-area adaptive adjustment", minimum=0, maximum=1., value=1., step=0.1)308 bsz_ = gr.Slider(label="Number of Samples to generate", minimum=1, maximum=4, value=1, step=1)309 seed_ = gr.Slider(label="Seed", minimum=-1, maximum=999999999, value=-1, step=1)310 311 final_run_btn = gr.Button("Generate ! 😺")312 313 layout_path = gr.Textbox(label="layout_path", visible=False)314 all_prompts = gr.Textbox(label="all_prompts", visible=False)315 316 with gr.Column():317 out_image = gr.Gallery(label="Result", columns=2, height='auto')318 319 button_run.click(process_sketch, inputs=[canvas_data], outputs=[post_sketch, binary_matrixes, *color_row, *colors], _js=get_js_colors, queue=False)320 321 get_genprompt_run.click(process_prompts, inputs=[binary_matrixes, *prompts], outputs=[gen_prompt_vis, general_prompt], queue=False)322 323 final_run_btn.click(process_generation, inputs=[binary_matrixes, seed_, creg_, sreg_, sizereg_, bsz_, general_prompt, *prompts], outputs=out_image)324 325 gr.Examples(326 examples=[['0.png', '***'.join([val_prompt[0]['textual_condition']] + val_prompt[0]['segment_descriptions']), 381940206],327 ['1.png', '***'.join([val_prompt[1]['textual_condition']] + val_prompt[1]['segment_descriptions']), 307504592],328 ['5.png', '***'.join([val_prompt[5]['textual_condition']] + val_prompt[5]['segment_descriptions']), 114972190]],329 inputs=[layout_path, all_prompts, seed_],330 outputs=[post_sketch, binary_matrixes, *color_row, *colors, *prompts, gen_prompt_vis, general_prompt, seed_],331 fn=process_example,332 run_on_click=True,333 label='😺 Examples 😺',334 )335 336 demo.load(None, None, None, _js=load_js)337 338demo.launch(debug=True)