Class VAEEncoder

python_coreml_stable_diffusion/torch2coreml.py:600–610

Wrapper nn.Module wrapper for pipe.encode() method

    }

    class VAEEncoder(nn.Module):
        """ Wrapper nn.Module wrapper for pipe.encode() method
        """

        def __init__(self):
            super().__init__()
            self.quant_conv = pipe.vae.quant_conv.to(dtype=torch.float32)
            self.encoder = pipe.vae.encoder.to(dtype=torch.float32)

        def forward(self, x):
            return self.quant_conv(self.encoder(x))

    baseline_encoder = VAEEncoder().eval()