Class TransparentVAEDecoder

lib_layerdiffuse/vae.py:334–411

class TransparentVAEDecoder(torch.nn.Module):
    def __init__(self, filename, dtype=torch.float16, *args, **kwargs):
        super().__init__(*args, **kwargs)
        sd = sf.load_file(filename)
        model = UNet1024(in_channels=3, out_channels=4)
        model.load_state_dict(sd, strict=True)
        model.to(dtype=dtype)
        model.eval()
        self.model = model
        self.dtype = dtype
        return

    @torch.no_grad()
    def estimate_single_pass(self, pixel, latent):
        y = self.model(pixel, latent)
        return y

    @torch.no_grad()
    def estimate_augmented(self, pixel, latent):
        args = [
            [False, 0], [False, 1], [False, 2], [False, 3], [True, 0], [True, 1], [True, 2], [True, 3],
        ]

        result = []

        for flip, rok in tqdm(args):
            feed_pixel = pixel.clone()
            feed_latent = latent.clone()

            if flip:
                feed_pixel = torch.flip(feed_pixel, dims=(3,))
                feed_latent = torch.flip(feed_latent, dims=(3,))

            feed_pixel = torch.rot90(feed_pixel, k=rok, dims=(2, 3))
            feed_latent = torch.rot90(feed_latent, k=rok, dims=(2, 3))

            eps = self.estimate_single_pass(feed_pixel, feed_latent).clip(0, 1)
            eps = torch.rot90(eps, k=-rok, dims=(2, 3))

            if flip:
                eps = torch.flip(eps, dims=(3,))

            result += [eps]

        result = torch.stack(result, dim=0)
        median = torch.median(result, dim=0).values
        return median

    @torch.no_grad()
    def forward(self, sd_vae, latent):
        pixel = sd_vae.decode(latent).sample
        pixel = (pixel * 0.5 + 0.5).clip(0, 1).to(self.dtype)
        latent = latent.to(self.dtype)
        result_list = []
        vis_list = []

        for i in range(int(latent.shape[0])):
            y = self.estimate_augmented(pixel[i:i + 1], latent[i:i + 1])