diff --git a/packages/ai/python/colorize.py b/packages/ai/python/colorize.py index 6aebbd00..0698dca1 100644 --- a/packages/ai/python/colorize.py +++ b/packages/ai/python/colorize.py @@ -17,23 +17,25 @@ def emit_progress(percent, stage): print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + DDCOLOR_MODEL_PATH = os.environ.get( "DDCOLOR_MODEL_PATH", - "/opt/models/ddcolor/ddcolor.onnx", + os.path.join(_MODELS_BASE, "ddcolor", "ddcolor.onnx"), ) # OpenCV DNN fallback model paths (lightweight ~17 MB) OPENCV_PROTO_PATH = os.environ.get( "OPENCV_COLORIZE_PROTO", - "/opt/models/colorize-opencv/colorization_deploy_v2.prototxt", + os.path.join(_MODELS_BASE, "colorize-opencv", "colorization_deploy_v2.prototxt"), ) OPENCV_MODEL_PATH = os.environ.get( "OPENCV_COLORIZE_MODEL", - "/opt/models/colorize-opencv/colorization_release_v2.caffemodel", + os.path.join(_MODELS_BASE, "colorize-opencv", "colorization_release_v2.caffemodel"), ) OPENCV_POINTS_PATH = os.environ.get( "OPENCV_COLORIZE_POINTS", - "/opt/models/colorize-opencv/pts_in_hull.npy", + os.path.join(_MODELS_BASE, "colorize-opencv", "pts_in_hull.npy"), ) diff --git a/packages/ai/python/detect_faces.py b/packages/ai/python/detect_faces.py index 855350ad..9c9507e9 100644 --- a/packages/ai/python/detect_faces.py +++ b/packages/ai/python/detect_faces.py @@ -11,8 +11,10 @@ def emit_progress(percent, stage): # ── Model path for new mp.tasks API ───────────────────────────────── +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite" -_DOCKER_MODEL_PATH = "/opt/models/mediapipe/blaze_face_short_range.tflite" +_DOCKER_MODEL_PATH = os.path.join(_MODELS_BASE, "mediapipe", "blaze_face_short_range.tflite") _LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") _LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.tflite") diff --git a/packages/ai/python/enhance_faces.py b/packages/ai/python/enhance_faces.py index 23131972..cde11b7f 100644 --- a/packages/ai/python/enhance_faces.py +++ b/packages/ai/python/enhance_faces.py @@ -26,21 +26,23 @@ def emit_progress(percent, stage): print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + GFPGAN_MODEL_PATH = os.environ.get( "GFPGAN_MODEL_PATH", - "/opt/models/gfpgan/GFPGANv1.3.pth", + os.path.join(_MODELS_BASE, "gfpgan", "GFPGANv1.3.pth"), ) CODEFORMER_MODEL_PATH = os.environ.get( "CODEFORMER_MODEL_PATH", - "/opt/models/codeformer/codeformer.pth", + os.path.join(_MODELS_BASE, "codeformer", "codeformer.pth"), ) # ── Model path for new mp.tasks API ───────────────────────────────── _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite" -_DOCKER_MODEL_PATH = "/opt/models/mediapipe/blaze_face_short_range.tflite" +_DOCKER_MODEL_PATH = os.path.join(_MODELS_BASE, "mediapipe", "blaze_face_short_range.tflite") _LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") _LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.tflite") diff --git a/packages/ai/python/face_landmarks.py b/packages/ai/python/face_landmarks.py index 16351a67..344e9bf9 100644 --- a/packages/ai/python/face_landmarks.py +++ b/packages/ai/python/face_landmarks.py @@ -78,8 +78,10 @@ def detect_with_solutions(img_array): # ── New API: mp.tasks (mediapipe >= 0.10.30) ─────────────────────── +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task" -_DOCKER_MODEL_PATH = "/opt/models/mediapipe/face_landmarker.task" +_DOCKER_MODEL_PATH = os.path.join(_MODELS_BASE, "mediapipe", "face_landmarker.task") MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") MODEL_PATH = os.path.join(MODEL_DIR, "face_landmarker.task") diff --git a/packages/ai/python/inpaint.py b/packages/ai/python/inpaint.py index 3d345849..33b1e143 100644 --- a/packages/ai/python/inpaint.py +++ b/packages/ai/python/inpaint.py @@ -12,7 +12,8 @@ def emit_progress(percent, stage): # Resolve the LaMa ONNX model path. # Docker places it at /opt/models/lama/lama_fp32.onnx. # For local dev, check a user-writable cache dir. -LAMA_MODEL_DIR = os.environ.get("LAMA_MODEL_DIR", "/opt/models/lama") +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") +LAMA_MODEL_DIR = os.environ.get("LAMA_MODEL_DIR", os.path.join(_MODELS_BASE, "lama")) LAMA_MODEL_PATH = os.path.join(LAMA_MODEL_DIR, "lama_fp32.onnx") LAMA_LOCAL_CACHE = os.path.join(os.path.expanduser("~"), ".cache", "ashim", "lama") LAMA_LOCAL_PATH = os.path.join(LAMA_LOCAL_CACHE, "lama_fp32.onnx") diff --git a/packages/ai/python/noise_removal.py b/packages/ai/python/noise_removal.py index f2b09b8c..2fdb310a 100644 --- a/packages/ai/python/noise_removal.py +++ b/packages/ai/python/noise_removal.py @@ -9,15 +9,17 @@ def emit_progress(percent, stage): print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + # Model paths - Docker locations as defaults, with env var overrides SCUNET_MODEL_PATH = os.environ.get( "SCUNET_MODEL_PATH", - "/opt/models/scunet/scunet_color_real_psnr.pth", + os.path.join(_MODELS_BASE, "scunet", "scunet_color_real_psnr.pth"), ) NAFNET_MODEL_PATH = os.environ.get( "NAFNET_MODEL_PATH", - "/opt/models/nafnet/NAFNet-SIDD-width64.pth", + os.path.join(_MODELS_BASE, "nafnet", "NAFNet-SIDD-width64.pth"), ) # Local cache for dev installs diff --git a/packages/ai/python/red_eye_removal.py b/packages/ai/python/red_eye_removal.py index e4443530..27ecc6e5 100644 --- a/packages/ai/python/red_eye_removal.py +++ b/packages/ai/python/red_eye_removal.py @@ -11,8 +11,10 @@ def emit_progress(percent, stage): # ── Model path for new mp.tasks API ───────────────────────────────── +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + _FACE_MESH_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task" -_DOCKER_MODEL_PATH = "/opt/models/mediapipe/face_landmarker.task" +_DOCKER_MODEL_PATH = os.path.join(_MODELS_BASE, "mediapipe", "face_landmarker.task") _LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") _LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "face_landmarker.task") diff --git a/packages/ai/python/restore.py b/packages/ai/python/restore.py index 5d1fb8ad..2a03d62d 100644 --- a/packages/ai/python/restore.py +++ b/packages/ai/python/restore.py @@ -22,12 +22,14 @@ def emit_progress(percent, stage): # ── Model paths ─────────────────────────────────────────────────────── -LAMA_MODEL_DIR = os.environ.get("LAMA_MODEL_DIR", "/opt/models/lama") +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + +LAMA_MODEL_DIR = os.environ.get("LAMA_MODEL_DIR", os.path.join(_MODELS_BASE, "lama")) LAMA_MODEL_PATH = os.path.join(LAMA_MODEL_DIR, "lama_fp32.onnx") LAMA_LOCAL_CACHE = os.path.join(os.path.expanduser("~"), ".cache", "ashim", "lama") LAMA_LOCAL_PATH = os.path.join(LAMA_LOCAL_CACHE, "lama_fp32.onnx") -CODEFORMER_MODEL_DIR = os.environ.get("CODEFORMER_MODEL_DIR", "/opt/models/codeformer") +CODEFORMER_MODEL_DIR = os.environ.get("CODEFORMER_MODEL_DIR", os.path.join(_MODELS_BASE, "codeformer")) CODEFORMER_MODEL_PATH = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.onnx") CODEFORMER_LOCAL_CACHE = os.path.join( os.path.expanduser("~"), ".cache", "ashim", "codeformer" @@ -35,7 +37,7 @@ CODEFORMER_LOCAL_CACHE = os.path.join( CODEFORMER_LOCAL_PATH = os.path.join(CODEFORMER_LOCAL_CACHE, "codeformer.onnx") DDCOLOR_MODEL_PATH = os.environ.get( - "DDCOLOR_MODEL_PATH", "/opt/models/ddcolor/ddcolor.onnx" + "DDCOLOR_MODEL_PATH", os.path.join(_MODELS_BASE, "ddcolor", "ddcolor.onnx") ) LAMA_MODEL_SIZE = 512 @@ -220,7 +222,7 @@ def _get_codeformer_path(): # ── Model path for new mp.tasks API ───────────────────────────────── _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite" -_FACE_DETECT_DOCKER_PATH = "/opt/models/mediapipe/blaze_face_short_range.tflite" +_FACE_DETECT_DOCKER_PATH = os.path.join(_MODELS_BASE, "mediapipe", "blaze_face_short_range.tflite") _FACE_DETECT_LOCAL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") _FACE_DETECT_LOCAL_PATH = os.path.join(_FACE_DETECT_LOCAL_DIR, "blaze_face_short_range.tflite") diff --git a/packages/ai/python/upscale.py b/packages/ai/python/upscale.py index c0107902..8792992f 100644 --- a/packages/ai/python/upscale.py +++ b/packages/ai/python/upscale.py @@ -26,14 +26,16 @@ def emit_progress(percent, stage): print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") + REALESRGAN_MODEL_PATH = os.environ.get( "REALESRGAN_MODEL_PATH", - "/opt/models/realesrgan/RealESRGAN_x4plus.pth", + os.path.join(_MODELS_BASE, "realesrgan", "RealESRGAN_x4plus.pth"), ) GFPGAN_MODEL_PATH = os.environ.get( "GFPGAN_MODEL_PATH", - "/opt/models/gfpgan/GFPGANv1.3.pth", + os.path.join(_MODELS_BASE, "gfpgan", "GFPGANv1.3.pth"), )