Update app.py

app.py

@@ -158,23 +158,34 @@ def infer(
     if latent_file:  # Check if a latent file is provided
         sd_image_a = torch.load(latent_file.name)  # Load the latent
         print("-- using latent file --")
-    print('-- generating image --')
-    with torch.no_grad():
-      sd_image = pipe(
-        prompt=enhanced_prompt,  # This conversion is fine
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator
-      ).images[0]
-    print('-- got image --')
+        print('-- generating image --')
+        with torch.no_grad():
+          sd_image = pipe(
+            prompt=enhanced_prompt,  # This conversion is fine
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            latent=sd_image_a,
+            generator=generator
+          ).images[0]
+    else:
+        with torch.no_grad():
+          sd_image = pipe(
+            prompt=enhanced_prompt,  # This conversion is fine
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator
+          ).images[0]
+        print('-- got image --')
     image_path = f"sd35m_{seed}.png"
     sd_image.save(image_path,optimize=False,compress_level=0)
     upload_to_ftp(image_path)
-    
-    # Convert the generated image to a tensor
+        # Convert the generated image to a tensor
     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():