Update app.py

app.py

@@ -41,6 +41,9 @@ def load_dit_model(dit_size):
     # Load checkpoint
     checkpoint = torch.load(ckpt_path, map_location="cpu")
     model.load_state_dict(checkpoint["model_state_dict"])
+
+    # Use half precision to speed up sampling
+    model = model.half()
     
     return model
 
@@ -98,7 +101,7 @@ class DiffusionSampler:
         s_null = torch.tensor([model.num_styles] * num_samples, device=self.device, dtype=torch.long)
         
         # Start with random latents
-        latents = torch.randn((num_samples, 4, 32, 32), device=self.device)
+        latents = torch.randn((num_samples, 4, 32, 32), device=self.device, dtype=torch.float16)
         
         # Use classifier-free guidance for better quality
         cfg_scale = 2.5
@@ -135,7 +138,9 @@ class DiffusionSampler:
         
         # Decode latents to images
         self.load_vae()
-        latents = latents / self.vae.config.scaling_factor
+
+        # Convert latents back to float32 for vae decoding
+        latents = latents.to(dtype=torch.float16) / self.vae.config.scaling_factor
         latents = latents.to(self.device)
         
         progress(0.95, desc="Decoding images...")
@@ -167,9 +172,9 @@ def generate_random_seed():
     return random.randint(0, 2**32 - 1)
 
 MODEL_SAMPLE_LIMITS = {
-    "S": {"min":1, "max": 18, "default": 4},
-    "B": {"min":1, "max": 9, "default": 4},
-    "L": {"min":1, "max": 3, "default": 1}
+    "S": {"min":1, "max": 16, "default": 4},
+    "B": {"min":1, "max": 12, "default": 4},
+    "L": {"min":1, "max": 4, "default": 1}
 }
 
 def update_sample_slider(dit_size):