Run controller updates inside @spaces.GPU to keep CUDA ops in GPU emulation

app.py

@@ -176,23 +176,18 @@ class MagicTimeController:
         height_slider,
         seed_textbox,
     ):
-        # Switch adapters on CPU side first (no GPU needed); this also avoids
-        # pickling `self` into the ZeroGPU worker, which fails because Swift
-        # adapters attach unpicklable local `device_hook` closures on Linear.
-        if self.selected_motion_module != motion_module_dropdown: self.update_motion_module(motion_module_dropdown)
-        if self.selected_motion_module != motion_module_dropdown: self.update_motion_module_2(motion_module_dropdown)
-        if self.selected_dreambooth != dreambooth_dropdown: self.update_dreambooth(dreambooth_dropdown)
-
-        while self.text_encoder is None or self.unet is None:
-            self.update_dreambooth(dreambooth_dropdown, motion_module_dropdown)
-
+        # Delegate to a module-level @spaces.GPU function so `self` (the
+        # controller, which contains Swift-modified nn.Modules with unpicklable
+        # local `device_hook` closures on Linear) is NOT pickled into the
+        # ZeroGPU worker. All GPU work happens via the global `controller`.
         return _magictime_gpu(
+            dreambooth_dropdown,
+            motion_module_dropdown,
             prompt_textbox,
             negative_prompt_textbox,
             int(width_slider),
             int(height_slider),
             str(seed_textbox),
-            dreambooth_dropdown,
         )
 
 controller = MagicTimeController()
@@ -200,15 +195,22 @@ controller = MagicTimeController()
 
 @spaces.GPU(duration=120)
 def _magictime_gpu(
+    dreambooth_dropdown,
+    motion_module_dropdown,
     prompt_textbox,
     negative_prompt_textbox,
     width_slider,
     height_slider,
     seed_textbox,
-    dreambooth_dropdown,
 ):
-    # Use the module-level `controller` global so we don't pickle `self`
-    # (which contains Swift-modified nn.Modules with unpicklable closure hooks).
+    # Use the module-level `controller` global so we don't pickle `self`.
+    if controller.selected_motion_module != motion_module_dropdown: controller.update_motion_module(motion_module_dropdown)
+    if controller.selected_motion_module != motion_module_dropdown: controller.update_motion_module_2(motion_module_dropdown)
+    if controller.selected_dreambooth != dreambooth_dropdown: controller.update_dreambooth(dreambooth_dropdown)
+
+    while controller.text_encoder is None or controller.unet is None:
+        controller.update_dreambooth(dreambooth_dropdown, motion_module_dropdown)
+
     torch.cuda.empty_cache()
     time.sleep(1)