mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
Siddharth Kumar Sah
85b1cfc10a
Complete rebrand from Stirling-Image to ashim following the project move to https://github.com/ashim-hq/ashim. Changes across 117 files: - Package scope: @stirling-image/* → @ashim/* - GitHub URLs: stirling-image/stirling-image → ashim-hq/ashim - Docker Hub: stirlingimage/stirling-image → ashimhq/ashim - GitHub Pages: stirling-image.github.io → ashim-hq.github.io - All branding text: "Stirling Image" → "ashim" - Docker service/volumes/user: stirling → ashim - Database: stirling.db → ashim.db - localStorage keys: stirling-token → ashim-token - Environment variables: STIRLING_GPU → ASHIM_GPU - Python cache dirs: .cache/stirling-image → .cache/ashim - SVG filter IDs, test prefixes, and all other references
612 lines
24 KiB
Python
612 lines
24 KiB
Python
"""Pre-download and verify all ML models for the Docker image.
|
|
|
|
This script runs at Docker build time. Any failure exits non-zero,
|
|
failing the build. No silent fallbacks.
|
|
"""
|
|
import os
|
|
import sys
|
|
import urllib.request
|
|
|
|
# Some servers (e.g. Berkeley) block the default Python-urllib User-Agent.
|
|
_opener = urllib.request.build_opener()
|
|
_opener.addheaders = [("User-Agent", "ashim/1.0")]
|
|
urllib.request.install_opener(_opener)
|
|
|
|
# Force CPU mode during build - no GPU driver available at build time.
|
|
# Must be set before any ML library import.
|
|
os.environ["PADDLE_DEVICE"] = "cpu"
|
|
os.environ["FLAGS_use_cuda"] = "0"
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
|
|
|
LAMA_MODEL_DIR = "/opt/models/lama"
|
|
LAMA_MODEL_URL = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx"
|
|
LAMA_MODEL_PATH = os.path.join(LAMA_MODEL_DIR, "lama_fp32.onnx")
|
|
LAMA_MIN_SIZE = 100_000_000 # ~200 MB
|
|
|
|
REALESRGAN_MODEL_DIR = "/opt/models/realesrgan"
|
|
REALESRGAN_MODEL_URL = (
|
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
|
|
)
|
|
REALESRGAN_MODEL_PATH = os.path.join(REALESRGAN_MODEL_DIR, "RealESRGAN_x4plus.pth")
|
|
REALESRGAN_MIN_SIZE = 60_000_000 # ~67 MB
|
|
|
|
GFPGAN_MODEL_DIR = "/opt/models/gfpgan"
|
|
GFPGAN_MODEL_URL = (
|
|
"https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
|
|
)
|
|
GFPGAN_MODEL_PATH = os.path.join(GFPGAN_MODEL_DIR, "GFPGANv1.3.pth")
|
|
GFPGAN_MIN_SIZE = 300_000_000 # ~332 MB
|
|
|
|
CODEFORMER_MODEL_DIR = "/opt/models/codeformer"
|
|
CODEFORMER_MODEL_URL = (
|
|
"https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth"
|
|
)
|
|
CODEFORMER_MODEL_PATH = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.pth")
|
|
CODEFORMER_MIN_SIZE = 350_000_000 # ~375 MB
|
|
|
|
DDCOLOR_MODEL_DIR = "/opt/models/ddcolor"
|
|
DDCOLOR_MODEL_URL = (
|
|
"https://huggingface.co/piddnad/DDColor-models/resolve/main/ddcolor_paper_tiny.pth"
|
|
)
|
|
DDCOLOR_ONNX_PATH = os.path.join(DDCOLOR_MODEL_DIR, "ddcolor.onnx")
|
|
DDCOLOR_MIN_SIZE = 50_000_000 # ~220 MB ONNX
|
|
|
|
SCUNET_MODEL_DIR = "/opt/models/scunet"
|
|
SCUNET_MODEL_URL = (
|
|
"https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth"
|
|
)
|
|
SCUNET_MODEL_PATH = os.path.join(SCUNET_MODEL_DIR, "scunet_color_real_psnr.pth")
|
|
SCUNET_MIN_SIZE = 3_000_000 # ~4 MB
|
|
|
|
CODEFORMER_MODEL_DIR = "/opt/models/codeformer"
|
|
CODEFORMER_ONNX_PATH = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.onnx")
|
|
CODEFORMER_MIN_SIZE = 100_000_000 # ~377 MB
|
|
|
|
NAFNET_MODEL_DIR = "/opt/models/nafnet"
|
|
NAFNET_MODEL_URL = (
|
|
"https://huggingface.co/mikestealth/nafnet-models/resolve/main/"
|
|
"NAFNet-SIDD-width64.pth"
|
|
)
|
|
NAFNET_MODEL_PATH = os.path.join(NAFNET_MODEL_DIR, "NAFNet-SIDD-width64.pth")
|
|
NAFNET_MIN_SIZE = 60_000_000 # ~67 MB
|
|
|
|
MEDIAPIPE_MODEL_DIR = "/opt/models/mediapipe"
|
|
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_MODEL_PATH = os.path.join(MEDIAPIPE_MODEL_DIR, "blaze_face_short_range.tflite")
|
|
FACE_DETECT_MIN_SIZE = 100_000 # ~200 KB
|
|
FACE_LANDMARKER_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
|
|
FACE_LANDMARKER_MODEL_PATH = os.path.join(MEDIAPIPE_MODEL_DIR, "face_landmarker.task")
|
|
FACE_LANDMARKER_MIN_SIZE = 1_000_000 # ~7 MB
|
|
|
|
FACEXLIB_MODEL_DIR = "/opt/models/gfpgan/facelib"
|
|
FACEXLIB_DET_URL = "https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth"
|
|
FACEXLIB_DET_PATH = os.path.join(FACEXLIB_MODEL_DIR, "detection_Resnet50_Final.pth")
|
|
FACEXLIB_DET_MIN_SIZE = 100_000_000 # ~104 MB
|
|
FACEXLIB_PARSE_URL = "https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth"
|
|
FACEXLIB_PARSE_PATH = os.path.join(FACEXLIB_MODEL_DIR, "parsing_parsenet.pth")
|
|
FACEXLIB_PARSE_MIN_SIZE = 80_000_000 # ~85 MB
|
|
|
|
OPENCV_COLORIZE_DIR = "/opt/models/colorize-opencv"
|
|
OPENCV_PROTO_URL = "https://raw.githubusercontent.com/richzhang/colorization/caffe/colorization/models/colorization_deploy_v2.prototxt"
|
|
OPENCV_PROTO_PATH = os.path.join(OPENCV_COLORIZE_DIR, "colorization_deploy_v2.prototxt")
|
|
OPENCV_CAFFE_URL = "https://huggingface.co/spaces/BilalSardar/Black-N-White-To-Color/resolve/main/colorization_release_v2.caffemodel"
|
|
OPENCV_CAFFE_PATH = os.path.join(OPENCV_COLORIZE_DIR, "colorization_release_v2.caffemodel")
|
|
OPENCV_CAFFE_MIN_SIZE = 100_000_000 # ~129 MB
|
|
OPENCV_POINTS_URL = "https://raw.githubusercontent.com/richzhang/colorization/caffe/colorization/resources/pts_in_hull.npy"
|
|
OPENCV_POINTS_PATH = os.path.join(OPENCV_COLORIZE_DIR, "pts_in_hull.npy")
|
|
|
|
REMBG_MODELS = [
|
|
"u2net",
|
|
"isnet-general-use",
|
|
"bria-rmbg",
|
|
"birefnet-general-lite",
|
|
"birefnet-portrait",
|
|
"birefnet-general",
|
|
"birefnet-matting",
|
|
]
|
|
|
|
# PaddleOCR PP-OCRv5 HuggingFace model repos to pre-download.
|
|
# These are the models used by PaddleOCR(ocr_version="PP-OCRv5").
|
|
# Downloaded via huggingface_hub to avoid initializing the PaddlePaddle
|
|
# inference engine, which segfaults under QEMU emulation at build time.
|
|
PADDLEOCR_MODELS = [
|
|
"PaddlePaddle/PP-OCRv5_server_det",
|
|
"PaddlePaddle/PP-OCRv5_server_rec",
|
|
"PaddlePaddle/PP-OCRv5_mobile_det",
|
|
"PaddlePaddle/PP-OCRv5_mobile_rec",
|
|
"PaddlePaddle/latin_PP-OCRv5_mobile_rec",
|
|
"PaddlePaddle/korean_PP-OCRv5_mobile_rec",
|
|
"PaddlePaddle/PP-LCNet_x1_0_textline_ori",
|
|
]
|
|
|
|
PADDLEOCR_VL_MODEL = "PaddlePaddle/PaddleOCR-VL-1.5"
|
|
|
|
# PaddleX stores models here by default
|
|
PADDLEX_MODEL_DIR = os.path.expanduser("~/.paddlex/official_models")
|
|
|
|
|
|
def _register_birefnet_matting():
|
|
"""Register BiRefNet-matting ONNX session for Ultra quality mode."""
|
|
import os
|
|
import pooch
|
|
from rembg.sessions import sessions_class
|
|
from rembg.sessions.birefnet_general import BiRefNetSessionGeneral
|
|
|
|
class BiRefNetMattingSession(BiRefNetSessionGeneral):
|
|
@classmethod
|
|
def download_models(cls, *args, **kwargs):
|
|
fname = f"{cls.name(*args, **kwargs)}.onnx"
|
|
pooch.retrieve(
|
|
"https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-matting-epoch_100.onnx",
|
|
None, # Skip checksum for GitHub release assets
|
|
fname=fname,
|
|
path=cls.u2net_home(*args, **kwargs),
|
|
progressbar=True,
|
|
)
|
|
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
|
|
|
|
@classmethod
|
|
def name(cls, *args, **kwargs):
|
|
return "birefnet-matting"
|
|
|
|
sessions_class.append(BiRefNetMattingSession)
|
|
|
|
|
|
def download_rembg_models():
|
|
"""Download all rembg ONNX models."""
|
|
print("=== Downloading rembg models ===")
|
|
from rembg import new_session
|
|
|
|
_register_birefnet_matting()
|
|
|
|
for model in REMBG_MODELS:
|
|
print(f" Downloading {model}...")
|
|
new_session(model)
|
|
print(f" {model} ready")
|
|
print(f"All {len(REMBG_MODELS)} rembg models downloaded.\n")
|
|
|
|
|
|
def download_lama_model():
|
|
"""Download LaMa ONNX inpainting model from HuggingFace."""
|
|
print("=== Downloading LaMa ONNX model ===")
|
|
os.makedirs(LAMA_MODEL_DIR, exist_ok=True)
|
|
print(f" Downloading from {LAMA_MODEL_URL}...")
|
|
urllib.request.urlretrieve(LAMA_MODEL_URL, LAMA_MODEL_PATH)
|
|
|
|
size = os.path.getsize(LAMA_MODEL_PATH)
|
|
assert size > LAMA_MIN_SIZE, (
|
|
f"LaMa model too small: {size} bytes (expected > {LAMA_MIN_SIZE})"
|
|
)
|
|
print(f" lama_fp32.onnx downloaded ({size / 1_000_000:.1f} MB)\n")
|
|
|
|
|
|
def download_realesrgan_model():
|
|
"""Download RealESRGAN_x4plus.pth pretrained weights."""
|
|
print("=== Downloading RealESRGAN model ===")
|
|
os.makedirs(REALESRGAN_MODEL_DIR, exist_ok=True)
|
|
print(f" Downloading from {REALESRGAN_MODEL_URL}...")
|
|
urllib.request.urlretrieve(REALESRGAN_MODEL_URL, REALESRGAN_MODEL_PATH)
|
|
|
|
size = os.path.getsize(REALESRGAN_MODEL_PATH)
|
|
assert size > REALESRGAN_MIN_SIZE, (
|
|
f"RealESRGAN model too small: {size} bytes (expected > {REALESRGAN_MIN_SIZE})"
|
|
)
|
|
print(f" RealESRGAN_x4plus.pth downloaded ({size / 1_000_000:.1f} MB)\n")
|
|
|
|
|
|
def download_gfpgan_model():
|
|
"""Download GFPGANv1.3.pth pretrained weights for face enhancement."""
|
|
print("=== Downloading GFPGAN model ===")
|
|
os.makedirs(GFPGAN_MODEL_DIR, exist_ok=True)
|
|
print(f" Downloading from {GFPGAN_MODEL_URL}...")
|
|
urllib.request.urlretrieve(GFPGAN_MODEL_URL, GFPGAN_MODEL_PATH)
|
|
|
|
size = os.path.getsize(GFPGAN_MODEL_PATH)
|
|
assert size > GFPGAN_MIN_SIZE, (
|
|
f"GFPGAN model too small: {size} bytes (expected > {GFPGAN_MIN_SIZE})"
|
|
)
|
|
print(f" GFPGANv1.3.pth downloaded ({size / 1_000_000:.1f} MB)\n")
|
|
|
|
|
|
def download_codeformer_model():
|
|
"""Download codeformer.pth pretrained weights for face enhancement."""
|
|
print("=== Downloading CodeFormer model ===")
|
|
os.makedirs(CODEFORMER_MODEL_DIR, exist_ok=True)
|
|
print(f" Downloading from {CODEFORMER_MODEL_URL}...")
|
|
urllib.request.urlretrieve(CODEFORMER_MODEL_URL, CODEFORMER_MODEL_PATH)
|
|
|
|
size = os.path.getsize(CODEFORMER_MODEL_PATH)
|
|
assert size > CODEFORMER_MIN_SIZE, (
|
|
f"CodeFormer model too small: {size} bytes (expected > {CODEFORMER_MIN_SIZE})"
|
|
)
|
|
print(f" codeformer.pth downloaded ({size / 1_000_000:.1f} MB)\n")
|
|
|
|
|
|
def download_ddcolor_model():
|
|
"""Download pre-exported DDColor ONNX model for AI photo colorization.
|
|
|
|
Uses the pre-converted ONNX model from HuggingFace (facefusion repo)
|
|
for direct inference via onnxruntime without needing PyTorch.
|
|
"""
|
|
print("=== Downloading DDColor ONNX model ===")
|
|
os.makedirs(DDCOLOR_MODEL_DIR, exist_ok=True)
|
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
print(" Downloading DDColor ONNX from HuggingFace...")
|
|
downloaded_path = hf_hub_download(
|
|
repo_id="facefusion/models-3.0.0",
|
|
filename="ddcolor.onnx",
|
|
local_dir=DDCOLOR_MODEL_DIR,
|
|
)
|
|
|
|
# huggingface_hub downloads to local_dir/filename
|
|
actual_path = os.path.join(DDCOLOR_MODEL_DIR, "ddcolor.onnx")
|
|
if not os.path.exists(actual_path) and os.path.exists(downloaded_path):
|
|
os.rename(downloaded_path, actual_path)
|
|
|
|
size = os.path.getsize(actual_path)
|
|
assert size > DDCOLOR_MIN_SIZE, (
|
|
f"DDColor model too small: {size} bytes (expected > {DDCOLOR_MIN_SIZE})"
|
|
)
|
|
print(f" DDColor ONNX model ready ({size / 1_000_000:.1f} MB)\n")
|
|
|
|
|
|
|
|
def download_codeformer_onnx_model():
|
|
"""Download CodeFormer ONNX model for AI face restoration.
|
|
|
|
Uses the pre-converted ONNX model from HuggingFace (facefusion repo)
|
|
for direct inference via onnxruntime without needing PyTorch.
|
|
"""
|
|
print("=== Downloading CodeFormer ONNX model ===")
|
|
os.makedirs(CODEFORMER_MODEL_DIR, exist_ok=True)
|
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
print(" Downloading CodeFormer ONNX from HuggingFace...")
|
|
downloaded_path = hf_hub_download(
|
|
repo_id="facefusion/models-3.0.0",
|
|
filename="codeformer.onnx",
|
|
local_dir=CODEFORMER_MODEL_DIR,
|
|
)
|
|
|
|
actual_path = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.onnx")
|
|
if not os.path.exists(actual_path) and os.path.exists(downloaded_path):
|
|
os.rename(downloaded_path, actual_path)
|
|
|
|
size = os.path.getsize(actual_path)
|
|
assert size > CODEFORMER_MIN_SIZE, (
|
|
f"CodeFormer model too small: {size} bytes (expected > {CODEFORMER_MIN_SIZE})"
|
|
)
|
|
print(f" CodeFormer ONNX model ready ({size / 1_000_000:.1f} MB)\n")
|
|
|
|
|
|
def download_paddleocr_models():
|
|
"""Pre-download PaddleOCR PP-OCRv5 model weights from HuggingFace.
|
|
|
|
Uses huggingface_hub.snapshot_download() to fetch model files directly
|
|
into the PaddleX cache directory. This avoids initializing PaddlePaddle's
|
|
C++ inference engine, which segfaults under QEMU emulation (arm64 host
|
|
building amd64 image).
|
|
"""
|
|
print("=== Downloading PaddleOCR PP-OCRv5 models ===")
|
|
from huggingface_hub import snapshot_download
|
|
|
|
os.makedirs(PADDLEX_MODEL_DIR, exist_ok=True)
|
|
|
|
for repo_id in PADDLEOCR_MODELS:
|
|
model_name = repo_id.split("/", 1)[1]
|
|
local_dir = os.path.join(PADDLEX_MODEL_DIR, model_name)
|
|
print(f" Downloading {model_name}...")
|
|
snapshot_download(repo_id=repo_id, local_dir=local_dir)
|
|
print(f" {model_name} ready")
|
|
print(f"All {len(PADDLEOCR_MODELS)} PaddleOCR PP-OCRv5 models downloaded.\n")
|
|
|
|
|
|
def download_paddleocr_vl_model():
|
|
"""Pre-download PaddleOCR-VL 1.5 model weights from HuggingFace."""
|
|
print("=== Downloading PaddleOCR-VL 1.5 model ===")
|
|
from huggingface_hub import snapshot_download
|
|
|
|
model_name = PADDLEOCR_VL_MODEL.split("/", 1)[1]
|
|
local_dir = os.path.join(PADDLEX_MODEL_DIR, model_name)
|
|
print(f" Downloading {model_name} (~1.93 GB)...")
|
|
snapshot_download(repo_id=PADDLEOCR_VL_MODEL, local_dir=local_dir)
|
|
print(f" {model_name} ready\n")
|
|
|
|
|
|
def download_scunet_model():
|
|
"""Download SCUNet real-noise denoising model."""
|
|
print(f"Downloading SCUNet model to {SCUNET_MODEL_PATH}...")
|
|
os.makedirs(SCUNET_MODEL_DIR, exist_ok=True)
|
|
urllib.request.urlretrieve(SCUNET_MODEL_URL, SCUNET_MODEL_PATH)
|
|
size = os.path.getsize(SCUNET_MODEL_PATH)
|
|
assert size > SCUNET_MIN_SIZE, (
|
|
f"SCUNet model too small: {size} bytes (expected >{SCUNET_MIN_SIZE})"
|
|
)
|
|
print(f" SCUNet model downloaded: {size:,} bytes")
|
|
|
|
|
|
def download_nafnet_model():
|
|
"""Download NAFNet SIDD width-64 denoising model."""
|
|
print(f"Downloading NAFNet model to {NAFNET_MODEL_PATH}...")
|
|
os.makedirs(NAFNET_MODEL_DIR, exist_ok=True)
|
|
urllib.request.urlretrieve(NAFNET_MODEL_URL, NAFNET_MODEL_PATH)
|
|
size = os.path.getsize(NAFNET_MODEL_PATH)
|
|
assert size > NAFNET_MIN_SIZE, (
|
|
f"NAFNet model too small: {size} bytes (expected >{NAFNET_MIN_SIZE})"
|
|
)
|
|
print(f" NAFNet model downloaded: {size:,} bytes")
|
|
|
|
|
|
def download_facexlib_models():
|
|
"""Download face detection and parsing models used by GFPGAN and CodeFormer.
|
|
|
|
These are auxiliary models from facexlib that GFPGAN/CodeFormer download
|
|
on first use via basicsr. Pre-downloading prevents runtime network access.
|
|
"""
|
|
print("=== Downloading facexlib auxiliary models ===")
|
|
os.makedirs(FACEXLIB_MODEL_DIR, exist_ok=True)
|
|
|
|
print(f" Downloading detection_Resnet50_Final.pth...")
|
|
urllib.request.urlretrieve(FACEXLIB_DET_URL, FACEXLIB_DET_PATH)
|
|
size = os.path.getsize(FACEXLIB_DET_PATH)
|
|
assert size > FACEXLIB_DET_MIN_SIZE, (
|
|
f"Face detection model too small: {size} bytes (expected > {FACEXLIB_DET_MIN_SIZE})"
|
|
)
|
|
print(f" detection_Resnet50_Final.pth downloaded ({size / 1_000_000:.1f} MB)")
|
|
|
|
print(f" Downloading parsing_parsenet.pth...")
|
|
urllib.request.urlretrieve(FACEXLIB_PARSE_URL, FACEXLIB_PARSE_PATH)
|
|
size = os.path.getsize(FACEXLIB_PARSE_PATH)
|
|
assert size > FACEXLIB_PARSE_MIN_SIZE, (
|
|
f"Face parsing model too small: {size} bytes (expected > {FACEXLIB_PARSE_MIN_SIZE})"
|
|
)
|
|
print(f" parsing_parsenet.pth downloaded ({size / 1_000_000:.1f} MB)\n")
|
|
|
|
|
|
def download_opencv_colorize_models():
|
|
"""Download OpenCV DNN colorization models (Zhang et al.).
|
|
|
|
Three files needed for the lightweight OpenCV colorizer fallback.
|
|
"""
|
|
print("=== Downloading OpenCV colorization models ===")
|
|
os.makedirs(OPENCV_COLORIZE_DIR, exist_ok=True)
|
|
|
|
for url, path, name in [
|
|
(OPENCV_PROTO_URL, OPENCV_PROTO_PATH, "colorization_deploy_v2.prototxt"),
|
|
(OPENCV_CAFFE_URL, OPENCV_CAFFE_PATH, "colorization_release_v2.caffemodel"),
|
|
(OPENCV_POINTS_URL, OPENCV_POINTS_PATH, "pts_in_hull.npy"),
|
|
]:
|
|
print(f" Downloading {name}...")
|
|
urllib.request.urlretrieve(url, path)
|
|
size = os.path.getsize(path)
|
|
print(f" {name} downloaded ({size / 1_000_000:.1f} MB)")
|
|
|
|
# Verify the caffemodel (the big one)
|
|
size = os.path.getsize(OPENCV_CAFFE_PATH)
|
|
assert size > OPENCV_CAFFE_MIN_SIZE, (
|
|
f"Caffemodel too small: {size} bytes (expected > {OPENCV_CAFFE_MIN_SIZE})"
|
|
)
|
|
print("OpenCV colorization models downloaded.\n")
|
|
|
|
|
|
def download_mediapipe_task_models():
|
|
"""Download MediaPipe tasks API model files for face detection and landmarks.
|
|
|
|
These models are used by the mp.tasks fallback when mp.solutions is
|
|
unavailable (mediapipe >= 0.10.30). Pre-downloading ensures the Docker
|
|
image works fully airgapped.
|
|
"""
|
|
print("=== Downloading MediaPipe task models ===")
|
|
os.makedirs(MEDIAPIPE_MODEL_DIR, exist_ok=True)
|
|
|
|
for url, path, name, min_size in [
|
|
(FACE_DETECT_MODEL_URL, FACE_DETECT_MODEL_PATH,
|
|
"blaze_face_short_range", FACE_DETECT_MIN_SIZE),
|
|
(FACE_LANDMARKER_MODEL_URL, FACE_LANDMARKER_MODEL_PATH,
|
|
"face_landmarker", FACE_LANDMARKER_MIN_SIZE),
|
|
]:
|
|
print(f" Downloading {name}...")
|
|
urllib.request.urlretrieve(url, path)
|
|
size = os.path.getsize(path)
|
|
assert size > min_size, (
|
|
f"{name} model too small: {size} bytes (expected > {min_size})"
|
|
)
|
|
print(f" {name} downloaded ({size / 1_000_000:.1f} MB)")
|
|
print("MediaPipe task models downloaded.\n")
|
|
|
|
|
|
def verify_mediapipe():
|
|
"""Verify MediaPipe face detection models are bundled in the wheel."""
|
|
print("=== Verifying MediaPipe models ===")
|
|
import mediapipe as mp
|
|
|
|
try:
|
|
for selection in [0, 1]:
|
|
label = "short-range" if selection == 0 else "full-range"
|
|
print(f" Verifying {label} model (selection={selection})...")
|
|
detector = mp.solutions.face_detection.FaceDetection(
|
|
model_selection=selection, min_detection_confidence=0.5
|
|
)
|
|
detector.close()
|
|
print(f" {label} model OK")
|
|
except AttributeError:
|
|
# mediapipe >= 0.10.30 removed mp.solutions; verify tasks API instead
|
|
print(" mp.solutions unavailable, verifying mp.tasks API...")
|
|
options = mp.tasks.vision.FaceDetectorOptions(
|
|
base_options=mp.tasks.BaseOptions(
|
|
model_asset_path=FACE_DETECT_MODEL_PATH
|
|
),
|
|
running_mode=mp.tasks.vision.RunningMode.IMAGE,
|
|
min_detection_confidence=0.5,
|
|
)
|
|
detector = mp.tasks.vision.FaceDetector.create_from_options(options)
|
|
detector.close()
|
|
print(" mp.tasks FaceDetector OK")
|
|
|
|
print("MediaPipe models verified.\n")
|
|
|
|
|
|
def smoke_test():
|
|
"""Final verification that all ML libraries and models are loadable.
|
|
|
|
GPU-dependent libraries (paddlepaddle-gpu, torch CUDA) cannot be imported
|
|
at build time because the CUDA driver is only available at runtime. We verify
|
|
CPU-only imports and check that model files exist on disk.
|
|
"""
|
|
print("=== Running smoke test ===")
|
|
|
|
# CPU-only imports that work on all platforms at build time
|
|
from PIL import Image
|
|
import cv2
|
|
import numpy
|
|
from rembg import new_session
|
|
print(" CPU imports OK (Pillow, cv2, numpy, rembg)")
|
|
|
|
# MediaPipe is CPU-only, should always import
|
|
import mediapipe as mp
|
|
print(" MediaPipe import OK")
|
|
|
|
# LaMa model file must exist
|
|
assert os.path.exists(LAMA_MODEL_PATH), (
|
|
f"LaMa model missing: {LAMA_MODEL_PATH}"
|
|
)
|
|
assert os.path.getsize(LAMA_MODEL_PATH) > LAMA_MIN_SIZE, (
|
|
"LaMa model file is too small"
|
|
)
|
|
print(" LaMa ONNX model file verified")
|
|
|
|
# RealESRGAN model file must exist
|
|
assert os.path.exists(REALESRGAN_MODEL_PATH), (
|
|
f"RealESRGAN model missing: {REALESRGAN_MODEL_PATH}"
|
|
)
|
|
assert os.path.getsize(REALESRGAN_MODEL_PATH) > REALESRGAN_MIN_SIZE, (
|
|
"RealESRGAN model file is too small"
|
|
)
|
|
print(" RealESRGAN model file verified")
|
|
|
|
# GFPGAN model file must exist
|
|
assert os.path.exists(GFPGAN_MODEL_PATH), (
|
|
f"GFPGAN model missing: {GFPGAN_MODEL_PATH}"
|
|
)
|
|
assert os.path.getsize(GFPGAN_MODEL_PATH) > GFPGAN_MIN_SIZE, (
|
|
"GFPGAN model file is too small"
|
|
)
|
|
print(" GFPGAN model file verified")
|
|
|
|
# CodeFormer model file must exist
|
|
assert os.path.exists(CODEFORMER_MODEL_PATH), (
|
|
f"CodeFormer model missing: {CODEFORMER_MODEL_PATH}"
|
|
)
|
|
assert os.path.getsize(CODEFORMER_MODEL_PATH) > CODEFORMER_MIN_SIZE, (
|
|
"CodeFormer model file is too small"
|
|
)
|
|
print(" CodeFormer model file verified")
|
|
|
|
# DDColor ONNX model must exist
|
|
assert os.path.exists(DDCOLOR_ONNX_PATH), (
|
|
f"DDColor model missing: {DDCOLOR_ONNX_PATH}"
|
|
)
|
|
assert os.path.getsize(DDCOLOR_ONNX_PATH) > DDCOLOR_MIN_SIZE, (
|
|
"DDColor model file is too small"
|
|
)
|
|
print(" DDColor ONNX model file verified")
|
|
|
|
# CodeFormer ONNX model must exist
|
|
assert os.path.exists(CODEFORMER_ONNX_PATH), (
|
|
f"CodeFormer model missing: {CODEFORMER_ONNX_PATH}"
|
|
)
|
|
assert os.path.getsize(CODEFORMER_ONNX_PATH) > CODEFORMER_MIN_SIZE, (
|
|
"CodeFormer model file is too small"
|
|
)
|
|
print(" CodeFormer ONNX model file verified")
|
|
|
|
# SCUNet model file must exist
|
|
assert os.path.exists(SCUNET_MODEL_PATH), f"SCUNet model not found: {SCUNET_MODEL_PATH}"
|
|
assert os.path.getsize(SCUNET_MODEL_PATH) > SCUNET_MIN_SIZE
|
|
print(" SCUNet model file verified")
|
|
|
|
# NAFNet model file must exist
|
|
assert os.path.exists(NAFNET_MODEL_PATH), f"NAFNet model not found: {NAFNET_MODEL_PATH}"
|
|
assert os.path.getsize(NAFNET_MODEL_PATH) > NAFNET_MIN_SIZE
|
|
print(" NAFNet model file verified")
|
|
|
|
# PaddleOCR model directories must exist
|
|
for repo_id in PADDLEOCR_MODELS:
|
|
model_name = repo_id.split("/", 1)[1]
|
|
model_dir = os.path.join(PADDLEX_MODEL_DIR, model_name)
|
|
assert os.path.isdir(model_dir), f"PaddleOCR model missing: {model_dir}"
|
|
print(f" PaddleOCR models verified ({len(PADDLEOCR_MODELS)} models)")
|
|
|
|
# PaddleOCR-VL model directory must exist
|
|
vl_name = PADDLEOCR_VL_MODEL.split("/", 1)[1]
|
|
vl_dir = os.path.join(PADDLEX_MODEL_DIR, vl_name)
|
|
assert os.path.isdir(vl_dir), f"PaddleOCR-VL model missing: {vl_dir}"
|
|
print(" PaddleOCR-VL model verified")
|
|
|
|
# MediaPipe task models must exist (for mp.tasks fallback)
|
|
assert os.path.exists(FACE_DETECT_MODEL_PATH), (
|
|
f"MediaPipe face detector model missing: {FACE_DETECT_MODEL_PATH}"
|
|
)
|
|
assert os.path.getsize(FACE_DETECT_MODEL_PATH) > FACE_DETECT_MIN_SIZE
|
|
print(" MediaPipe face detector model verified")
|
|
|
|
assert os.path.exists(FACE_LANDMARKER_MODEL_PATH), (
|
|
f"MediaPipe face landmarker model missing: {FACE_LANDMARKER_MODEL_PATH}"
|
|
)
|
|
assert os.path.getsize(FACE_LANDMARKER_MODEL_PATH) > FACE_LANDMARKER_MIN_SIZE
|
|
print(" MediaPipe face landmarker model verified")
|
|
|
|
# Facexlib auxiliary models must exist (for GFPGAN/CodeFormer)
|
|
assert os.path.exists(FACEXLIB_DET_PATH), (
|
|
f"Facexlib detection model missing: {FACEXLIB_DET_PATH}"
|
|
)
|
|
assert os.path.getsize(FACEXLIB_DET_PATH) > FACEXLIB_DET_MIN_SIZE
|
|
assert os.path.exists(FACEXLIB_PARSE_PATH), (
|
|
f"Facexlib parsing model missing: {FACEXLIB_PARSE_PATH}"
|
|
)
|
|
assert os.path.getsize(FACEXLIB_PARSE_PATH) > FACEXLIB_PARSE_MIN_SIZE
|
|
print(" Facexlib auxiliary models verified")
|
|
|
|
# OpenCV colorization models must exist
|
|
assert os.path.exists(OPENCV_PROTO_PATH), (
|
|
f"OpenCV colorize prototxt missing: {OPENCV_PROTO_PATH}"
|
|
)
|
|
assert os.path.exists(OPENCV_CAFFE_PATH), (
|
|
f"OpenCV colorize caffemodel missing: {OPENCV_CAFFE_PATH}"
|
|
)
|
|
assert os.path.getsize(OPENCV_CAFFE_PATH) > OPENCV_CAFFE_MIN_SIZE
|
|
assert os.path.exists(OPENCV_POINTS_PATH), (
|
|
f"OpenCV colorize points missing: {OPENCV_POINTS_PATH}"
|
|
)
|
|
print(" OpenCV colorization models verified")
|
|
|
|
print("Smoke test passed.\n")
|
|
|
|
|
|
def main():
|
|
print("Pre-downloading all ML models...\n")
|
|
download_lama_model()
|
|
download_rembg_models()
|
|
download_realesrgan_model()
|
|
download_gfpgan_model()
|
|
download_codeformer_model()
|
|
download_ddcolor_model()
|
|
download_codeformer_onnx_model()
|
|
download_paddleocr_models()
|
|
download_paddleocr_vl_model()
|
|
download_scunet_model()
|
|
download_nafnet_model()
|
|
download_facexlib_models()
|
|
download_opencv_colorize_models()
|
|
download_mediapipe_task_models()
|
|
verify_mediapipe()
|
|
smoke_test()
|
|
print("All models downloaded and verified.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|