Update app.py

app.py

@@ -18,6 +18,9 @@ latent_scale_factor = 0.18215  # Same as in DiTTrainer
 # For tracking progress in UI
 global_progress = 0
 
+# Enable half precision inference
+USE_HALF_PRECISION = True
+
 def load_dit_model(dit_size):
     """Load DiT model of specified size"""
     #ckpt_path = f"./ckpts/DiT_{dit_size}_final.pth"
@@ -45,8 +48,9 @@ def load_dit_model(dit_size):
     return model
 
 class DiffusionSampler:
-    def __init__(self, device="cuda" if torch.cuda.is_available() else "cpu"):
+    def __init__(self, device="cuda" if torch.cuda.is_available() else "cpu", use_half = USE_HALF_PRECISION):
         self.device = device
+        self.use_half = use_half
         self.vae = None
         
         # Pre-compute diffusion parameters
@@ -65,6 +69,14 @@ class DiffusionSampler:
         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)
+
+        # Convert to half precision if needed
+        if self.use_half:
+            self.sqrt_alphas_cumprod = self.sqrt_alphas_cumprod.half()
+            self.sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alpha_cumprod.half()
+            self.sqrt_recip_alphas = self.sqrt_recip_alphas.half()
+            self.betas = self.betas.half()
+            self.posterior_variance = self.posterior_variance.half()
         
     def load_vae(self):
         """Load VAE model (done lazily to save memory until needed)"""
@@ -99,6 +111,8 @@ class DiffusionSampler:
         
         # Start with random latents
         latents = torch.randn((num_samples, 4, 32, 32), device=self.device)
+        if self.use_half:
+            latents = latents.half()
         
         # Use classifier-free guidance for better quality
         cfg_scale = 2.5
@@ -135,6 +149,10 @@ class DiffusionSampler:
         
         # Decode latents to images
         self.load_vae()
+
+        # Convert back to float
+        if self.use_half:
+            latents = latents.float()
         latents = latents / self.vae.config.scaling_factor
         latents = latents.to(self.device)