Method __call__

infer.py:30–72

(
        self,
        prompt: Union[torch.FloatTensor, PIL.Image.Image],
        glyph: Union[torch.FloatTensor, PIL.Image.Image],
        mask_image: Union[torch.FloatTensor, PIL.Image.Image],
        mask: Union[torch.FloatTensor, PIL.Image.Image],
        num_inference_steps: int = 50,
        device=None
    )

    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[torch.FloatTensor, PIL.Image.Image],
        glyph: Union[torch.FloatTensor, PIL.Image.Image],
        mask_image: Union[torch.FloatTensor, PIL.Image.Image],
        mask: Union[torch.FloatTensor, PIL.Image.Image],
        num_inference_steps: int = 50,
        device=None
    ):
        if mask_image is None:
            raise ValueError("`mask_image` input cannot be undefined.")

        batch_size = prompt.shape[0]
        vae.to(device)
        unet.to(device)
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # Preprocess mask and image
        vae_scale_factor = self.vae_scale_factor
        _, _, mask_height, mask_width = mask.size()
        mask = torch.nn.functional.interpolate(mask, size=[mask_width // vae_scale_factor, mask_height // vae_scale_factor])

        glyph_latents = vae.encode(glyph).latent_dist.sample() * vae.config.scaling_factor
        masked_image_latents = vae.encode(mask_image).latent_dist.sample() * vae.config.scaling_factor

        shape = (batch_size, vae.config.latent_channels, mask_height // vae_scale_factor, mask_width // vae_scale_factor)
        latents = randn_tensor(shape, generator=torch.manual_seed(20), device=device) * self.scheduler.init_noise_sigma

        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for t in timesteps:
                latent_model_input = latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                # glyph_latents
                sample = torch.cat([latent_model_input, masked_image_latents, glyph_latents, mask], dim=1)
                noise_pred = unet(sample=sample, timestep=t, encoder_hidden_states=prompt, ).sample
                latents = self.scheduler.step(noise_pred, t, latents).prev_sample
                progress_bar.update()

        pred_latents = latents / vae.config.scaling_factor
        image_vae = vae.decode(pred_latents).sample
        image = (image_vae / 2 + 0.5) * 255.0
        return image, image_vae

# Initialize models
vae = AutoencoderKL.from_pretrained("./model/vae_ft/checkpoint-150000/vae")