From 940c59d6c98136bf96a1003b7e4ce4ff357ba230 Mon Sep 17 00:00:00 2001
From: Alex <alex@matrixwide.com>
Date: Sun, 18 Oct 2026 09:10:55 -0400
Subject: [PATCH] Lock GPU pin with warm CUDA alloc before ultralytics
 select_device runs

Ultralytics select_device() rewrites CUDA_VISIBLE_DEVICES on first
predict() when the CUDA context is uninitialized: --gpu 1 silently
became a two-GPU cuda:0,1 context that OOMed on GPU0 (Strata's card).
Initializing torch and touching cuda once after setting CVD locks the
pin; both scripts now use device cuda:0 within the pin.
---
 anime_facesaver.py | 17 ++++++++++++++---
 audit_stills.py    | 10 +++++++++-
 2 files changed, 23 insertions(+), 4 deletions(-)

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