Update all files for EO-VAE

_eo_vae/distributions.py

@@ -0,0 +1,22 @@
+# MIT License - Based on CompVis latent-diffusion / diffusers
+# Diagonal Gaussian for VAE latent distribution
+
+import torch
+
+
+class DiagonalGaussianDistribution:
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None:
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean, device=self.parameters.device)
+
+    def sample(self) -> torch.Tensor:
+        return self.mean + self.std * torch.randn_like(self.mean, device=self.parameters.device)
+
+    def mode(self) -> torch.Tensor:
+        return self.mean