diff --git a/anime_facesaver.py b/anime_facesaver.py index 9e586b9..1bfb2e1 100644 --- a/anime_facesaver.py +++ b/anime_facesaver.py @@ -121,14 +121,25 @@ def parse_arguments(): def resolve_gpu(gpu_arg): - """Pin CUDA before torch initializes. Returns torch device string.""" + """Pin CUDA before torch initializes. Returns torch device string. + + NOTE: ultralytics' select_device() rewrites CUDA_VISIBLE_DEVICES on + first predict() if the CUDA context is not yet initialized, silently + un-pinning us onto GPU0. We force-initialize a tiny CUDA allocation + here so our CVD pin is locked in before ultralytics can touch it. + """ if gpu_arg.lower() == "cpu": os.environ["CUDA_VISIBLE_DEVICES"] = "-1" return "cpu" if gpu_arg.lower() == "auto": return None # let ultralytics/torch decide os.environ["CUDA_VISIBLE_DEVICES"] = gpu_arg - return "cuda" + import torch + if torch.cuda.is_available(): + torch.zeros(1, device="cuda") # lock the context to our pin + else: + return "cpu" + return "cuda:0" # within our CVD, physical GPU N is always cuda:0 def ensure_directory(directory): @@ -300,7 +311,7 @@ def main(): from ultralytics import YOLO import cv2 - if device == "cuda" and not torch.cuda.is_available(): + if device == "cuda:0" and not torch.cuda.is_available(): print("CUDA requested but not visible after pinning " f"(--gpu {args.gpu}). Falling back to CPU.") device = "cpu" diff --git a/audit_stills.py b/audit_stills.py index 7908cd0..57876dd 100644 --- a/audit_stills.py +++ b/audit_stills.py @@ -40,8 +40,16 @@ def main(): os.environ["CUDA_VISIBLE_DEVICES"] = "-1" device = "cpu" elif args.gpu.lower() != "auto": + # Lock the pin with a warm CUDA alloc BEFORE ultralytics runs + # select_device(), which would otherwise rewrite + # CUDA_VISIBLE_DEVICES and un-pin us onto GPU0. os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu - device = "cuda" + import torch + if torch.cuda.is_available(): + torch.zeros(1, device="cuda") + device = "cuda:0" + else: + device = "cpu" else: device = None