Method encode

lib_layerdiffuse/vae.py:280–292

(self, img_rgba, img_rgb, padded_img_rgb, use_offset=True)

        return origin_pixel, y

    def encode(self, img_rgba, img_rgb, padded_img_rgb, use_offset=True):
        a_bchw_01 = img_rgba[:, 3:, :, :]
        vae_feed = img_rgb.to(device=self.sd_vae.device, dtype=self.sd_vae.dtype)
        latent_dist = self.sd_vae.encode(vae_feed).latent_dist
        offset_feed = torch.cat([padded_img_rgb, a_bchw_01], dim=1).to(device=self.sd_vae.device, dtype=self.dtype)
        offset = self.encoder(offset_feed) * self.alpha
        if use_offset:
            latent = dist_sample_deterministic(dist=latent_dist, perturbation=offset)
            latent = self.sd_vae.config.scaling_factor * (latent - self.sd_vae.config.shift_factor)
        else:
            latent = latent_dist.sample()
            latent = self.sd_vae.config.scaling_factor * (latent - self.sd_vae.config.shift_factor)
        return latent
    
    def forward(self, img_rgba, img_rgb, padded_img_rgb, use_offset=True):
        return self.decode(self.encode(img_rgba, img_rgb, padded_img_rgb, use_offset))