Update app.py

app.py

@@ -62,12 +62,12 @@ class DiffusionSampler:
         self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), self.alphas_cumprod[:-1]])
         self.posterior_variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
         
-        # Move to device
-        self.sqrt_alphas_cumprod = self.sqrt_alphas_cumprod.to(self.device)
-        self.sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alpha_cumprod.to(self.device)
-        self.sqrt_recip_alphas = self.sqrt_recip_alphas.to(self.device)
-        self.betas = self.betas.to(self.device)
-        self.posterior_variance = self.posterior_variance.to(self.device)
+        # Move to device in half precision
+        self.sqrt_alphas_cumprod = self.sqrt_alphas_cumprod.to(self.device).half()
+        self.sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alpha_cumprod.to(self.device).half()
+        self.sqrt_recip_alphas = self.sqrt_recip_alphas.to(self.device).half()
+        self.betas = self.betas.to(self.device).half()
+        self.posterior_variance = self.posterior_variance.to(self.device).half()
         
     def load_vae(self):
         """Load VAE model (done lazily to save memory until needed)"""