update code

app.py

@@ -4,6 +4,7 @@ import sys
 import warnings
 import logging
 import spaces
+
 # Configure logging settings
 logging.basicConfig(
     level=logging.INFO,
@@ -52,6 +53,19 @@ if _get_output(["nvcc", "--version"]) is None:
     logging.info("installCUDA: %s" % _get_output(["nvcc", "--version"]))
 else:
     logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))
+
+import torch
+
+# 设置当前默认 GPU 设备（推荐在 CUDA 初始化前设置）
+torch.cuda.set_device(0)
+
+# 显式初始化 CUDA（通常是可选的，但在多线程中有助于避免问题）
+torch.cuda.init()
+
+# 测试
+print("CUDA available:", torch.cuda.is_available())
+print("Current device:", torch.cuda.current_device())
+print("Device name:", torch.cuda.get_device_name(0))
 print("CUDA_HOME =", os.environ.get("CUDA_HOME"))
 from torch.utils.cpp_extension import CUDA_HOME
 print("CUDA_HOME from PyTorch:", CUDA_HOME)

data_process/lib/faceverse_process/core/FaceVerseModel_v3.py

@@ -104,7 +104,7 @@ def get_renderer(img_size, device, R=None, T=None, K=None, orthoCam=False, raste
 
 
 class FaceVerseModel(nn.Module):
-    def __init__(self, model_dict, batch_size=1, device='cuda:1', expr_52=True, **kargs):
+    def __init__(self, model_dict, batch_size=1, device='cuda', expr_52=True, **kargs):
         super(FaceVerseModel, self).__init__()
 
         self.batch_size = batch_size