[Admin maintenance] Support new ZeroGPU hardware

Thank you so much for having shared this Space with the community on this demo. We have upgraded the ZeroGPU infra-structure to run on modern blackwell architecture.
For that, we need to upgrade your demo to support that. This PR fixes your demo to work with the new architecture. As this is something we broke on our end, we may merge this PR autonomously. If this breaks unexpectedly or brings unintended consequences, feel free to revert, modify or otherwise. Any issues you can email apolinario@huggingface.co

app.py

@@ -166,7 +166,6 @@ class MagicTimeController:
         _, unexpected = self.unet_model.load_state_dict(motion_module_state_dict, strict=False)
         assert len(unexpected) == 0
 
-    @spaces.GPU(duration=120)
     def magictime(
         self,
         dreambooth_dropdown,
@@ -177,66 +176,90 @@ class MagicTimeController:
         height_slider,
         seed_textbox,
     ):
-        torch.cuda.empty_cache()
-        time.sleep(1)
+        # 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),
+        )
 
-        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)
+controller = MagicTimeController()
 
-        if is_xformers_available(): self.unet.enable_xformers_memory_efficient_attention()
 
-        pipeline = MagicTimePipeline(
-            vae=self.vae, text_encoder=self.text_encoder, tokenizer=self.tokenizer, unet=self.unet,
-            scheduler=DDIMScheduler(**OmegaConf.to_container(self.inference_config.noise_scheduler_kwargs))
-        ).to(device)     
+@spaces.GPU(duration=120)
+def _magictime_gpu(
+    dreambooth_dropdown,
+    motion_module_dropdown,
+    prompt_textbox,
+    negative_prompt_textbox,
+    width_slider,
+    height_slider,
+    seed_textbox,
+):
+    # 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)
 
-        if int(seed_textbox) > 0: seed = int(seed_textbox)
-        else: seed = int(random_seed())
-        torch.manual_seed(seed)
-        
-        assert seed == torch.initial_seed()
-        print(f"### seed: {seed}")
-        
-        generator = torch.Generator(device=device)
-        generator.manual_seed(seed)
-        
-        sample = pipeline(
-            prompt_textbox,
-            negative_prompt     = negative_prompt_textbox,
-            num_inference_steps = 25,
-            guidance_scale      = 8.,
-            width               = width_slider,
-            height              = height_slider,
-            video_length        = 16,
-            generator           = generator,
-        ).videos
-
-        save_sample_path = os.path.join(self.savedir, f"sample.mp4")
-        save_videos_grid(sample, save_sample_path)
-    
-        json_config = {
-            "prompt": prompt_textbox,
-            "n_prompt": negative_prompt_textbox,
-            "width": width_slider,
-            "height": height_slider,
-            "seed": seed,
-            "dreambooth": dreambooth_dropdown,
-        }
-
-        # 修复：将字典序列化为 JSON 字符串
-        json_config_str = json.dumps(json_config, indent=4)
+    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)
-        
-        # 修复：直接返回字符串以配合 gr.Code 组件
-        return save_sample_path, json_config_str
+    torch.cuda.empty_cache()
+    time.sleep(1)
 
-controller = MagicTimeController()
+    if is_xformers_available(): controller.unet.enable_xformers_memory_efficient_attention()
+
+    pipeline = MagicTimePipeline(
+        vae=controller.vae, text_encoder=controller.text_encoder, tokenizer=controller.tokenizer, unet=controller.unet,
+        scheduler=DDIMScheduler(**OmegaConf.to_container(controller.inference_config.noise_scheduler_kwargs))
+    ).to(device)
+
+    if int(seed_textbox) > 0: seed = int(seed_textbox)
+    else: seed = int(random_seed())
+    torch.manual_seed(seed)
+
+    assert seed == torch.initial_seed()
+    print(f"### seed: {seed}")
+
+    generator = torch.Generator(device=device)
+    generator.manual_seed(seed)
+
+    sample = pipeline(
+        prompt_textbox,
+        negative_prompt     = negative_prompt_textbox,
+        num_inference_steps = 25,
+        guidance_scale      = 8.,
+        width               = width_slider,
+        height              = height_slider,
+        video_length        = 16,
+        generator           = generator,
+    ).videos
+
+    save_sample_path = os.path.join(controller.savedir, f"sample.mp4")
+    save_videos_grid(sample, save_sample_path)
+
+    json_config = {
+        "prompt": prompt_textbox,
+        "n_prompt": negative_prompt_textbox,
+        "width": width_slider,
+        "height": height_slider,
+        "seed": seed,
+        "dreambooth": dreambooth_dropdown,
+    }
+
+    json_config_str = json.dumps(json_config, indent=4)
+
+    torch.cuda.empty_cache()
+    time.sleep(1)
+
+    return save_sample_path, json_config_str
 
 def ui():
     with gr.Blocks(css=css) as demo:

requirements.txt

@@ -1,10 +1,6 @@
-# torch==2.2.2
-# torchvision==0.17.2
-# torchaudio==2.2.2
-# xformers==0.0.25.post1
-torch==2.7.1
-torchvision==0.22.1
-torchaudio==2.7.1
+torch==2.8.0
+torchvision==0.23.0
+torchaudio==2.8.0
 imageio==2.27.0
 imageio[ffmpeg]
 imageio[pyav]
@@ -15,11 +11,8 @@ accelerate==0.28.0
 diffusers==0.11.1
 transformers==4.38.2
 huggingface_hub==0.25.2
-# huggingface_hub==0.33.5
-gradio>=3.50.2
 gdown
-triton
 einops
 omegaconf
 safetensors
-spaces
+spaces