| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from copy import deepcopy | |
| from torchvision.transforms.functional import rgb_to_grayscale | |
| import segmentation_models_pytorch as smp | |
| from diffusers import StableDiffusionInpaintPipeline | |
| from diffusers.utils.torch_utils import randn_tensor | |
| from transformers import PretrainedConfig, PreTrainedModel | |
| class SDGrayInpaintConfig(PretrainedConfig): | |
| model_type = "sd_gray_inpaint" | |
| def __init__( | |
| self, | |
| base_model="stabilityai/stable-diffusion-2-inpainting", | |
| height=512, | |
| width=512, | |
| **kwargs | |
| ): | |
| self.base_model=base_model | |
| self.height=height | |
| self.width=width | |
| super().__init__(**kwargs) | |
| class SDGrayInpaintModel(PreTrainedModel): | |
| config_class = SDGrayInpaintConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| pipe = StableDiffusionInpaintPipeline.from_pretrained(config.base_model) | |
| self.mask_predictor = smp.Unet( | |
| encoder_name="mit_b4", | |
| encoder_weights="imagenet", | |
| in_channels=3, | |
| classes=1, | |
| ) | |
| self.image_processor = pipe.image_processor | |
| self.scheduler = pipe.scheduler | |
| self.unet = pipe.unet | |
| self.vae = pipe.vae | |
| self.prompt_embeds = nn.Parameter(torch.randn(1,77,1024)) | |
| self.height=config.height | |
| self.width=config.width | |
| def forward( | |
| self, | |
| images_gray_masked, | |
| masks=None, | |
| num_inference_steps=250, | |
| seed=42, | |
| input_type='pil', | |
| output_type='pil' | |
| ): | |
| generator = torch.Generator() | |
| generator.manual_seed(seed) | |
| if input_type=='pil': | |
| images_gray_masked = self.image_processor.preprocess(images_gray_masked, height=self.height, width=self.width).float() | |
| elif input_type=='pt': | |
| images_gray_masked=images_gray_masked | |
| else: | |
| raise ValueError('unsupported input_type') | |
| images_gray_masked = images_gray_masked.to(self.vae.device) | |
| if masks is None: | |
| masks_logits = self.mask_predictor(images_gray_masked) | |
| masks = (torch.sigmoid(masks_logits)>0.5)*1. | |
| masks = masks.float().to(self.vae.device) | |
| images_gray_masked = (1-masks) * images_gray_masked | |
| B, C, H, W = images_gray_masked.shape | |
| prompt_embeds = self.prompt_embeds.repeat(B,1,1) | |
| scheduler = deepcopy(self.scheduler) | |
| scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=self.vae.device) | |
| masked_image_latents = self.vae.encode(images_gray_masked).latent_dist.mode() * self.vae.config.scaling_factor | |
| mask_latents = F.interpolate(masks, size=(self.unet.config.sample_size, self.unet.config.sample_size)) | |
| latents = randn_tensor(masked_image_latents.shape, generator=generator).to(self.device) * self.scheduler.init_noise_sigma | |
| for t in scheduler.timesteps: | |
| latents = scheduler.scale_model_input(latents, t) | |
| latent_model_input = torch.cat([latents, mask_latents, masked_image_latents], dim=1) | |
| noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds)[0] | |
| latents = scheduler.step(noise_pred, t, latents)[0] | |
| latents = latents / self.vae.config.scaling_factor | |
| images_gray_restored = self.vae.decode(latents.detach())[0] | |
| images_gray_restored = images_gray_masked * (1-masks) + images_gray_restored.detach() * masks | |
| images_gray_restored = rgb_to_grayscale(images_gray_restored) | |
| if output_type=='pil': | |
| images_gray_restored = self.image_processor.postprocess(images_gray_restored) | |
| elif output_type=='np': | |
| images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'np') | |
| elif output_type=='pt': | |
| images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'pt') | |
| elif output_type=='none': | |
| images_gray_restored = images_gray_restored | |
| else: | |
| raise ValueError('unsupported output_type') | |
| return images_gray_restored | |