Move models from CPU to GPU in worker process (like LivePortrait)

app.py

@@ -600,6 +600,39 @@ def run_graio_demo(args):
     # @spaces.GPU 装饰器会自动处理 GPU 初始化，不需要手动初始化
     @spaces.GPU(duration=360)
     def gpu_wrapped_generate_video(*args, **kwargs):
+        # 在 worker 进程中将模型移动到 GPU（如果模型在 CPU 上）
+        # 参考 LivePortrait: 在 worker 进程中直接使用 .to("cuda")
+        if torch.cuda.is_available() and device == -1:
+            try:
+                logging.info("Moving models from CPU to GPU in worker process...")
+                cuda_device = torch.device("cuda")
+                
+                # 移动主模型到 GPU
+                if hasattr(wan_a2v, 'model') and wan_a2v.model is not None:
+                    wan_a2v.model = wan_a2v.model.to(cuda_device)
+                
+                # 移动 VAE 模型到 GPU
+                if hasattr(wan_a2v, 'vae') and wan_a2v.vae is not None:
+                    if hasattr(wan_a2v.vae, 'model'):
+                        wan_a2v.vae.model = wan_a2v.vae.model.to(cuda_device)
+                    # 移动 VAE 的 mean 和 std 张量
+                    if hasattr(wan_a2v.vae, 'mean'):
+                        wan_a2v.vae.mean = wan_a2v.vae.mean.to(cuda_device)
+                    if hasattr(wan_a2v.vae, 'std'):
+                        wan_a2v.vae.std = wan_a2v.vae.std.to(cuda_device)
+                
+                # 移动 CLIP 模型到 GPU
+                if hasattr(wan_a2v, 'clip') and wan_a2v.clip is not None:
+                    if hasattr(wan_a2v.clip, 'model'):
+                        wan_a2v.clip.model = wan_a2v.clip.model.to(cuda_device)
+                
+                # 更新设备信息
+                wan_a2v.device = cuda_device
+                
+                logging.info("Models moved to GPU successfully")
+            except Exception as e:
+                logging.warning(f"Failed to move models to GPU: {e}")
+        
         return generate_video(*args, **kwargs)
 
     def toggle_audio_inputs(person_num):