diff --git a/docs/features.md b/docs/features.md index 565ef9c..6818cc9 100644 --- a/docs/features.md +++ b/docs/features.md @@ -51,7 +51,7 @@ 1. 先运行 `uv run python prepare_vendor.py`，按 `vendor.lock.toml` 下载并校验固定版本的 第三方构建输入；再运行 `uv run python pack.py`，分别构建 `windows-x64`、`macos-arm64`、 - `linux-x64` 三个平台的 Extension ZIP（如 `SnapForm.v0.10.12.macos-arm64.zip`）。 + `linux-x64` 三个平台的 Extension ZIP（如 `SnapForm.v0.11.0.macos-arm64.zip`）。 每个包只含对应平台的 uv 二进制，不再携带 AI 项目的源码包。 2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 3. 在 Extension Preferences 中确认 Storage Root 与 Device。Device 是正式功能 @@ -63,12 +63,14 @@ 固定显示 `Environment`、`ONNX Runtime` 两个用户阶段；FastAPI 等内部安装步骤合并在 `ONNX Runtime` 中。 安装完成后直接写入 Environment Manifest，不再额外导入全部组件进行验证。推荐模型属于安装 后的自动步骤，不计入阶段编号；`DA3 Base` 等其他模型可稍后在 `Models` 面板按需下载。 - 状态文字使用 `阶段/总数 阶段名 | 当前操作 | 预计剩余时间`，依赖安装会结合预设耗时与实际 - 进度修正预期；模型下载根据字节数和速度计算剩余时间，例如 + 状态文字使用 `阶段/总数 阶段名 | 当前操作 | 预计剩余时间`。依赖安装只显示阶段与已解析的 + 包总数，不再承诺剩余时间；模型下载根据字节数和速度计算剩余时间，例如 `Downloading DA3 Small | 14.2 MB/s | About 2 min left`。重试或切换镜像时进度不会倒退。 Environment 安装后即使模型下载中断也不影响 Service 使用，可稍后在 `Models` 面板单独补下。 - 安装不做启动网络预检：pip 安装失败会自动依次重试清华、阿里镜像，模型下载由 worker - 自行探测回退到 HF 镜像；需要手动指定镜像时可通过 `PIP_INDEX_URL` 环境变量覆盖。 + 安装不做启动网络预检：依赖安装默认依次尝试清华、阿里镜像并以 pypi.org 兜底；模型下载由 + worker 自行探测回退到 OmooLab 镜像，切换结果写入 service.json 的 models_from_r2 字段， + 本次周期内后续模型直接走 R2；需要手动指定依赖镜像时可通过 + `UV_INDEX_URL` 环境变量覆盖。 5. 在 3D View 按 `N`，打开 `SnapForm` 页签。 6. 如需其他模型，展开 `Models` panel，找到模型并点击该行的 `Download`。 7. 展开默认收起的 `Debug`，选择 `Output Type` 和本次执行使用的模型。 diff --git a/docs/service.md b/docs/service.md index fadc7b7..d2332c2 100644 --- a/docs/service.md +++ b/docs/service.md @@ -24,7 +24,7 @@ Environment、模型和任务结果默认保存在 `~/.snapform`： 扩展只内置对应平台的 `uv` 可执行文件。`Install Environment` 使用它创建环境并安装： - ONNX Runtime 1.24.4；Windows 使用 `onnxruntime-directml` -- NumPy 2.4.2、Pillow 12.1.1、OpenCV Headless 4.13.0.92、trimesh 4.11.2 +- NumPy 2.4.2、Pillow 12.1.1、PyAV 17.0.0、trimesh 4.11.2 - FastAPI 0.139.2、Uvicorn 0.51.0、huggingface-hub 1.4.1 运行环境不安装 PyTorch、TorchVision、xFormers，也不安装 DA3、MoGe、BEN2 或 utils3d 的 Python 源码包。模型推理和 SnapForm 所需的相机、点云、网格后处理均由 ONNX Runtime、NumPy 与内部实现完成。 @@ -41,14 +41,14 @@ Windows 当前发行环境提供 DirectML 与 CPU；`Auto` 优先使用 `DmlExec - DA3 Small、DA3 Base - BEN2 Base -首次安装自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`，其他模型在 `Models` panel 按需下载。模型按 Hugging Face 仓库末段名称保存在 `~/.snapform/models/`，并保留仓库内部结构。下载顺序为 Hugging Face，然后回退到 `https://models.omoolab.xyz//`；R2 文件会校验大小与 SHA-256。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 +首次安装自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`，其他模型在 `Models` panel 按需下载。模型按 Hugging Face 仓库末段名称保存在 `~/.snapform/models/`，并保留仓库内部结构。声明了文件清单与校验值的模型会直接走 Hugging Face `resolve` 直链流式下载，跳过文件清单与 ETag 预校验；连接源不可达时 6 秒内快速失败并回退到 `https://models.omoolab.xyz//`，并把 `models_from_r2` 写入 service.json，本次 service 周期内后续模型直接走 R2，不再重复等待 Hugging Face。R2 文件会校验大小与 SHA-256。未声明清单的仓库仍使用 `huggingface_hub` 快照下载。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 DA3 ONNX 使用模型的 camera decoder 恢复外参与内参；旧版 `Use Ray Pose` 参数只为任务协议兼容保留，不参与 ONNX 推理。MoGe-2 的 metric depth、normal、mask、normalized intrinsics，以及 BEN2 的 Alpha，均在 NumPy 后处理中恢复。 ## 安装网络与镜像 -安装器直接开始工作，不做阻塞式网络预检。依赖安装失败时会依次重试清华与阿里云 PyPI 镜像，也可用 `PIP_INDEX_URL` 或 Preferences 中的 `PyPI Index` 显式指定。Windows 系统代理会注入子进程；同时支持 `HTTP_PROXY`、`HTTPS_PROXY` 与 `ALL_PROXY`。 +安装器直接开始工作，不做阻塞式网络预检。依赖安装默认依次尝试清华、阿里云 PyPI 镜像并以 pypi.org 兜底，也可用 `UV_INDEX_URL` 显式指定。Windows 系统代理会注入子进程；同时支持 `HTTP_PROXY`、`HTTPS_PROXY` 与 `ALL_PROXY`。 GitHub 不参与用户安装流程。`prepare_vendor.py` 只为扩展打包准备各平台的 `uv` 二进制。 diff --git a/pyproject.toml b/pyproject.toml index 82eb642..f8e2df3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "snapform" -version = "0.10.12" +version = "0.11.0" description = "Fast image-guided structure building tools for Blender" readme = "README.md" requires-python = ">=3.11" diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index 081f2d4..4f80b2b 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -1,7 +1,7 @@ schema_version = "1.0.0" id = "snapform" -version = "0.10.12" +version = "0.11.0" name = "SnapForm" tagline = "Build spatial structures from images, fast" maintainer = "OmooLab" diff --git a/src/snapform/client.py b/src/snapform/client.py index 8a6f30c..a8724bb 100644 --- a/src/snapform/client.py +++ b/src/snapform/client.py @@ -18,6 +18,7 @@ from urllib.request import ProxyHandler, Request, build_opener import bpy +from .env_revision import MANIFEST_FILENAME, MANIFEST_REVISION from .models import ( BEN2_MODEL, get_downloadable_model, @@ -26,8 +27,6 @@ from .models import ( ) -MANIFEST_FILENAME = "manifest.json" -MANIFEST_REVISION = 7 SERVICE_ADDRESS = "127.0.0.1" DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" JOB_FOLDER_NAMES = { diff --git a/src/snapform/env_revision.py b/src/snapform/env_revision.py new file mode 100644 index 0000000..d0a6201 --- /dev/null +++ b/src/snapform/env_revision.py @@ -0,0 +1,4 @@ +"""Shared environment metadata for the Blender UI and the install subprocess.""" + +MANIFEST_FILENAME = "manifest.json" +MANIFEST_REVISION = 8 diff --git a/src/snapform/install_env.py b/src/snapform/install_env.py index 1a6cac6..ba6a70c 100644 --- a/src/snapform/install_env.py +++ b/src/snapform/install_env.py @@ -19,6 +19,7 @@ SCRIPT_DIRECTORY = Path(__file__).resolve().parent if str(SCRIPT_DIRECTORY) not in sys.path: sys.path.insert(0, str(SCRIPT_DIRECTORY)) +from env_revision import MANIFEST_REVISION from models import BEN2_MODEL, DOWNLOADABLE_MODELS, MOGE2_MODEL @@ -32,10 +33,9 @@ UVICORN_VERSION = "0.51.0" ONNXRUNTIME_VERSION = "1.24.4" NUMPY_VERSION = "2.4.2" PILLOW_VERSION = "12.1.1" -OPENCV_VERSION = "4.13.0.92" +AV_VERSION = "17.0.0" TRIMESH_VERSION = "4.11.2" HUGGINGFACE_HUB_VERSION = "1.4.1" -MANIFEST_REVISION = 7 _resolved_pypi_index = None DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key) RUNTIME_STAGE_COUNT = 2 @@ -45,15 +45,9 @@ MODELS_SLICE_START = 0.85 MODELS_SLICE_END = 1.0 RAMP_RATE = 0.02 CREEP_CAP = 0.60 -INSTALL_ESTIMATES = { - "ONNX Runtime": 3 * 60, -} -INSTALL_ESTIMATE_RANGES = { - "ONNX Runtime": AI_COMPONENTS_SLICE, -} CHECK_CODE = ( "import json; import fastapi, uvicorn, onnxruntime; " - "import cv2, numpy, PIL, trimesh; " + "import av, numpy, PIL, trimesh; " "print('SNAPFORM_DEVICES=' + json.dumps({" "'onnxruntime': onnxruntime.__version__, " "'providers': onnxruntime.get_available_providers()}))" @@ -76,7 +70,6 @@ def write_status( stage=None, stages=None, stage_label=None, - eta_seconds=None, ): payload = {"progress": progress, "message": message} if error: @@ -85,8 +78,6 @@ def write_status( payload["stage"] = stage payload["stages"] = stages payload["stage_label"] = stage_label - if eta_seconds is not None: - payload["eta_seconds"] = max(int(eta_seconds), 0) Path(path).write_text( json.dumps(payload, ensure_ascii=False), encoding="utf-8", @@ -103,7 +94,7 @@ def child_environment(environment=None): if env_key not in process_environment: process_environment[env_key] = value if _resolved_pypi_index: - process_environment["PIP_INDEX_URL"] = _resolved_pypi_index + process_environment["UV_INDEX_URL"] = _resolved_pypi_index if environment: process_environment.update(environment) return process_environment @@ -117,13 +108,6 @@ def run(command, environment=None): ) -def format_elapsed(seconds): - seconds = int(seconds) - if seconds < 60: - return f"{seconds}s" - return f"{seconds // 60}m {seconds % 60:02d}s" - - def parse_package_count(line): try: return int(line.split()[1]) @@ -131,26 +115,12 @@ def parse_package_count(line): return 0 -def estimated_remaining(label, progress, elapsed=0.0): - expected_seconds = INSTALL_ESTIMATES.get(label) - progress_range = INSTALL_ESTIMATE_RANGES.get(label) - if expected_seconds is None or progress_range is None: - return None - start, end = progress_range - if end <= start: - return 0.0 - remaining_fraction = 1.0 - min( - max((progress - start) / (end - start), 0.0), - 1.0, - ) - return max(expected_seconds * remaining_fraction - elapsed, 0.0) - - def resolve_indexes(args): - """???????????????????????????????????""" + """确定用户指定的 PyPI 源，优先命令行参数，其次 uv/pip 的环境变量。""" global _resolved_pypi_index _resolved_pypi_index = ( getattr(args, "pypi_index", None) + or os.environ.get("UV_INDEX_URL") or os.environ.get("PIP_INDEX_URL") ) @@ -206,11 +176,6 @@ def _run_install_once( stage=stage, stages=stages, stage_label=label, - eta_seconds=estimated_remaining( - label, - displayed, - clock() - started, - ), ) while True: @@ -233,7 +198,21 @@ def _run_install_once( if empty: if elapsed - last_heartbeat >= 5.0: last_heartbeat = elapsed - detail("Downloading packages") + # uv 的下载阶段没有增量进度，只报告已解析的包总数， + # 让用户对规模有预期，而不承诺剩余时间或已耗时。 + message = ( + f"Downloading {total} packages" + if total + else "Downloading packages" + ) + write_status( + status_path, + displayed, + message, + stage=stage, + stages=stages, + stage_label=label, + ) continue line = raw_line.rstrip() print(line) @@ -273,13 +252,20 @@ def run_install( end=1.0, clock=time.monotonic, ): - """Run a uv install command, retrying with mirrors when the default index fails.""" - attempts = [environment] - if _resolved_pypi_index in (None, PYPI_DEFAULT_INDEX): - attempts.extend( - {"PIP_INDEX_URL": mirror} for mirror in PYPI_MIRRORS - ) - for index, attempt_environment in enumerate(attempts): + """Run a uv install command, retrying across PyPI indexes on failure.""" + if _resolved_pypi_index is None: + # 未指定源时优先国内镜像，pypi.org 作为最终兜底。 + indexes = (*PYPI_MIRRORS, PYPI_DEFAULT_INDEX) + elif _resolved_pypi_index == PYPI_DEFAULT_INDEX: + # 明确选择官方源时，失败后回退国内镜像。 + indexes = (PYPI_DEFAULT_INDEX, *PYPI_MIRRORS) + else: + # 用户指定了特定源，不再回退。 + indexes = (_resolved_pypi_index,) + attempts = [ + dict(environment or {}, UV_INDEX_URL=index) for index in indexes + ] + for position, attempt_environment in enumerate(attempts): try: return _run_install_once( status_path, @@ -293,9 +279,12 @@ def run_install( clock, ) except subprocess.CalledProcessError: - if index + 1 == len(attempts): + if position + 1 == len(attempts): raise - print(f"Retrying with PyPI mirror: {PYPI_MIRRORS[0]}") + print( + "Retrying with next PyPI index: " + f"{attempts[position + 1]['UV_INDEX_URL']}" + ) def venv_python(venv_dir): @@ -330,7 +319,7 @@ def install_stages(args, venv_dir, python, uv): f"fastapi=={FASTAPI_VERSION}", f"huggingface-hub=={HUGGINGFACE_HUB_VERSION}", f"numpy=={NUMPY_VERSION}", - f"opencv-python-headless=={OPENCV_VERSION}", + f"av=={AV_VERSION}", runtime_package, f"pillow=={PILLOW_VERSION}", f"trimesh=={TRIMESH_VERSION}", @@ -465,7 +454,6 @@ def run_verify( stage=stage, stages=stages, stage_label=stage_label, - eta_seconds=max(60.0 - elapsed, 0.0), ) continue line = raw_line.rstrip() @@ -596,7 +584,6 @@ def run_stage( stage=stage, stages=stages, stage_label=label, - eta_seconds=estimated_remaining(label, start), ) installed = [] if kind == "install": diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 895488d..70e3131 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -1,15 +1,23 @@ import argparse import hashlib +import http.client import json import os import shutil +import socket +import sys import tempfile import threading import time -import sys from pathlib import Path from urllib.parse import quote -from urllib.request import Request, urlopen +from urllib.request import ( + HTTPSHandler, + ProxyHandler, + Request, + build_opener, + urlopen, +) SCRIPT_DIRECTORY = Path(__file__).resolve().parent @@ -40,7 +48,29 @@ EXPORT_FORMATS = ( ) HF_ENDPOINT = "https://huggingface.co" R2_ENDPOINT = "https://models.omoolab.xyz" -DOWNLOAD_TIMEOUT = 60.0 +# 连接（TCP + TLS 握手）超时，源不可达时快速失败切换；下载阶段不设读超时。 +CONNECT_TIMEOUT = 6.0 + + +class _ShortConnectHTTPSConnection(http.client.HTTPSConnection): + """HTTPS 连接：TCP + TLS 握手用短超时，读操作不设超时。""" + + def connect(self): + try: + self.timeout = CONNECT_TIMEOUT + super().connect() + finally: + self.timeout = None + # 连接建立后读操作不设超时，下载大文件不会被读超时中断。 + self.sock.settimeout(None) + + +class _ShortConnectHTTPSHandler(HTTPSHandler): + def https_open(self, req): + return self.do_open(_ShortConnectHTTPSConnection, req) + + +_download_opener = build_opener(ProxyHandler(), _ShortConnectHTTPSHandler()) class TaskCancelled(RuntimeError): @@ -107,34 +137,32 @@ def extract_video_frames( video_path, output_directory, maximum_frames, - cv2, + av, start_frame=1, cancel_check=None, ): output_directory = Path(output_directory) output_directory.mkdir(parents=True, exist_ok=True) - capture = cv2.VideoCapture(str(video_path)) - if not capture.isOpened(): - raise RuntimeError(f"Unable to open video: {video_path}") + try: + container = av.open(str(video_path)) + except Exception as error: + raise RuntimeError(f"Unable to open video: {video_path}") from error image_paths = [] skipped = 0 try: - while len(image_paths) < maximum_frames: + stream = container.streams.video[0] + for frame in container.decode(stream): check_cancelled(cancel_check) - success, frame = capture.read() - if not success: - break if skipped < start_frame - 1: skipped += 1 continue - image_path = ( - output_directory / f"frame.{len(image_paths) + 1:04d}.jpg" - ) - if not cv2.imwrite(str(image_path), frame): - raise RuntimeError(f"Unable to write video frame: {image_path}") + image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg" + frame.to_image().save(image_path, format="JPEG") image_paths.append(image_path) + if len(image_paths) >= maximum_frames: + break finally: - capture.release() + container.close() if not image_paths: raise RuntimeError(f"Unable to read any video frame from: {video_path}") @@ -162,7 +190,7 @@ def prepare_input( supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES)) raise RuntimeError(f"Unsupported input type; expected one of: {supported}") - import cv2 + import av temporary_directory = Path( tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root) @@ -172,7 +200,7 @@ def prepare_input( input_path, temporary_directory, video_frames, - cv2, + av, start_frame=start_frame, cancel_check=cancel_check, ) @@ -254,33 +282,6 @@ def unproject_camera_depth(depth, intrinsics): return (rays * depth.reshape(-1, 1)).reshape(height, width, 3) -def smooth_depth_map(depth, smoothing): - import numpy as np - - factor = min(max(float(smoothing) / 100.0, 0.0), 1.0) - if factor <= 0.0: - return np.asarray(depth, dtype=np.float32).copy() - import cv2 - - values = np.asarray(depth, dtype=np.float32) - valid = np.isfinite(values) & (values > 0.0) - if not valid.any(): - raise ValueError("DA3 returned no valid depth values") - low, high = np.percentile(values[valid], (5.0, 95.0)) - depth_range = max(float(high - low), 1e-6) - diameter = 3 + 2 * int(round(3.0 * factor)) - working_values = values.copy() - working_values[~valid] = float(np.median(values[valid])) - smoothed = cv2.bilateralFilter( - working_values, - diameter, - depth_range * (0.01 + 0.09 * factor), - 1.0 + 4.0 * factor, - ) - smoothed[~valid] = values[~valid] - return smoothed - - def compact_mesh(vertices, faces): import numpy as np @@ -961,8 +962,10 @@ def create_download_progress_class( class DownloadProgress(tqdm_type): def __init__(self, *args, **kwargs): description = str(kwargs.get("desc", "")) + # huggingface_hub 1.x 的字节进度条描述是 "Downloading (incomplete total...)"， + # 用前缀匹配而非精确匹配，否则真实下载进度会被忽略。 self.snapform_reports_progress = ( - description == "Downloading bytes" + description.startswith("Downloading") or description.startswith("Fetching ") ) super().__init__(*args, **kwargs) @@ -998,6 +1001,35 @@ def progress_range(args, staged): return 0.05, 0.90 +def _discovery_path(model_dir): + """storage_root/service.json，从 storage_root/models/ 反推 storage_root。""" + return Path(model_dir).parent.parent / "service.json" + + +def _models_from_r2(model_dir): + try: + payload = json.loads( + _discovery_path(model_dir).read_text(encoding="utf-8") + ) + except (OSError, ValueError, json.JSONDecodeError): + return False + return bool(payload.get("models_from_r2")) + + +def _mark_models_from_r2(model_dir): + """在 service.json 上记录 models_from_r2，保留文件里已有的其它字段。""" + path = _discovery_path(model_dir) + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError, json.JSONDecodeError): + payload = {} + payload["models_from_r2"] = True + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f"{path.name}.tmp") + temporary.write_text(json.dumps(payload), encoding="utf-8") + temporary.replace(path) + + def download_model(args): label = getattr(args, "label", None) or args.repository model_dir = Path(args.model_dir) @@ -1009,6 +1041,77 @@ def download_model(args): "stages": getattr(args, "stages", None), "stage_label": getattr(args, "stage_label", None), } + r2_directory = getattr(args, "r2_directory", "") + r2_files = getattr(args, "r2_file", []) + if r2_directory and r2_files: + _download_declared_or_fail( + args, + model_dir, + label, + start, + span, + stage_kwargs, + ) + return + _download_snapshot_or_fail(args, model_dir, label, start, span, stage_kwargs) + + +def _download_declared_or_fail(args, model_dir, label, start, span, stage_kwargs): + """Stream declared files from Hugging Face, then the R2 mirror, remembering fallback.""" + if _models_from_r2(model_dir): + sources = ( + ( + R2_ENDPOINT, + download_from_r2, + f"Downloading {label} | OmooLab mirror", + ), + ) + else: + sources = ( + (HF_ENDPOINT, download_from_hf, f"Downloading {label}"), + ( + R2_ENDPOINT, + download_from_r2, + f"Downloading {label} | OmooLab mirror", + ), + ) + failures = [] + for position, (source, downloader, message) in enumerate(sources): + write_status(args.status, start, message, **stage_kwargs) + try: + downloader( + args, + model_dir, + start=start, + span=span, + label=label, + stage_kwargs=stage_kwargs, + ) + _verify_model_complete( + model_dir, + getattr(args, "required_pattern", None), + ) + except Exception as error: + failures.append((source, error)) + continue + # 非首个源成功意味着 hf 已失败，记住本次 service 周期后续直接走 r2。 + if position > 0: + _mark_models_from_r2(model_dir) + write_status( + args.status, + start + span, + f"{label} ready", + **stage_kwargs, + ) + return + detail = "; ".join( + f"{source}: {error}" for source, error in failures + ) + raise RuntimeError(f"All download sources failed: {detail}") + + +def _download_snapshot_or_fail(args, model_dir, label, start, span, stage_kwargs): + """Fallback download for repositories without declared files, via huggingface_hub.""" failures = [] try: from huggingface_hub import snapshot_download @@ -1027,7 +1130,9 @@ def download_model(args): ) for endpoint in download_endpoints() if snapshot_download else (): os.environ["HF_ENDPOINT"] = endpoint - message = f"Downloading {label}" + # snapshot_download 先拉取文件清单并校验 ETag，还没有字节下载， + # 用 Checking 提示当前阶段，避免进度条看起来卡住。 + message = f"Checking {label} files" if endpoint != HF_ENDPOINT: message = f"{message} (mirror)" write_status(args.status, start, message, **stage_kwargs) @@ -1053,43 +1158,6 @@ def download_model(args): **stage_kwargs, ) return - # Hugging Face 失败后使用由项目维护的静态模型镜像。 - r2_directory = getattr(args, "r2_directory", "") - r2_files = getattr(args, "r2_file", []) - if not r2_directory or not r2_files: - detail = "; ".join( - f"{source}: {error}" for source, error in failures - ) - raise RuntimeError(f"All download sources failed: {detail}") - try: - write_status( - args.status, - start, - f"Downloading {label} | OmooLab mirror", - **stage_kwargs, - ) - download_from_r2( - args, - model_dir, - start=start, - span=span, - label=label, - stage_kwargs=stage_kwargs, - ) - _verify_model_complete( - model_dir, - getattr(args, "required_pattern", None), - ) - except Exception as error: - failures.append((R2_ENDPOINT, error)) - else: - write_status( - args.status, - start + span, - f"{label} ready", - **stage_kwargs, - ) - return detail = "; ".join( f"{source}: {error}" for source, error in failures ) @@ -1121,25 +1189,24 @@ def _verify_model_complete(model_dir, required_patterns=None): ) -def _download_r2_file( - directory, - filename, +def _download_file_from_url( + url, destination, expected_size, expected_checksum, progress=None, retries=2, ): - """Download and verify one OmooLab R2 model file atomically.""" - url = f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}" + """Download and verify one model file atomically from a direct URL.""" request = Request(url, headers={"User-Agent": "SnapForm"}) + opener = _download_opener.open if request.type == "https" else urlopen last_error = None for _attempt in range(retries): temporary = destination.with_name(f"{destination.name}.part") try: checksum = hashlib.sha256() downloaded = 0 - with urlopen(request, timeout=DOWNLOAD_TIMEOUT) as response: + with opener(request) as response: with open(temporary, "wb") as output: while True: chunk = response.read(1024 * 1024) @@ -1152,11 +1219,11 @@ def _download_r2_file( progress(len(chunk)) if downloaded != expected_size: raise RuntimeError( - f"Size mismatch for {filename}: " + f"Size mismatch for {destination.name}: " f"{downloaded} != {expected_size}" ) if checksum.hexdigest().lower() != expected_checksum.lower(): - raise RuntimeError(f"SHA-256 mismatch for {filename}") + raise RuntimeError(f"SHA-256 mismatch for {destination.name}") temporary.replace(destination) return except Exception as error: @@ -1168,20 +1235,37 @@ def _download_r2_file( raise last_error -def download_from_r2( - args, - model_dir, - start=0.05, - span=0.90, - label=None, - stage_kwargs=None, -): - """Download the declared model files from OmooLab R2.""" - repository = args.repository +def r2_file_url(directory, filename): + return f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}" + + +def hf_file_url(repository, filename): + """Hugging Face resolve URL for a direct, metadata-free download.""" + endpoint = (os.environ.get("HF_ENDPOINT") or HF_ENDPOINT).rstrip("/") + return f"{endpoint}/{repository}/resolve/main/{quote(filename)}" + + +def _parse_declared_files(args): selected = [] for value in args.r2_file: filename, size, checksum = value.split("|", 2) selected.append((filename, int(size), checksum)) + return selected + + +def _download_declared_files( + args, + model_dir, + url_builder, + start, + span, + label, + stage_kwargs, + retries=2, +): + """Download each declared file through url_builder with aggregated progress.""" + repository = args.repository + selected = _parse_declared_files(args) reporter = DownloadProgressReporter( args.status, repository, @@ -1201,16 +1285,57 @@ def download_from_r2( for path, size, checksum in selected: destination = Path(model_dir) / path destination.parent.mkdir(parents=True, exist_ok=True) - _download_r2_file( - args.r2_directory, - path, + _download_file_from_url( + url_builder(path), destination, size, checksum, progress=on_chunk, + retries=retries, ) +def download_from_r2( + args, + model_dir, + start=0.05, + span=0.90, + label=None, + stage_kwargs=None, +): + """Download the declared model files from OmooLab R2.""" + _download_declared_files( + args, + model_dir, + lambda filename: r2_file_url(args.r2_directory, filename), + start, + span, + label, + stage_kwargs, + ) + + +def download_from_hf( + args, + model_dir, + start=0.05, + span=0.90, + label=None, + stage_kwargs=None, +): + """Download the declared model files directly from Hugging Face resolve URLs.""" + _download_declared_files( + args, + model_dir, + lambda filename: hf_file_url(args.repository, filename), + start, + span, + label, + stage_kwargs, + retries=1, + ) + + def format_rate(rate): size = float(rate) if size <= 0: diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index 9d69a90..78d3f0a 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -60,6 +60,29 @@ class DownloadProgressTest(unittest.TestCase): self.assertAlmostEqual(status["progress"], 0.275) self.assertEqual(status["message"], "Downloading owner/model") + def test_reports_byte_progress_for_aggregated_bar_description(self): + # huggingface_hub 1.x 用 "Downloading (incomplete total...)" 作为字节进度条描述 + with tempfile.TemporaryDirectory() as directory: + status_path = Path(directory) / "status.json" + progress_class = self.worker.create_download_progress_class( + status_path, + "owner/model", + tqdm_type=FakeTqdm, + clock=lambda: 1.0, + ) + progress = progress_class( + desc="Downloading (incomplete total...)", + total=100, + initial=0, + unit="B", + ) + + progress.update(25) + + status = json.loads(status_path.read_text(encoding="utf-8")) + self.assertAlmostEqual(status["progress"], 0.275) + self.assertEqual(status["message"], "Downloading owner/model") + def test_reports_download_speed(self): class FakeTqdmWithRate(FakeTqdm): @property diff --git a/tests/test_install_env.py b/tests/test_install_env.py index b56b993..730d09a 100644 --- a/tests/test_install_env.py +++ b/tests/test_install_env.py @@ -521,7 +521,59 @@ class InstallEnvTest(unittest.TestCase): self.assertEqual(first.kwargs["stages"], 8) self.assertEqual(first.kwargs["stage_label"], "pytorch deps") - def test_run_install_retries_once_with_mirror(self): + def test_run_install_heartbeat_reports_package_count(self): + bootstrap = load_install_env() + + class SlowLines: + def __iter__(self): + yield "Resolved 3 packages in 1.0s\n" + time.sleep(0.6) + + class FakeProcess: + def __init__(self, lines): + self.stdout = lines + + def wait(self): + return 0 + + clock_state = {"now": 0.0} + + def fake_clock(): + clock_state["now"] += 6.0 + return clock_state["now"] + + fake_process = FakeProcess(SlowLines()) + with patch.object( + bootstrap.subprocess, + "Popen", + return_value=fake_process, + ), patch.object(bootstrap, "write_status") as mocked_write: + bootstrap._run_install_once( + "status.json", + 2, + 2, + "ONNX Runtime", + ["uv", "pip", "install", "onnxruntime"], + None, + 0.10, + 0.85, + fake_clock, + ) + + heartbeats = [ + call + for call in mocked_write.call_args_list + if call.args[2].startswith("Downloading") + ] + self.assertTrue(heartbeats) + self.assertTrue( + any("3 packages" in call.args[2] for call in heartbeats) + ) + for call in heartbeats: + self.assertNotIn("elapsed", call.args[2]) + self.assertNotIn("eta_seconds", call.kwargs) + + def test_run_install_retries_across_indexes(self): bootstrap = load_install_env() previous = bootstrap._resolved_pypi_index bootstrap._resolved_pypi_index = None @@ -546,32 +598,44 @@ class InstallEnvTest(unittest.TestCase): self.assertEqual(installed, ["numpy==1.26.4"]) self.assertEqual(mocked_install.call_count, 2) + # 未指定源时默认国内镜像优先，第一个镜像失败后切到第二个镜像。 + self.assertEqual( + mocked_install.call_args_list[0].args[5], + {"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[0]}, + ) self.assertEqual( mocked_install.call_args_list[1].args[5], - {"PIP_INDEX_URL": bootstrap.PYPI_MIRRORS[0]}, + {"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[1]}, ) - def test_formats_elapsed_time(self): + def test_run_install_uses_only_configured_index(self): bootstrap = load_install_env() - - self.assertEqual(bootstrap.format_elapsed(12), "12s") - self.assertEqual(bootstrap.format_elapsed(83), "1m 23s") - - def test_onnx_runtime_estimate_decreases_with_progress(self): - bootstrap = load_install_env() - - first = bootstrap.estimated_remaining( - "ONNX Runtime", - bootstrap.AI_COMPONENTS_SLICE[0], - ) - second = bootstrap.estimated_remaining( - "ONNX Runtime", - sum(bootstrap.AI_COMPONENTS_SLICE) / 2, + previous = bootstrap._resolved_pypi_index + bootstrap._resolved_pypi_index = ( + "https://pypi.tuna.tsinghua.edu.cn/simple" ) + try: + with patch.object( + bootstrap, + "_run_install_once", + return_value=["numpy==1.26.4"], + ) as mocked_install: + installed = bootstrap.run_install( + "status.json", + 3, + 8, + "ai deps", + ["uv", "pip", "install", "x"], + ) + finally: + bootstrap._resolved_pypi_index = previous - self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["ONNX Runtime"]) - self.assertGreater(second, 0) - self.assertLess(second, first) + self.assertEqual(installed, ["numpy==1.26.4"]) + self.assertEqual(mocked_install.call_count, 1) + self.assertEqual( + mocked_install.call_args.args[5], + {"UV_INDEX_URL": "https://pypi.tuna.tsinghua.edu.cn/simple"}, + ) def test_resolve_indexes_uses_configured_pypi_index(self): bootstrap = load_install_env() @@ -584,6 +648,18 @@ class InstallEnvTest(unittest.TestCase): ) def test_resolve_indexes_uses_environment_override(self): + bootstrap = load_install_env() + with patch.dict( + os.environ, + {"UV_INDEX_URL": "https://pypi.example/simple"}, + ): + bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None)) + self.assertEqual( + bootstrap._resolved_pypi_index, + "https://pypi.example/simple", + ) + + def test_resolve_indexes_falls_back_to_pip_index_url(self): bootstrap = load_install_env() with patch.dict( os.environ, @@ -607,7 +683,7 @@ class InstallEnvTest(unittest.TestCase): bootstrap._resolved_pypi_index = previous_pypi self.assertEqual( - env["PIP_INDEX_URL"], + env["UV_INDEX_URL"], "https://pypi.tuna.tsinghua.edu.cn/simple", ) self.assertEqual( diff --git a/tests/test_worker.py b/tests/test_worker.py index 6895293..fd5198b 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -1,6 +1,7 @@ import importlib.util import json import os +import shutil import sys import tempfile import unittest @@ -25,6 +26,49 @@ def load_worker(): return module +def _make_fake_av(frame_count): + """Build a fake av module yielding frame_count integer-valued frames.""" + class FakeImage: + def __init__(self, value): + self.value = value + + def save(self, path, format=None): + Path(path).write_text(str(self.value), encoding="utf-8") + + class FakeFrame: + def __init__(self, value): + self.value = value + + def to_image(self): + return FakeImage(self.value) + + class FakeContainer: + def __init__(self): + self.position = 0 + self.closed = False + self.streams = SimpleNamespace(video=[object()]) + + def decode(self, _stream): + while self.position < frame_count: + frame = FakeFrame(self.position) + self.position += 1 + yield frame + + def close(self): + self.closed = True + + class FakeAv: + def __init__(self): + self.containers = [] + + def open(self, _path): + container = FakeContainer() + self.containers.append(container) + return container + + return FakeAv() + + class WorkerTest(unittest.TestCase): @classmethod def setUpClass(cls): @@ -177,92 +221,32 @@ class WorkerTest(unittest.TestCase): self.assertIsNone(temporary_directory) def test_extract_video_frames_reads_frames_sequentially(self): - class FakeCapture: - def __init__(self): - self.position = 0 - - def isOpened(self): - return True - - def read(self): - if self.position >= 100: - return False, None - frame = self.position - self.position += 1 - return True, frame - - def release(self): - return None - - class FakeCv2: - def __init__(self): - self.captures = [] - - def VideoCapture(self, _path): - capture = FakeCapture() - self.captures.append(capture) - return capture - - def imwrite(self, path, frame): - Path(path).write_text(str(frame), encoding="utf-8") - return True - with tempfile.TemporaryDirectory() as directory: root = Path(directory) - fake_cv2 = FakeCv2() + fake_av = _make_fake_av(100) image_paths = self.worker.extract_video_frames( root / "clip.mkv", root / "out", 10, - fake_cv2, + fake_av, ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) - self.assertEqual(len(fake_cv2.captures), 1) - self.assertEqual(fake_cv2.captures[0].position, 10) + self.assertEqual(len(fake_av.containers), 1) + self.assertEqual(fake_av.containers[0].position, 10) def test_extract_video_frames_starts_after_offset(self): - class FakeCapture: - def __init__(self): - self.position = 0 - - def isOpened(self): - return True - - def read(self): - if self.position >= 100: - return False, None - frame = self.position - self.position += 1 - return True, frame - - def release(self): - return None - - class FakeCv2: - def __init__(self): - self.captures = [] - - def VideoCapture(self, _path): - capture = FakeCapture() - self.captures.append(capture) - return capture - - def imwrite(self, path, frame): - Path(path).write_text(str(frame), encoding="utf-8") - return True - with tempfile.TemporaryDirectory() as directory: root = Path(directory) - fake_cv2 = FakeCv2() + fake_av = _make_fake_av(100) image_paths = self.worker.extract_video_frames( root / "clip.mkv", root / "out", 10, - fake_cv2, + fake_av, start_frame=71, ) extracted = [ @@ -271,41 +255,16 @@ class WorkerTest(unittest.TestCase): ] self.assertEqual(extracted, [str(frame) for frame in range(70, 80)]) - self.assertEqual(fake_cv2.captures[0].position, 80) + self.assertEqual(fake_av.containers[0].position, 80) def test_extract_video_frames_tolerates_short_video_tail(self): - class ShortCapture: - def __init__(self): - self.position = 0 - - def isOpened(self): - return True - - def read(self): - if self.position >= 99: - return False, None - frame = self.position - self.position += 1 - return True, frame - - def release(self): - return None - - class FakeCv2: - def VideoCapture(self, _path): - return ShortCapture() - - def imwrite(self, path, frame): - Path(path).write_text(str(frame), encoding="utf-8") - return True - with tempfile.TemporaryDirectory() as directory: root = Path(directory) image_paths = self.worker.extract_video_frames( root / "clip.mkv", root / "out", 10, - FakeCv2(), + _make_fake_av(99), ) self.assertEqual( @@ -443,38 +402,120 @@ class WorkerTest(unittest.TestCase): self.worker.download_model(args) def test_download_model_reports_all_failed_sources(self): - fake_hugging_face = ModuleType("huggingface_hub") - - def fake_snapshot_download(**options): - raise OSError("direct down") - - fake_hugging_face.snapshot_download = fake_snapshot_download args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", status="status.json", allow_pattern=[], + r2_directory="DA3Small", + r2_file=["config.json|100|abc"], ) - with tempfile.TemporaryDirectory() as directory, patch.dict( - sys.modules, - {"huggingface_hub": fake_hugging_face}, - ), patch.object( + with tempfile.TemporaryDirectory() as directory, patch.object( self.worker, - "create_download_progress_class", - return_value=object, - ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( + "download_from_hf", + side_effect=OSError("direct down"), + ), patch.object( self.worker, "download_from_r2", side_effect=OSError("mirror down"), ): - args.r2_directory = "DA3Small" - args.r2_file = ["config.json|100|abc"] args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"): self.worker.download_model(args) + def test_download_model_prefers_direct_hf_then_r2(self): + args = SimpleNamespace( + repository="depth-anything/DA3-SMALL", + model_dir="model", + status="status.json", + allow_pattern=["config.json"], + r2_directory="DA3Small", + r2_file=["config.json|100|abc"], + ) + order = [] + + def fake_download_from_hf(*_args, **_kwargs): + order.append("hf") + raise OSError("hf down") + + def fake_download_from_r2(*_args, **_kwargs): + order.append("r2") + + with tempfile.TemporaryDirectory() as directory, patch.object( + self.worker, + "download_from_hf", + side_effect=fake_download_from_hf, + ), patch.object( + self.worker, + "download_from_r2", + side_effect=fake_download_from_r2, + ), patch.object( + self.worker, + "_verify_model_complete", + return_value=None, + ): + args.model_dir = str(Path(directory) / "model") + args.status = str(Path(directory) / "status.json") + self.worker.download_model(args) + + self.assertEqual(order, ["hf", "r2"]) + + def test_download_declared_skips_hf_when_r2_preferred(self): + root = PROJECT_ROOT / ".tmp-models-from-r2" + model_dir = root / "models" / "DA3" + args = SimpleNamespace( + repository="depth-anything/DA3-SMALL", + model_dir=str(model_dir), + status="status.json", + allow_pattern=[], + r2_directory="DA3Small", + r2_file=["config.json|100|abc"], + ) + order = [] + try: + self.worker._mark_models_from_r2(model_dir) + with patch.object( + self.worker, + "download_from_hf", + side_effect=lambda *a, **k: order.append("hf"), + ), patch.object( + self.worker, + "download_from_r2", + side_effect=lambda *a, **k: order.append("r2"), + ), patch.object( + self.worker, + "_verify_model_complete", + return_value=None, + ), patch.object(self.worker, "write_status"): + self.worker._download_declared_or_fail( + args, model_dir, "DA3 Small", 0.0, 1.0, {} + ) + finally: + shutil.rmtree(root, ignore_errors=True) + + self.assertEqual(order, ["r2"]) + + def test_models_from_r2_roundtrip_preserves_fields(self): + root = PROJECT_ROOT / ".tmp-models-from-r2" + model_dir = root / "models" / "DA3" + discovery = root / "service.json" + try: + discovery.parent.mkdir(parents=True, exist_ok=True) + discovery.write_text( + json.dumps({"port": 1, "token": "t"}), + encoding="utf-8", + ) + self.assertFalse(self.worker._models_from_r2(model_dir)) + self.worker._mark_models_from_r2(model_dir) + self.assertTrue(self.worker._models_from_r2(model_dir)) + payload = json.loads(discovery.read_text(encoding="utf-8")) + self.assertEqual(payload["port"], 1) + self.assertTrue(payload["models_from_r2"]) + finally: + shutil.rmtree(root, ignore_errors=True) + def test_download_from_r2_downloads_declared_files(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", @@ -488,8 +529,8 @@ class WorkerTest(unittest.TestCase): with tempfile.TemporaryDirectory() as directory, patch.object( self.worker, - "_download_r2_file", - side_effect=lambda *args, **kwargs: downloaded.append(args[1]), + "_download_file_from_url", + side_effect=lambda *args, **kwargs: downloaded.append(args[1].name), ), patch.object(self.worker, "write_status"): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") diff --git a/uv.lock b/uv.lock index a3d6970..3397aa2 100644 --- a/uv.lock +++ b/uv.lock @@ -205,7 +205,7 @@ wheels = [ [[package]] name = "snapform" -version = "0.10.12" +version = "0.11.0" source = { virtual = "." } [package.dev-dependencies]