diff --git a/tools/patch_comfy_sampler_capture.py b/tools/patch_comfy_sampler_capture.py
index e7202b7..6eb9804 100644
--- a/tools/patch_comfy_sampler_capture.py
+++ b/tools/patch_comfy_sampler_capture.py
@@ -10,7 +10,7 @@ if "import os\n" not in source:
 
 old = "            if callback is not None:\n                callback(x[\"i\"], x[\"denoised\"], x[\"x\"], total_steps)\n"
 new = (
-    "            capture_dir = os.getenv(\"H3_SAMPLER_CAPTURE_DIR\")\n"
+    "            capture_dir = os.getenv(\"H3_SAMPLER_CAPTURE_DIR\") or os.getenv(\"H3_CAPTURE_DIR\")\n"
     "            if capture_dir:\n"
     "                os.makedirs(capture_dir, exist_ok=True)\n"
     "                torch.cuda.synchronize(x[\"x\"].device)\n"