Update app.py

app.py

@@ -418,36 +418,40 @@ def apply_prompt(meta_data, new_prompt):
 
 @spaces.GPU(duration=30)
 def on_image_change(filepath):
-    # Extract the filename without extension
     filename = os.path.splitext(os.path.basename(filepath))[0]
 
-    if filename in ["example1", "example3", "example4"]:
+    if filename not in ["example1","example3","example4"]:
+        return filepath, None, None, None, None, None
 
-        meta_data_raw = load_state_from_file(f"assets/{filename}-turbo.pkl")
-
-        global weights
-        _, _, _, weights = meta_data_raw
-
-        global num_inference_steps
-        num_inference_steps = 10
-        scale_value = 7
+    # 1) load the raw state (might contain CUDA tensors!)
+    caption, real_latents, inversed_latents, saved_weights = load_state_from_file(
+        f"assets/{filename}-turbo.pkl"
+    )
 
-        if filename == "example1":
-            scale_value = 8
-            new_prompt = "a photo of a tree, summer, colourful"
+    # 2) immediately move all tensors to CPU
+    real_cpu    = real_latents.detach().cpu()
+    inversed_cpu= [x.detach().cpu() for x in inversed_latents]
 
-        elif filename == "example3":
-            scale_value = 6
-            new_prompt = "a realistic photo of a female warrior, flowing dark purple or black hair, bronze shoulder armour, leather chest piece, sky background with clouds"
+    # 3) repack a truly CPU-only state
+    cpu_meta_data = (caption, real_cpu, inversed_cpu, saved_weights)
 
-        elif filename == "example4":
-            scale_value = 13
-            new_prompt = "a photo of plastic bottle on some sand, beach background, sky background"
+    # 4) your existing logic can then run on GPU, but it will
+    #    consume only CPU tensors and return only CPU tensors
+    _, _, _, global_weights = cpu_meta_data
+    global weights; weights = global_weights
+    num_inference_steps = 10
+    if filename == "example1":
+        scale_value, new_prompt = 8, "a photo of a tree, summer, colourful"
+    elif filename == "example3":
+        scale_value, new_prompt = 6, "a realistic photo of a female warrior, ..."
+    else:
+        scale_value, new_prompt = 13, "a photo of plastic bottle on some sand, ..."
 
-        update_scale(scale_value)
-        img = apply_prompt(meta_data_raw, new_prompt)
+    update_scale(scale_value)
+    img = apply_prompt(cpu_meta_data, new_prompt)
 
-    return filepath, img, meta_data_raw, num_inference_steps, scale_value, scale_value
+    # 5) return only CPU objects (img is still a numpy array)
+    return filepath, img, cpu_meta_data, num_inference_steps, scale_value, scale_value
 
 
 def update_value(value, layer_type, resolution, depth):