关于推理的问题 #7

Activity

1. chenhaoxing commented on Jan 29, 2024

   

   chenhaoxing

   on Jan 29, 2024

   Owner

   More actions

   class StableDiffusionDocPipeline(DiffusionPipeline):
       def __init__(
           self,
           vae: AutoencoderKL,
           text_encoder: VisionEncoderDecoderModel,
           tokenizer: TrOCRProcessor,
           unet: UNet2DConditionModel,
           scheduler: KarrasDiffusionSchedulers,
       ):
           super().__init__()
           self.register_modules(
               vae=vae,
               text_encoder=text_encoder,
               tokenizer=tokenizer,
               unet=unet,
               scheduler=scheduler,
           )
   
           self.vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
   
       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def _execution_device(self):
           r"""
           Returns the device on which the pipeline's models will be executed. After calling
           `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
           hooks.
           """
           if not hasattr(self.unet, "_hf_hook"):
               return self.device
           for module in self.unet.modules():
               if (
                   hasattr(module, "_hf_hook")
                   and hasattr(module._hf_hook, "execution_device")
                   and module._hf_hook.execution_device is not None
               ):
                   return torch.device(module._hf_hook.execution_device)
           return self.device
           
       @torch.no_grad()
       def __call__(
           self,
           prompt: Union[torch.FloatTensor, PIL.Image.Image] = None,
           location: Optional[np.ndarray]= None,
           image: Union[torch.FloatTensor, PIL.Image.Image] = None,
           mask_image: Union[torch.FloatTensor, PIL.Image.Image] = None,
           mask: Union[torch.FloatTensor, PIL.Image.Image] = None,
           height: Optional[int] = None,
           width: Optional[int] = None,
           num_inference_steps: int = 50,
           guidance_scale: float = 7.5,
           num_images_per_prompt: Optional[int] = 1,
           eta: float = 0.0,
           generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
           latents: Optional[torch.FloatTensor] = None,
           prompt_embeds: Optional[torch.FloatTensor] = None,
           output_type: Optional[str] = "pil",
           return_dict: bool = True,
           # callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
           # callback_steps: int = 1,
       ):
           if image is None:
               raise ValueError("`image` input cannot be undefined.")
   
           if mask_image is None:
               raise ValueError("`mask_image` input cannot be undefined.")
   
           device = self._execution_device
           # 2. Encode input prompt
           # if type(prompt) == "string":
           #     ttf_imgs = draw_text(image[0].shape[:2][::-1], prompt, location)
   
           prompt_images = self.tokenizer(images=prompt, return_tensors="pt").pixel_values.to(device)
           ocr_feature = self.text_encoder(prompt_images)
           prompt_embeds = ocr_feature.last_hidden_state.detach()
           
           # 3. Define call parameters
           batch_size = prompt_embeds.shape[0]
   
           # 4. Preprocess mask and image
           
           # 5. set timesteps
           self.scheduler.set_timesteps(num_inference_steps, device=device)
           timesteps = self.scheduler.timesteps
   
           # 6. Prepare latent variables
           vae = self.vae
           latents = vae.encode(image).latent_dist.sample()
           latents = latents * vae.config.scaling_factor
           noise = torch.randn_like(latents)
   
           # 7. Prepare mask latent variables
           # Rex: prepare mask && mask latent as input of UNET  
           vae_scale_factor = self.vae_scale_factor     
           width, height, *_ = mask.size()[::-1]         
           mask = torch.nn.functional.interpolate(
               mask, size=[width // vae_scale_factor, height // vae_scale_factor, *_][:-2][::-1]
           )
           
           masked_image_latents = vae.encode(mask_image).latent_dist.sample()
           masked_image_latents = masked_image_latents * vae.config.scaling_factor
   
           shape = (1, vae.config.latent_channels, height // vae_scale_factor, width // vae_scale_factor)
   
           latents = randn_tensor(shape, generator=torch.manual_seed(20), device=device,)
           
           # scale the initial noise by the standard deviation required by the scheduler
           latents = latents * self.scheduler.init_noise_sigma
           
           # 10. Denoising loop
           # num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
           with self.progress_bar(total=num_inference_steps) as progress_bar:
               for i, t in enumerate(timesteps):
                   # expand the latents if we are doing classifier free guidance
                   latent_model_input = latents
   
                   # concat latents, mask, masked_image_latents in the channel dimension
                   latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                   latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)
   
                   # predict the noise residual
                   noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample
   
                   # compute the previous noisy sample x_t -> x_t-1
                   latents = self.scheduler.step(noise_pred, t, latents).prev_sample
   
                   progress_bar.update()
                   
           # 11. Post-processing
           pred_latents = 1 / vae.config.scaling_factor * latents
           image_vae = vae.decode(pred_latents).sample
   
           # 13. Convert to PIL
           image = (image_vae / 2 + 0.5) * 255.0
           image = image.cpu().permute(0, 2, 3, 1).float().detach().numpy()
           return image, image_vae
2. prefixRAINSTARsuffix commented on Feb 1, 2024

   

   prefixRAINSTARsuffix

   on Feb 1, 2024

   Author

   More actions

   @chenhaoxing 谢谢你的代码。

   你在文中提到 “After dimension adjustment through a convolution layer, the feature vector 𝑧ˆ𝑡 = Conv(𝑧′𝑡) is fed into the UNet”，但是我在训练代码和推理代码里都没有看到这个卷积层，而是直接把 [latent_model_input, mask, masked_image_latents] 送入了 UNet，这样不会有维度不匹配问题吗。

   我尝试去复现这一部分，但是生成的图片质量很差，请问具体是怎么做的呢。