Upload region_control.py

region_control.py

@@ -0,0 +1,208 @@
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler
+
+# suppress partial model loading warning
+logging.set_verbosity_error()
+
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+import argparse
+import numpy as np
+from PIL import Image
+
+
+def seed_everything(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    # torch.backends.cudnn.deterministic = True
+    # torch.backends.cudnn.benchmark = True
+
+
+def get_views(panorama_height, panorama_width, window_size=64, stride=8):
+    panorama_height /= 8
+    panorama_width /= 8
+    num_blocks_height = (panorama_height - window_size) // stride + 1
+    num_blocks_width = (panorama_width - window_size) // stride + 1
+    total_num_blocks = int(num_blocks_height * num_blocks_width)
+    views = []
+    for i in range(total_num_blocks):
+        h_start = int((i // num_blocks_width) * stride)
+        h_end = h_start + window_size
+        w_start = int((i % num_blocks_width) * stride)
+        w_end = w_start + window_size
+        views.append((h_start, h_end, w_start, w_end))
+    return views
+
+
+class MultiDiffusion(nn.Module):
+    def __init__(self, device, sd_version='2.0', hf_key=None):
+        super().__init__()
+
+        self.device = device
+        self.sd_version = sd_version
+
+        print(f'[INFO] loading stable diffusion...')
+        if hf_key is not None:
+            print(f'[INFO] using hugging face custom model key: {hf_key}')
+            model_key = hf_key
+        elif self.sd_version == '2.1':
+            model_key = "stabilityai/stable-diffusion-2-1-base"
+        elif self.sd_version == '2.0':
+            model_key = "stabilityai/stable-diffusion-2-base"
+        elif self.sd_version == '1.5':
+            model_key = "runwayml/stable-diffusion-v1-5"
+        else:
+            raise ValueError(f'Stable-diffusion version {self.sd_version} not supported.')
+
+        # Create model
+        self.vae = AutoencoderKL.from_pretrained(model_key, subfolder="vae").to(self.device)
+        self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer")
+        self.text_encoder = CLIPTextModel.from_pretrained(model_key, subfolder="text_encoder").to(self.device)
+        self.unet = UNet2DConditionModel.from_pretrained(model_key, subfolder="unet").to(self.device)
+
+        self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
+
+        print(f'[INFO] loaded stable diffusion!')
+
+    @torch.no_grad()
+    def get_random_background(self, n_samples):
+        # sample random background with a constant rgb value
+        backgrounds = torch.rand(n_samples, 3, device=self.device)[:, :, None, None].repeat(1, 1, 512, 512)
+        return torch.cat([self.encode_imgs(bg.unsqueeze(0)) for bg in backgrounds])
+
+    @torch.no_grad()
+    def get_text_embeds(self, prompt, negative_prompt):
+        # 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, text_embeddings])
+        return text_embeddings
+
+    @torch.no_grad()
+    def encode_imgs(self, imgs):
+        imgs = 2 * imgs - 1
+        posterior = self.vae.encode(imgs).latent_dist
+        latents = posterior.sample() * 0.18215
+        return latents
+
+    @torch.no_grad()
+    def decode_latents(self, latents):
+        latents = 1 / 0.18215 * latents
+        imgs = self.vae.decode(latents).sample
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+        return imgs
+
+    @torch.no_grad()
+    def generate(self, masks, prompts, negative_prompts='', height=512, width=2048, num_inference_steps=50,
+                      guidance_scale=7.5, bootstrapping=20):
+
+        # get bootstrapping backgrounds
+        # can move this outside of the function to speed up generation. i.e., calculate in init
+        bootstrapping_backgrounds = self.get_random_background(bootstrapping)
+
+        # Prompts -> text embeds
+        text_embeds = self.get_text_embeds(prompts, negative_prompts)  # [2 * len(prompts), 77, 768]
+
+        # Define panorama grid and get views
+        latent = torch.randn((1, self.unet.in_channels, height // 8, width // 8), device=self.device)
+        noise = latent.clone().repeat(len(prompts) - 1, 1, 1, 1)
+        views = get_views(height, width)
+        count = torch.zeros_like(latent)
+        value = torch.zeros_like(latent)
+
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        with torch.autocast('cuda'):
+            for i, t in enumerate(self.scheduler.timesteps):
+                count.zero_()
+                value.zero_()
+
+                for h_start, h_end, w_start, w_end in views:
+                    masks_view = masks[:, :, h_start:h_end, w_start:w_end]
+                    latent_view = latent[:, :, h_start:h_end, w_start:w_end].repeat(len(prompts), 1, 1, 1)
+                    if i < bootstrapping:
+                        bg = bootstrapping_backgrounds[torch.randint(0, bootstrapping, (len(prompts) - 1,))]
+                        bg = self.scheduler.add_noise(bg, noise[:, :, h_start:h_end, w_start:w_end], t)
+                        latent_view[1:] = latent_view[1:] * masks_view[1:] + bg * (1 - masks_view[1:])
+
+                    # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+                    latent_model_input = torch.cat([latent_view] * 2)
+
+                    # predict the noise residual
+                    noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeds)['sample']
+
+                    # perform guidance
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                    # compute the denoising step with the reference model
+                    latents_view_denoised = self.scheduler.step(noise_pred, t, latent_view)['prev_sample']
+
+                    value[:, :, h_start:h_end, w_start:w_end] += (latents_view_denoised * masks_view).sum(dim=0,
+                                                                                                          keepdims=True)
+                    count[:, :, h_start:h_end, w_start:w_end] += masks_view.sum(dim=0, keepdims=True)
+
+                # take the MultiDiffusion step
+                latent = torch.where(count > 0, value / count, value)
+
+        # Img latents -> imgs
+        imgs = self.decode_latents(latent)  # [1, 3, 512, 512]
+        img = T.ToPILImage()(imgs[0].cpu())
+        return img
+
+
+def preprocess_mask(mask_path, h, w, device):
+    mask = np.array(Image.open(mask_path).convert("L"))
+    mask = mask.astype(np.float32) / 255.0
+    mask = mask[None, None]
+    mask[mask < 0.5] = 0
+    mask[mask >= 0.5] = 1
+    mask = torch.from_numpy(mask).to(device)
+    mask = torch.nn.functional.interpolate(mask, size=(h, w), mode='nearest')
+    return mask
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--mask_paths', type=list)
+    parser.add_argument('--bg_prompt', type=str)
+    parser.add_argument('--bg_negative', type=str)  # 'artifacts, blurry, smooth texture, bad quality, distortions, unrealistic, distorted image'
+    parser.add_argument('--fg_prompts', type=list)
+    parser.add_argument('--fg_negative', type=list)  # 'artifacts, blurry, smooth texture, bad quality, distortions, unrealistic, distorted image'
+    parser.add_argument('--sd_version', type=str, default='2.0', choices=['1.5', '2.0'],
+                        help="stable diffusion version")
+    parser.add_argument('--H', type=int, default=768)
+    parser.add_argument('--W', type=int, default=512)
+    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--steps', type=int, default=50)
+    parser.add_argument('--bootstrapping', type=int, default=20)
+    opt = parser.parse_args()
+
+    seed_everything(opt.seed)
+
+    device = torch.device('cuda')
+
+    sd = MultiDiffusion(device, opt.sd_version)
+
+    fg_masks = torch.cat([preprocess_mask(mask_path, opt.H // 8, opt.W // 8, device) for mask_path in opt.mask_paths])
+    bg_mask = 1 - torch.sum(fg_masks, dim=0, keepdim=True)
+    bg_mask[bg_mask < 0] = 0
+    masks = torch.cat([bg_mask, fg_masks])
+
+    prompts = [opt.bg_prompt] + opt.fg_prompts
+    neg_prompts = [opt.bg_negative] + opt.fg_negative
+
+    img = sd.generate(masks, prompts, neg_prompts, opt.H, opt.W, opt.steps, bootstrapping=opt.bootstrapping)
+
+    # save image
+    img.save('out.png')