enforce bf16

app.py

@@ -56,7 +56,8 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
         # Convert latent channels to 64 as expected by the transformer
         latent_channels = pipe.vae.config.latent_channels
         if latent_channels != 64:
-            latents = torch.nn.Conv2d(latent_channels, 64, kernel_size=1).to(device)(latents)
+            conv = torch.nn.Conv2d(latent_channels, 64, kernel_size=1).to(device, dtype=dtype)
+            latents = conv(latents)
 
         # Reshape latents to match the transformer's input expectations
         latents = latents.view(1, 64, height // 8, width // 8)
@@ -87,6 +88,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
     return image, seed
 
 
+
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",