Fix Stateless GPU environment: avoid CUDA init in main process, initialize GPU in worker process

app.py

@@ -384,35 +384,46 @@ def run_graio_demo(args):
 
     os.makedirs(args.audio_save_dir, exist_ok=True)
 
-    # 运行时动态检测 GPU 可用性（参考 Meigen-MultiTalk）
-    # 记录 GPU 信息，以便在 worker 进程中使用相同的 GPU
+    # 在 Stateless GPU 环境中，主进程不能初始化 CUDA
+    # GPU 检测和初始化将在 worker 进程中进行（通过 @spaces.GPU 装饰器）
+    # 在主进程中，我们使用 CPU 加载模型，然后在 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_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.")
+    # 检查是否在 Stateless GPU 环境中（通过检查 SPACE_ID 和 spaces 模块）
+    is_stateless_gpu = os.environ.get("SPACE_ID") is not None
+    
+    if is_stateless_gpu:
+        # Stateless GPU 环境：主进程不能初始化 CUDA，使用 CPU 加载模型
+        logging.info("Stateless GPU environment detected. Loading models on CPU in main process.")
+        logging.info("GPU will be initialized in worker process via @spaces.GPU decorator.")
+        device = -1  # 使用 CPU 加载模型
+    else:
+        # 本地环境：可以正常检测和使用 GPU
+        if torch.cuda.is_available():
+            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
+                else:
+                    logging.warning("CUDA is available but no GPU devices found. Using CPU.")
+                    device = -1  # 使用 CPU
+            except Exception as e:
+                logging.warning(f"GPU detection error: {e}. Using CPU.")
                 device = -1  # 使用 CPU
-        except Exception as e:
-            logging.warning(f"GPU detection error: {e}. Using CPU.")
+        else:
+            logging.warning("No CUDA-compatible GPU found. Using CPU (slower).")
             device = -1  # 使用 CPU
-    else:
-        logging.warning("No CUDA-compatible GPU found. Using CPU (slower).")
-        device = -1  # 使用 CPU
 
     logging.info("Creating AnyTalker pipeline.")
     # 加载模型
@@ -437,15 +448,17 @@ 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():
+        # 在 worker 进程中设置 GPU 设备（Stateless GPU 环境）
+        if torch.cuda.is_available():
             try:
-                torch.cuda.set_device(gpu_device_id)
+                # 如果记录了 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 {gpu_device_id}: {e}")
+                logging.warning(f"Failed to set GPU device: {e}")
         
         input_data = {}
         input_data["prompt"] = img2vid_prompt
@@ -598,25 +611,30 @@ 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
-        if gpu_device_id is not None:
-            try:
-                if torch.cuda.is_available():
-                    # 设置到初始化时记录的 GPU 设备
-                    torch.cuda.set_device(gpu_device_id)
+        # 在 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 using GPU device {current_device} ({current_gpu_name}) - matching initialization")
+                    logging.info(f"Worker process initialized GPU: device {current_device} ({current_gpu_name})")
                     
-                    # 验证 GPU 名称是否匹配（如果记录了名称）
-                    if gpu_name and current_gpu_name != gpu_name:
-                        logging.warning(f"GPU name mismatch: init={gpu_name}, worker={current_gpu_name}")
+                    # 如果模型是在 CPU 上加载的，需要移动到 GPU
+                    if device == -1:
+                        logging.info("Moving models from CPU to GPU in worker process...")
+                        # 注意：这里需要确保模型已经加载，并且可以移动到 GPU
+                        # 由于模型是在主进程加载的，可能需要重新加载或移动
                 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")
+                    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)
 

wan/audio2video_multiID.py

@@ -67,11 +67,11 @@ class WanAF2V:
             use_half (`bool`, *optional*, defaults to False):
                 Whether to use half precision (float16/bfloat16) for model inference. Reduces memory usage.
         """
-        # 如果 CUDA 不可用，自动回退到 CPU
-        if torch.cuda.is_available():
-            self.device = torch.device(f"cuda:{device_id}")
-        else:
+        # 如果 device_id 为 -1 或 CUDA 不可用，使用 CPU
+        if device_id == -1 or not torch.cuda.is_available():
             self.device = torch.device("cpu")
+        else:
+            self.device = torch.device(f"cuda:{device_id}")
         self.config = config
         self.rank = rank
         self.t5_cpu = t5_cpu