Simplify GPU handling following LivePortrait pattern: CPU in main process, GPU in worker via @spaces.GPU

app.py

@@ -384,14 +384,9 @@ def run_graio_demo(args):
 
     os.makedirs(args.audio_save_dir, exist_ok=True)
 
-    # 在 Stateless GPU 环境中，主进程不能初始化 CUDA
-    # GPU 检测和初始化将在 worker 进程中进行（通过 @spaces.GPU 装饰器）
-    # 在主进程中，我们使用 CPU 加载模型，然后在 worker 进程中移动到 GPU
-    gpu_device_id = None
-    gpu_name = None
-    gpu_uuid = None
-    
-    # 检查是否在 Stateless GPU 环境中（通过检查 SPACE_ID 和 spaces 模块）
+    # 参考 LivePortrait: 在 Stateless GPU 环境中，主进程不能初始化 CUDA
+    # 参考: https://huggingface.co/spaces/KlingTeam/LivePortrait/blob/main/app.py
+    # 在主进程中，我们使用 CPU 加载模型，GPU 将在 worker 进程中通过 @spaces.GPU 装饰器初始化
     is_stateless_gpu = os.environ.get("SPACE_ID") is not None
     
     if is_stateless_gpu:
@@ -405,16 +400,10 @@ def run_graio_demo(args):
             try:
                 num_gpus = torch.cuda.device_count()
                 if num_gpus > 0:
-                    gpu_device_id = local_rank if world_size > 1 else 0
-                    torch.cuda.set_device(gpu_device_id)
-                    gpu_name = torch.cuda.get_device_name(gpu_device_id)
-                    # 尝试获取 GPU UUID（如果可用）
-                    try:
-                        gpu_uuid = torch.cuda.get_device_properties(gpu_device_id).uuid
-                    except:
-                        pass
-                    logging.info(f"GPU AVAILABLE: {num_gpus} GPU(s), Device ID: {gpu_device_id}, Name: {gpu_name}, UUID: {gpu_uuid}")
-                    device = gpu_device_id
+                    device = local_rank if world_size > 1 else 0
+                    torch.cuda.set_device(device)
+                    gpu_name = torch.cuda.get_device_name(device)
+                    logging.info(f"GPU AVAILABLE: {num_gpus} GPU(s), Device ID: {device}, Name: {gpu_name}")
                 else:
                     logging.warning("CUDA is available but no GPU devices found. Using CPU.")
                     device = -1  # 使用 CPU
@@ -448,17 +437,16 @@ def run_graio_demo(args):
 
     def generate_video(img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3,
                     sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector):
-        # 在 worker 进程中设置 GPU 设备（Stateless GPU 环境）
+        # 参考 LivePortrait: 在 worker 进程中直接使用 cuda 设备
+        # 参考: https://huggingface.co/spaces/KlingTeam/LivePortrait/blob/main/src/gradio_pipeline.py
+        # @spaces.GPU 装饰器已经初始化了 GPU，这里直接使用即可
         if torch.cuda.is_available():
             try:
-                # 如果记录了 GPU 设备 ID，使用它；否则使用默认设备 0
-                target_device = gpu_device_id if gpu_device_id is not None else 0
-                torch.cuda.set_device(target_device)
                 current_device = torch.cuda.current_device()
                 current_gpu_name = torch.cuda.get_device_name(current_device)
                 logging.info(f"Using GPU device {current_device} ({current_gpu_name}) for inference")
             except Exception as e:
-                logging.warning(f"Failed to set GPU device: {e}")
+                logging.warning(f"GPU check error: {e}")
         
         input_data = {}
         input_data["prompt"] = img2vid_prompt
@@ -609,33 +597,9 @@ def run_graio_demo(args):
 
     # 使用 @spaces.GPU 装饰器包装 generate_video 函数（参考 LivePortrait）
     # 参考: https://huggingface.co/spaces/KlingTeam/LivePortrait/blob/main/app.py
+    # @spaces.GPU 装饰器会自动处理 GPU 初始化，不需要手动初始化
     @spaces.GPU(duration=360)
     def gpu_wrapped_generate_video(*args, **kwargs):
-        # 在 worker 进程中初始化 GPU（Stateless GPU 环境）
-        worker_gpu_device_id = None
-        try:
-            if torch.cuda.is_available():
-                # 在 worker 进程中检测 GPU
-                num_gpus = torch.cuda.device_count()
-                if num_gpus > 0:
-                    worker_gpu_device_id = 0  # 使用第一个 GPU
-                    torch.cuda.set_device(worker_gpu_device_id)
-                    current_device = torch.cuda.current_device()
-                    current_gpu_name = torch.cuda.get_device_name(current_device)
-                    logging.info(f"Worker process initialized GPU: device {current_device} ({current_gpu_name})")
-                    
-                    # 如果模型是在 CPU 上加载的，需要移动到 GPU
-                    if device == -1:
-                        logging.info("Moving models from CPU to GPU in worker process...")
-                        # 注意：这里需要确保模型已经加载，并且可以移动到 GPU
-                        # 由于模型是在主进程加载的，可能需要重新加载或移动
-                else:
-                    logging.warning("No GPU devices found in worker process")
-            else:
-                logging.warning("CUDA not available in worker process")
-        except RuntimeError as e:
-            logging.warning(f"GPU initialization error in worker process: {e}")
-        
         return generate_video(*args, **kwargs)
 
     def toggle_audio_inputs(person_num):