Method encode

stable_diffusion/stable_diffusion/vae.py:260–267

(self, x)

        return self.decoder(self.post_quant_proj(z))

    def encode(self, x):
        x = self.encoder(x)
        x = self.quant_proj(x)
        mean, logvar = x.split(2, axis=-1)
        mean = mean * self.scaling_factor
        logvar = logvar + 2 * math.log(self.scaling_factor)

        return mean, logvar

    def __call__(self, x, key=None):
        mean, logvar = self.encode(x)