Update app.py

app.py

@@ -158,10 +158,10 @@ def infer(
     if latent_file:  # Check if a latent file is provided
         initial_latents = pipe.prepare_latents(
             batch_size=1,
-            num_channels_latents=pipe.unet.in_channels,
-            height=pipe.unet.sample_size[0],
-            width=pipe.unet.sample_size[1],
-            dtype=pipe.unet.dtype,
+            num_channels_latents=pipe.transformer.in_channels,
+            height=pipe.transformer.sample_size[0],
+            width=pipe.transformer.sample_size[1],
+            dtype=pipe.transformer.dtype,
             device=pipe.device,
             generator=generator,
         )
@@ -205,10 +205,10 @@ def infer(
         generated_latents = vae.encode(generated_image_tensor.to(torch.bfloat16)).latent_dist.sample().mul_(0.18215)
         initial_latents = pipe.prepare_latents(
             batch_size=1,
-            num_channels_latents=pipe.unet.in_channels,
-            height=pipe.unet.sample_size[0],
-            width=pipe.unet.sample_size[1],
-            dtype=pipe.unet.dtype,
+            num_channels_latents=pipe.transformer.in_channels,
+            height=pipe.transformer.sample_size[0],
+            width=pipe.transformer.sample_size[1],
+            dtype=pipe.transformer.dtype,
             device=pipe.device,
             generator=generator,
             )