| 12 | from tqdm import tqdm |
| 13 | |
| 14 | class StableDiffusionPipeline(DiffusionPipeline): |
| 15 | def __init__(self, vae: AutoencoderKL, unet: UNet2DConditionModel, scheduler: DDPMScheduler): |
| 16 | super().__init__() |
| 17 | self.register_modules(vae=vae, unet=unet, scheduler=scheduler) |
| 18 | self.vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1) |
| 19 | |
| 20 | @property |
| 21 | def _execution_device(self): |
| 22 | if not hasattr(self.unet, "_hf_hook"): |
| 23 | return self.device |
| 24 | for module in self.unet.modules(): |
| 25 | if hasattr(module, "_hf_hook") and hasattr(module._hf_hook, "execution_device"): |
| 26 | return torch.device(module._hf_hook.execution_device) |
| 27 | return self.device |
| 28 | |
| 29 | @torch.no_grad() |
| 30 | def __call__( |
| 31 | self, |
| 32 | prompt: Union[torch.FloatTensor, PIL.Image.Image], |
| 33 | glyph: Union[torch.FloatTensor, PIL.Image.Image], |
| 34 | mask_image: Union[torch.FloatTensor, PIL.Image.Image], |
| 35 | mask: Union[torch.FloatTensor, PIL.Image.Image], |
| 36 | num_inference_steps: int = 50, |
| 37 | device=None |
| 38 | ): |
| 39 | if mask_image is None: |
| 40 | raise ValueError("`mask_image` input cannot be undefined.") |
| 41 | |
| 42 | batch_size = prompt.shape[0] |
| 43 | vae.to(device) |
| 44 | unet.to(device) |
| 45 | self.scheduler.set_timesteps(num_inference_steps, device=device) |
| 46 | timesteps = self.scheduler.timesteps |
| 47 | |
| 48 | # Preprocess mask and image |
| 49 | vae_scale_factor = self.vae_scale_factor |
| 50 | _, _, mask_height, mask_width = mask.size() |
| 51 | mask = torch.nn.functional.interpolate(mask, size=[mask_width // vae_scale_factor, mask_height // vae_scale_factor]) |
| 52 | |
| 53 | glyph_latents = vae.encode(glyph).latent_dist.sample() * vae.config.scaling_factor |
| 54 | masked_image_latents = vae.encode(mask_image).latent_dist.sample() * vae.config.scaling_factor |
| 55 | |
| 56 | shape = (batch_size, vae.config.latent_channels, mask_height // vae_scale_factor, mask_width // vae_scale_factor) |
| 57 | latents = randn_tensor(shape, generator=torch.manual_seed(20), device=device) * self.scheduler.init_noise_sigma |
| 58 | |
| 59 | with self.progress_bar(total=num_inference_steps) as progress_bar: |
| 60 | for t in timesteps: |
| 61 | latent_model_input = latents |
| 62 | latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) |
| 63 | # glyph_latents |
| 64 | sample = torch.cat([latent_model_input, masked_image_latents, glyph_latents, mask], dim=1) |
| 65 | noise_pred = unet(sample=sample, timestep=t, encoder_hidden_states=prompt, ).sample |
| 66 | latents = self.scheduler.step(noise_pred, t, latents).prev_sample |
| 67 | progress_bar.update() |
| 68 | |
| 69 | pred_latents = latents / vae.config.scaling_factor |
| 70 | image_vae = vae.decode(pred_latents).sample |
| 71 | image = (image_vae / 2 + 0.5) * 255.0 |