Track GPU device during initialization and use same GPU in worker process and inference

app.py

@@ -385,14 +385,25 @@ def run_graio_demo(args):
     os.makedirs(args.audio_save_dir, exist_ok=True)
 
     # 运行时动态检测 GPU 可用性（参考 Meigen-MultiTalk）
+    # 记录 GPU 信息，以便在 worker 进程中使用相同的 GPU
+    gpu_device_id = None
+    gpu_name = None
+    gpu_uuid = None
+    
     if torch.cuda.is_available():
         try:
             num_gpus = torch.cuda.device_count()
             if num_gpus > 0:
-                gpu_name = torch.cuda.get_device_name(0)
-                logging.info(f"GPU AVAILABLE: {num_gpus} GPU(s), Name: {gpu_name}")
-                # 使用 GPU
-                device = local_rank if world_size > 1 else 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
             else:
                 logging.warning("CUDA is available but no GPU devices found. Using CPU.")
                 device = -1  # 使用 CPU
@@ -426,6 +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):
+        # 确保使用初始化时记录的 GPU 设备
+        if gpu_device_id is not None and torch.cuda.is_available():
+            try:
+                torch.cuda.set_device(gpu_device_id)
+                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 {gpu_device_id}: {e}")
+        
         input_data = {}
         input_data["prompt"] = img2vid_prompt
         input_data["cond_image"] = img2vid_image
@@ -577,16 +598,25 @@ def run_graio_demo(args):
     # 参考: https://huggingface.co/spaces/KlingTeam/LivePortrait/blob/main/app.py
     @spaces.GPU(duration=360)
     def gpu_wrapped_generate_video(*args, **kwargs):
-        # 在 worker 进程中检查 GPU 可用性
-        try:
-            if torch.cuda.is_available():
-                # 尝试访问 GPU 以确保它已准备好
-                _ = torch.cuda.current_device()
-                logging.info(f"GPU ready in worker process: {torch.cuda.get_device_name(0)}")
-            else:
-                logging.warning("GPU not available in worker process, but continuing...")
-        except RuntimeError as e:
-            logging.warning(f"GPU initialization error in worker process: {e}. Continuing anyway...")
+        # 在 worker 进程中确保使用初始化时记录的 GPU
+        if gpu_device_id is not None:
+            try:
+                if torch.cuda.is_available():
+                    # 设置到初始化时记录的 GPU 设备
+                    torch.cuda.set_device(gpu_device_id)
+                    current_device = torch.cuda.current_device()
+                    current_gpu_name = torch.cuda.get_device_name(current_device)
+                    logging.info(f"Worker process using GPU device {current_device} ({current_gpu_name}) - matching initialization")
+                    
+                    # 验证 GPU 名称是否匹配（如果记录了名称）
+                    if gpu_name and current_gpu_name != gpu_name:
+                        logging.warning(f"GPU name mismatch: init={gpu_name}, worker={current_gpu_name}")
+                else:
+                    logging.warning("GPU not available in worker process, but continuing...")
+            except RuntimeError as e:
+                logging.warning(f"GPU initialization error in worker process: {e}. Continuing anyway...")
+        else:
+            logging.info("No GPU device ID recorded, using default device")
         
         return generate_video(*args, **kwargs)