Update app.py

app.py

@@ -62,13 +62,13 @@ 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 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()
-        
+        # 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)
+    
     def load_vae(self):
         """Load VAE model (done lazily to save memory until needed)"""
         if self.vae is None:
@@ -102,6 +102,13 @@ class DiffusionSampler:
         
         # Start with random latents
         latents = torch.randn((num_samples, 4, 32, 32), device=self.device, dtype=torch.float16)
+
+        # Convert diffusion parameters to half precision for compatibility
+        sqrt_alphas_cumprod_half = self.sqrt_alphas_cumprod.half()
+        sqrt_one_minus_alpha_cumprod_half = self.sqrt_one_minus_alpha_cumprod.half()
+        sqrt_recip_alphas_half = self.sqrt_recip_alphas.half()
+        betas_half = self.betas.half()
+        posterior_variance_half = self.posterior_variance.half()
         
         # Use classifier-free guidance for better quality
         cfg_scale = 2.5
@@ -118,10 +125,10 @@ class DiffusionSampler:
                 
                 t = torch.full((num_samples,), t_val, device=self.device, dtype=torch.long)
 
-                sqrt_recip_alphas_t = self.sqrt_recip_alphas[t].view(-1, 1, 1, 1)
-                sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alpha_cumprod[t].view(-1, 1, 1, 1)
-                beta_t = self.betas[t].view(-1, 1, 1, 1)
-                posterior_variance_t = self.posterior_variance[t].view(-1, 1, 1, 1)
+                sqrt_recip_alphas_t = sqrt_recip_alphas_half[t].view(-1, 1, 1, 1)
+                sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alpha_cumprod_half[t].view(-1, 1, 1, 1)
+                beta_t = betas_half[t].view(-1, 1, 1, 1)
+                posterior_variance_t = posterior_variance_half[t].view(-1, 1, 1, 1)
 
                 # Get noise prediction with classifier-free guidance
                 eps_theta_cond = model(latents, t, g_cond, s_cond)
@@ -130,7 +137,7 @@ class DiffusionSampler:
 
                 # Update latents
                 mean = sqrt_recip_alphas_t * (latents - (beta_t / sqrt_one_minus_alphas_cumprod_t) * eps_theta)
-                noise = torch.randn_like(latents)
+                noise = torch.randn_like(latents,dtype=torch.float16)
                 if t_val == 0:
                     latents = mean
                 else: