[Admin maintenance] Support new ZeroGPU hardware (#5)

- [Admin maintenance] Support new ZeroGPU hardware (8b8c48f79ac89d1c80bf6a94634ef1917f2f5ed9)


Co-authored-by: Apolinário from multimodal AI art <multimodalart@users.noreply.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