Update app.py

app.py

@@ -156,7 +156,17 @@ def infer(
     print('-- filtered prompt --')
     print(enhanced_prompt)
     if latent_file:  # Check if a latent file is provided
-        sd_image_a = torch.load(latent_file.name).to('cuda')  # Load the latent
+        initial_latents = pipeline.prepare_latents(
+            batch_size=1,
+            num_channels_latents=pipeline.unet.in_channels,
+            height=pipeline.unet.sample_size[0],
+            width=pipeline.unet.sample_size[1],
+            dtype=pipeline.unet.dtype,
+            device=pipeline.device,
+            generator=generator,
+        )
+        sd_image_a = torch.load(latent_file.name)  # Load the latent
+        initial_latents += sd_image_a
         #sd_image_b = pipe.vae.encode(sd_image_a.to(torch.bfloat16)).latent_dist.sample().mul_(0.18215)
         print("-- using latent file --")
         print('-- generating image --')
@@ -168,7 +178,7 @@ def infer(
             num_inference_steps=num_inference_steps,
             width=width,
             height=height,
-            latents=sd_image_a,
+            latents=initial_latents,
             generator=generator
           ).images[0]
     else:
@@ -193,9 +203,19 @@ def infer(
     # Encode the generated image into latents
     with torch.no_grad():
         generated_latents = vae.encode(generated_image_tensor.to(torch.bfloat16)).latent_dist.sample().mul_(0.18215)
-    latent_path = f"sd35m_{seed}.pt"
+        initial_latents = pipeline.prepare_latents(
+            batch_size=1,
+            num_channels_latents=pipeline.unet.in_channels,
+            height=pipeline.unet.sample_size[0],
+            width=pipeline.unet.sample_size[1],
+            dtype=pipeline.unet.dtype,
+            device=pipeline.device,
+            generator=generator,
+            )
+        initial_latents += generated_latents
+        latent_path = f"sd35m_{seed}.pt"
     # Save the latents to a .pt file
-    torch.save(generated_latents, latent_path)
+    torch.save(initial_latents, latent_path)
     upload_to_ftp(latent_path)
     #refiner.scheduler.set_timesteps(num_inference_steps,device)
     refine = refiner(