Update app.py

app.py

@@ -259,7 +259,7 @@ def reconstruct(input_img: Image.Image, caption: str):
     image = pipe.vae.decode(
         latents / pipe.vae.config.scaling_factor, return_dict=False
     )[0]
-    img_np = image.squeeze(0).float().permute(1, 2, 0).cpu()
+    img_np = image.squeeze(0).float().permute(1, 2, 0).detach().cpu()
     img_np = ((img_np / 2 + 0.5).clamp(0, 1).numpy() * 255).astype(np.uint8)
 
     update_scale(12)                       # initial cross-attn value
@@ -303,7 +303,7 @@ def apply_prompt(meta_data: Any, new_prompt: str):
         latents.unsqueeze(0) / pipe.vae.config.scaling_factor,
         return_dict=False
     )[0]
-    img_np = image.squeeze(0).float().permute(1, 2, 0).cpu()
+    img_np = image.squeeze(0).float().permute(1, 2, 0).detach().cpu()
     img_np = ((img_np / 2 + 0.5).clamp(0, 1).numpy() * 255).astype(np.uint8)
 
     pipe.to("cpu"); torch.cuda.empty_cache()