Update app.py

app.py

@@ -215,13 +215,13 @@ def reconstruct(input_img, caption):
 
         update_scale(12)
 
-        real_gpu = real_image_initial_latents.detach()
-        inversed_gpu = [x.detach() for x in inversed_latents]
+        real_cpu = real_image_initial_latents.detach().cpu()
+        inversed_cpu = [x.detach().cpu() for x in inversed_latents]
 
         return image_np, caption, 12, [
             caption,
-            real_gpu,
-            inversed_gpu,
+            real_cpu,
+            inversed_cpu,
             weights
         ]
 
@@ -358,10 +358,21 @@ def apply_prompt(meta_data, new_prompt):
     pipe, _, scheduler = load_pipeline()
     pipe.to("cuda")
 
-    caption, real_image_initial_latents, inversed_latents, saved_weights = meta_data
+    caption, real_latents_cpu, inversed_latents_cpu, saved_weights = meta_data
+    
     # overwrite the global so your processor can see it
     global weights
     weights = saved_weights
+
+    # move everything onto CUDA (and match dtype if needed)
+    device = next(pipe.unet.parameters()).device
+    dtype  = next(pipe.unet.parameters()).dtype
+
+    real_latents    = real_latents_cpu.to(device=device, dtype=dtype)
+    inversed_latents = [x.to(device=device, dtype=dtype) for x in inversed_latents_cpu]
+
+    # now all your latents live on CUDA, so the callback won't mix devices
+    initial_latents = torch.cat([real_latents] * 2)
     
     negative_prompt = ""