icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update addon tests, bootstrap runtime, download progress, and worker functionality

Browse Source 
- Modified status messages and stage labels in `test_addon.py` for clarity and consistency.
- Added tests to ensure proper host termination and command execution in `test_addon.py`.
- Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks.
- Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting.
- Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints.
- Bumped version in `uv.lock` to 0.9.30.
This commit is contained in:
马南
2026-08-07 13:11:04 +08:00
parent e0cb995473
commit d08d0e9c1a
17 changed files with 1799 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
src/snapform/worker.py
View File

@@ -1,9 +1,12 @@
import argparse
import json
import os
import shutil
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
@@ -21,6 +24,9 @@ VIDEO_SUFFIXES = {
}
SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
HF_ENDPOINT = "https://huggingface.co/"
HF_MIRRORS = ("https://hf-mirror.com",)
HF_PROBE_TIMEOUT = 5.0
class TaskCancelled(RuntimeError):
@@ -32,6 +38,34 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled")
def reachable(url, timeout=HF_PROBE_TIMEOUT):
request = urllib.request.Request(
url,
method="HEAD",
headers={"User-Agent": "SnapForm-Runtime-Installer"},
)
try:
with urllib.request.urlopen(request, timeout=timeout):
return True
except urllib.error.HTTPError:
return True
except Exception:
return False
def resolve_hf_endpoint():
"""Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror)."""
configured = os.environ.get("HF_ENDPOINT")
if configured:
return configured, False
if reachable(HF_ENDPOINT):
return HF_ENDPOINT, False
for mirror in HF_MIRRORS:
if reachable(mirror):
return mirror, True
return HF_ENDPOINT, False
def write_status(
path,
progress,
@@ -159,6 +193,7 @@ class DownloadProgressReporter:
stage=None,
stages=None,
stage_label=None,
label=None,
):
self.status_path = status_path
self.repository = repository
@@ -169,10 +204,11 @@ class DownloadProgressReporter:
self.stage = stage
self.stages = stages
self.stage_label = stage_label
self.label = label
self.progress = start
self.last_write_time = 0.0
def report(self, current, total):
def report(self, current, total, rate=None):
if not total or current <= 0:
return
@@ -186,10 +222,18 @@ class DownloadProgressReporter:
return
self.progress = progress
self.last_write_time = now
speed = format_rate(rate) if rate else ""
if self.stage is not None:
message = f"downloading {speed}" if speed else "downloading"
else:
prefix = self.label or self.repository
message = f"downloading {prefix}"
if speed:
message = f"{message} {speed}"
write_status(
self.status_path,
progress,
f"Downloading {self.repository}",
message,
stage=self.stage,
stages=self.stages,
stage_label=self.stage_label,
@@ -206,6 +250,7 @@ def create_download_progress_class(
stage=None,
stages=None,
stage_label=None,
label=None,
):
if tqdm_type is None:
from tqdm.auto import tqdm as tqdm_type
@@ -219,6 +264,7 @@ def create_download_progress_class(
stage,
stages,
stage_label,
label,
)
class DownloadProgress(tqdm_type):
@@ -233,26 +279,40 @@ def create_download_progress_class(
def update(self, amount=1):
result = super().update(amount)
if self.snapform_reports_progress:
reporter.report(self.n, self.total)
reporter.report(self.n, self.total, self._download_rate())
return result
def refresh(self, *args, **kwargs):
result = super().refresh(*args, **kwargs)
if self.snapform_reports_progress:
reporter.report(self.n, self.total)
reporter.report(self.n, self.total, self._download_rate())
return result
def _download_rate(self):
try:
return getattr(self, "format_dict", {}).get("rate")
except Exception:
return None
return DownloadProgress
def download_model(args):
from huggingface_hub import snapshot_download
def progress_range(args, staged):
start = getattr(args, "progress_start", None)
end = getattr(args, "progress_end", None)
if start is not None and end is not None:
return float(start), float(end) - float(start)
if staged:
return 0.0, 1.0
return 0.05, 0.90
def download_model(args):
label = getattr(args, "label", None) or args.repository
model_dir = Path(args.model_dir)
model_dir.mkdir(parents=True, exist_ok=True)
staged = getattr(args, "stage", None) is not None
start = 0.0 if staged else 0.05
span = 1.0 if staged else 0.90
start, span = progress_range(args, staged)
stage_kwargs = {
"stage": getattr(args, "stage", None),
"stages": getattr(args, "stages", None),
@@ -261,14 +321,24 @@ def download_model(args):
write_status(
args.status,
start,
f"Downloading {args.repository}",
"checking download source",
**stage_kwargs,
)
endpoint, used_mirror = resolve_hf_endpoint()
if endpoint != os.environ.get("HF_ENDPOINT"):
os.environ["HF_ENDPOINT"] = endpoint
from huggingface_hub import snapshot_download
message = f"downloading {label}"
if used_mirror:
message = f"{message} (mirror)"
write_status(args.status, start, message, **stage_kwargs)
progress_class = create_download_progress_class(
args.status,
args.repository,
start=start,
span=span,
label=label,
**stage_kwargs,
)
snapshot_download(
@@ -279,12 +349,26 @@ def download_model(args):
)
write_status(
args.status,
1.0,
"Model is ready",
start + span,
"model ready",
**stage_kwargs,
)
def format_rate(rate):
size = float(rate)
if size <= 0:
return ""
units = ("B/s", "KB/s", "MB/s", "GB/s")
unit_index = 0
while size >= 1024 and unit_index < len(units) - 1:
size /= 1024
unit_index += 1
if unit_index == 0:
return f"{size:.0f} {units[unit_index]}"
return f"{size:.1f} {units[unit_index]}"
def select_device(requested, torch):
if requested == "cpu":
return torch.device("cpu")
@@ -292,7 +376,15 @@ def select_device(requested, torch):
if not torch.cuda.is_available():
raise RuntimeError("CUDA was requested but is not available")
return torch.device("cuda")
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
if requested == "mps":
if not torch.backends.mps.is_available():
raise RuntimeError("MPS was requested but is not available")
return torch.device("mps")
if torch.cuda.is_available():
return torch.device("cuda")
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def infer_foreground(
@@ -404,18 +496,36 @@ def exported_paths(output_directory, export_format):
return []
def npz_is_complete(path):
"""Return True only when the npz archive can be fully read."""
try:
import numpy as np
with np.load(path, allow_pickle=False) as archive:
return "depth" in archive.files
except Exception:
return False
def wait_for_exported_paths(
output_directory,
export_format,
timeout=30.0,
cancel_check=None,
):
# 上游 NPZ exporter 使用后台线程，因此需要等待文件实际落盘。
# DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整，
# 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
check_cancelled(cancel_check)
result_paths = exported_paths(output_directory, export_format)
if result_paths:
if (
export_format in ("mini_npz", "npz")
and not npz_is_complete(result_paths[0])
):
time.sleep(0.1)
continue
return result_paths
time.sleep(0.05)
return []
@@ -549,6 +659,9 @@ def parser():
download.add_argument("--stage", type=int, default=None)
download.add_argument("--stages", type=int, default=None)
download.add_argument("--stage-label", default=None)
download.add_argument("--label", default=None)
download.add_argument("--progress-start", type=float, default=None)
download.add_argument("--progress-end", type=float, default=None)
background = commands.add_parser("remove-background")
background.add_argument("--input", required=True)
@@ -557,7 +670,7 @@ def parser():
background.add_argument("--status", required=True)
background.add_argument(
"--device",
choices=("auto", "cuda", "cpu"),
choices=("auto", "cuda", "mps", "cpu"),
required=True,
)
@@ -598,7 +711,7 @@ def parser():
reconstruction.add_argument("--video-frames", type=int, required=True)
reconstruction.add_argument(
"--device",
choices=("auto", "cuda", "cpu"),
choices=("auto", "cuda", "mps", "cpu"),
required=True,
)
reconstruction.add_argument(