Update app.py

app.py

@@ -189,7 +189,7 @@ def infer(
     generated_image_tensor = torch.tensor([np.array(sd_image).transpose(2, 0, 1)]).to('cuda') / 255.0
     # Encode the generated image into latents
     with torch.no_grad():
-        generated_latents = vae.encode(generated_image_tensor).latent_dist.sample().mul_(0.18215)
+        generated_latents = vae.encode(generated_image_tensor.to(torch.bfloat16)).latent_dist.sample().mul_(0.18215)
     latent_path = f"sd35m_{seed}.pt"
     # Save the latents to a .pt file
     torch.save(generated_latents, latent_path)