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

优化 Runtime 安装进度体验

Browse Source 
- 状态栏与窗口强制刷新，进度更新无需移动鼠标即可实时显示
- 安装过程拆分为 12 个阶段，每段进度从 0 到 100%
- 安装阶段解析 uv 下载事件推进进度，模型下载阶段实时反映字节进度
This commit is contained in:
马南
2026-08-05 22:45:58 +08:00
parent 838af0ee96
commit e0cb995473
11 changed files with 552 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/snapform/blender_manifest.toml
View File

@@ -1,7 +1,7 @@
schema_version = "1.0.0"
id = "snapform"
version = "0.9.28"
version = "0.9.29"
name = "SnapForm"
tagline = "Build spatial structures from images, fast"
maintainer = "OmooLab"

421
src/snapform/bootstrap_runtime.py
View File

@@ -28,10 +28,22 @@ TORCH_INDEXES = {
DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key)
def write_status(path, progress, message, error=None):
def write_status(
path,
progress,
message,
error=None,
stage=None,
stages=None,
stage_label=None,
):
payload = {"progress": progress, "message": message}
if error:
payload["error"] = error
if stage is not None:
payload["stage"] = stage
payload["stages"] = stages
payload["stage_label"] = stage_label
Path(path).write_text(
json.dumps(payload, ensure_ascii=False),
encoding="utf-8",
@@ -50,185 +62,282 @@ def run(command, environment=None):
)
def run_install(status_path, stage, stages, label, command, environment=None):
"""Run a uv install command and translate its verbose package events into progress."""
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
process_environment = os.environ.copy()
process_environment["UV_LOG"] = "debug"
if environment:
process_environment.update(environment)
process = subprocess.Popen(
[str(value) for value in command],
env=process_environment,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding="utf-8",
errors="replace",
)
progress = 0.0
downloaded = 0
for raw_line in process.stdout:
line = raw_line.rstrip()
print(line)
lowered = line.lower()
if "resolved" in lowered and "package" in lowered:
new_progress, message = 0.05, "Resolving dependencies"
elif "downloaded" in lowered and "package" not in lowered:
downloaded += 1
new_progress = min(0.85, 0.10 + downloaded * 0.08)
message = f"Downloaded {downloaded} package(s)"
elif "prepared" in lowered:
new_progress, message = 0.70, "Preparing packages"
elif "installed" in lowered or "built" in lowered:
new_progress, message = 0.95, "Installing packages"
else:
continue
if new_progress <= progress:
continue
progress = new_progress
write_status(
status_path,
progress,
message,
stage=stage,
stages=stages,
stage_label=label,
)
return_code = process.wait()
if return_code != 0:
raise subprocess.CalledProcessError(return_code, command)
def environment_python(runtime_dir):
if sys.platform == "win32":
return Path(runtime_dir) / "Scripts" / "python.exe"
return Path(runtime_dir) / "bin" / "python"
def download_default_models(args, python):
models = {**MODELS, BEN2_MODEL.key: BEN2_MODEL}
download_script = (
"import json, sys; "
"from huggingface_hub import snapshot_download; "
"patterns = json.loads(sys.argv[3]); "
"snapshot_download(repo_id=sys.argv[1], local_dir=sys.argv[2], "
"allow_patterns=patterns or None)"
)
models_dir = Path(args.models_dir)
models_dir.mkdir(parents=True, exist_ok=True)
for index, model_key in enumerate(DEFAULT_MODEL_KEYS):
model = models[model_key]
progress = 0.90 + index * 0.03
write_status(args.status, progress, f"Downloading {model.label}")
run(
[
python,
"-c",
download_script,
model.repository,
models_dir / model.key.lower(),
json.dumps(model.download_patterns),
]
)
def install(args):
runtime_dir = Path(args.runtime_dir)
runtime_dir.parent.mkdir(parents=True, exist_ok=True)
write_status(args.status, 0.03, "Creating isolated Python Runtime")
run(
[
args.uv,
"venv",
"--clear",
"--python",
"3.12",
runtime_dir,
]
)
python = environment_python(runtime_dir)
write_status(args.status, 0.12, "Installing PyTorch")
run(
[
args.uv,
"pip",
"install",
"--python",
python,
f"torch=={TORCH_VERSION}",
f"torchvision=={TORCHVISION_VERSION}",
"--index-url",
TORCH_INDEXES[args.backend],
]
)
write_status(args.status, 0.38, "Installing matching xFormers")
xformers_command = [
args.uv,
"pip",
"install",
"--python",
python,
f"xformers=={XFORMERS_VERSION}",
]
if args.backend == "cu128":
xformers_command.extend(
["--index-url", TORCH_INDEXES[args.backend]]
)
else:
xformers_command.append("--no-deps")
run(xformers_command)
write_status(args.status, 0.50, "Installing Depth Anything 3")
def install_stages(args, runtime_dir, python):
uv = args.uv
torch_index = TORCH_INDEXES[args.backend]
da3_requirement = (
"depth-anything-3 @ "
"git+https://github.com/ByteDance-Seed/Depth-Anything-3.git"
f"@{args.da3_commit}"
)
run(
[
args.uv,
"pip",
"install",
"--python",
python,
da3_requirement,
f"addict=={ADDICT_VERSION}",
]
)
write_status(args.status, 0.68, "Installing BEN2")
ben2_requirement = (
"ben2 @ git+https://github.com/PramaLLC/BEN2.git"
f"@{args.ben2_commit}"
)
run(
[
args.uv,
"pip",
"install",
"--python",
python,
ben2_requirement,
]
xformers_command = [
uv,
"pip",
"install",
"-v",
"--python",
python,
f"xformers=={XFORMERS_VERSION}",
]
if args.backend == "cu128":
xformers_command.extend(["--index-url", torch_index])
else:
xformers_command.append("--no-deps")
gsplat_requirement = (
"gsplat @ git+https://github.com/nerfstudio-project/"
f"gsplat.git@{GSPLAT_COMMIT}"
)
write_status(args.status, 0.75, "Installing AI Host")
run(
[
args.uv,
"pip",
"install",
"--python",
python,
f"fastapi=={FASTAPI_VERSION}",
f"uvicorn=={UVICORN_VERSION}",
]
check_code = (
"import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"import fastapi, uvicorn; "
"import gsplat; "
"print(torch.__version__, torch.cuda.is_available())"
)
write_status(args.status, 0.78, "Installing gsplat build tools")
run(
[
args.uv,
"pip",
stages = [
(
"Create Python Runtime",
[uv, "venv", "--clear", "--python", "3.12", runtime_dir],
None,
"plain",
),
(
"Install PyTorch",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
f"torch=={TORCH_VERSION}",
f"torchvision=={TORCHVISION_VERSION}",
"--index-url",
torch_index,
],
None,
"install",
"--python",
),
("Install xFormers", xformers_command, None, "install"),
(
"Install Depth Anything 3",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
da3_requirement,
f"addict=={ADDICT_VERSION}",
],
None,
"install",
),
(
"Install BEN2",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
ben2_requirement,
],
None,
"install",
),
(
"Install AI Host",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
f"fastapi=={FASTAPI_VERSION}",
f"uvicorn=={UVICORN_VERSION}",
],
None,
"install",
),
(
"Install gsplat build tools",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
"setuptools",
"wheel",
"ninja",
],
None,
"install",
),
(
"Install gsplat",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
"--no-build-isolation",
gsplat_requirement,
],
# PLY 导出不调用 rasterizer，禁用本地扩展可避免 CUDA 与 MSVC 工具链耦合
{"BUILD_NO_CUDA": "1"},
"install",
),
("Check Runtime", [python, "-c", check_code], None, "plain"),
]
models = {**MODELS, BEN2_MODEL.key: BEN2_MODEL}
worker_script = Path(__file__).with_name("worker.py")
models_dir = Path(args.models_dir)
for model_key in DEFAULT_MODEL_KEYS:
model = models[model_key]
download_command = [
python,
"setuptools",
"wheel",
"ninja",
worker_script,
"download",
"--repository",
model.repository,
"--model-dir",
models_dir / model.key.lower(),
"--status",
args.status,
]
)
write_status(args.status, 0.84, "Installing gsplat")
# PLY 导出不调用 rasterizer，禁用本地扩展可避免 CUDA 与 MSVC 工具链耦合。
gsplat_environment = {"BUILD_NO_CUDA": "1"}
run(
[
args.uv,
"pip",
"install",
"--python",
python,
"--no-build-isolation",
for pattern in model.download_patterns:
download_command.extend(["--allow-pattern", pattern])
stages.append(
(
"gsplat @ git+https://github.com/nerfstudio-project/"
f"gsplat.git@{GSPLAT_COMMIT}"
),
],
gsplat_environment,
f"Download {model.label}",
download_command,
None,
"download",
)
)
return stages
def run_stage(status_path, stage, stages, label, command, environment, kind):
write_status(
status_path,
0.0,
"",
stage=stage,
stages=stages,
stage_label=label,
)
if kind == "install":
run_install(status_path, stage, stages, label, command, environment)
elif kind == "download":
run(
[
*command,
"--stage",
str(stage),
"--stages",
str(stages),
"--stage-label",
label,
],
environment,
)
else:
run(command, environment)
write_status(
status_path,
1.0,
"",
stage=stage,
stages=stages,
stage_label=label,
)
write_status(args.status, 0.88, "Checking Runtime")
run(
[
python,
"-c",
(
"import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"import fastapi, uvicorn; "
"import gsplat; "
"print(torch.__version__, torch.cuda.is_available())"
),
]
)
download_default_models(args, python)
def install(args):
runtime_dir = Path(args.runtime_dir)
runtime_dir.parent.mkdir(parents=True, exist_ok=True)
python = environment_python(runtime_dir)
stages = install_stages(args, runtime_dir, python)
for index, (label, command, environment, kind) in enumerate(stages, start=1):
run_stage(
args.status,
index,
len(stages),
label,
command,
environment,
kind,
)
ready_path = runtime_dir / "snapform_ready.json"
ready_path.write_text(

16
src/snapform/operators.py
View File

@@ -412,8 +412,20 @@ class RunJob(bpy.types.Operator):
return
settings = context.scene.snapform_settings
settings.progress = float(status.get("progress", settings.progress))
settings.status = status.get("message", settings.status)
self._redraw_ui(context)
settings.status = self._format_status(status, settings.status)
self._redraw_ui(context, force=True)
def _format_status(self, status, fallback):
message = status.get("message", fallback)
stage = status.get("stage")
stages = status.get("stages")
stage_label = status.get("stage_label")
if not stage or not stages:
return message
stage_text = f"Step {stage}/{stages}"
if stage_label:
stage_text = f"{stage_text} {stage_label}"
return f"{stage_text}: {message}" if message else stage_text
def _finish(self, context, job, return_code):
settings = context.scene.snapform_settings

14
src/snapform/status_bar.py
View File

@@ -1,3 +1,5 @@
import bpy
from .operators import active_job
@@ -45,3 +47,15 @@ def tag_ui_regions(context):
def redraw_ui(context, force=False):
tag_ui_regions(context)
if not force:
return
try:
bpy.ops.wm.redraw_timer(type="DRAW_WIN_SWAP", iterations=1)
except (AttributeError, RuntimeError, TypeError):
pass
workspace = getattr(context, "workspace", None)
if workspace is not None:
try:
workspace.status_text_set_internal(None)
except (AttributeError, RuntimeError, TypeError):
pass

83
src/snapform/worker.py
View File

@@ -32,10 +32,23 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled")
def write_status(path, progress, message, error=None, state=None):
def write_status(
path,
progress,
message,
error=None,
state=None,
stage=None,
stages=None,
stage_label=None,
):
payload = {"progress": progress, "message": message}
if error:
payload["error"] = error
if stage is not None:
payload["stage"] = stage
payload["stages"] = stages
payload["stage_label"] = stage_label
if state is None:
if error:
state = "failed"
@@ -136,12 +149,27 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
class DownloadProgressReporter:
def __init__(self, status_path, repository, clock=time.monotonic):
def __init__(
self,
status_path,
repository,
clock=time.monotonic,
start=0.05,
span=0.90,
stage=None,
stages=None,
stage_label=None,
):
self.status_path = status_path
self.repository = repository
self.clock = clock
self.lock = threading.Lock()
self.progress = 0.05
self.start = start
self.span = span
self.stage = stage
self.stages = stages
self.stage_label = stage_label
self.progress = start
self.last_write_time = 0.0
def report(self, current, total):
@@ -149,7 +177,7 @@ class DownloadProgressReporter:
return
factor = min(max(current / total, 0.0), 1.0)
progress = 0.05 + factor * 0.90
progress = self.start + factor * self.span
now = self.clock()
with self.lock:
if progress <= self.progress:
@@ -162,6 +190,9 @@ class DownloadProgressReporter:
self.status_path,
progress,
f"Downloading {self.repository}",
stage=self.stage,
stages=self.stages,
stage_label=self.stage_label,
)
@@ -170,11 +201,25 @@ def create_download_progress_class(
repository,
tqdm_type=None,
clock=time.monotonic,
start=0.05,
span=0.90,
stage=None,
stages=None,
stage_label=None,
):
if tqdm_type is None:
from tqdm.auto import tqdm as tqdm_type
reporter = DownloadProgressReporter(status_path, repository, clock)
reporter = DownloadProgressReporter(
status_path,
repository,
clock,
start,
span,
stage,
stages,
stage_label,
)
class DownloadProgress(tqdm_type):
def __init__(self, *args, **kwargs):
@@ -205,10 +250,26 @@ def download_model(args):
model_dir = Path(args.model_dir)
model_dir.mkdir(parents=True, exist_ok=True)
write_status(args.status, 0.05, f"Downloading {args.repository}")
staged = getattr(args, "stage", None) is not None
start = 0.0 if staged else 0.05
span = 1.0 if staged else 0.90
stage_kwargs = {
"stage": getattr(args, "stage", None),
"stages": getattr(args, "stages", None),
"stage_label": getattr(args, "stage_label", None),
}
write_status(
args.status,
start,
f"Downloading {args.repository}",
**stage_kwargs,
)
progress_class = create_download_progress_class(
args.status,
args.repository,
start=start,
span=span,
**stage_kwargs,
)
snapshot_download(
repo_id=args.repository,
@@ -216,7 +277,12 @@ def download_model(args):
tqdm_class=progress_class,
allow_patterns=args.allow_pattern or None,
)
write_status(args.status, 1.0, "Model is ready")
write_status(
args.status,
1.0,
"Model is ready",
**stage_kwargs,
)
def select_device(requested, torch):
@@ -480,6 +546,9 @@ def parser():
download.add_argument("--model-dir", required=True)
download.add_argument("--status", required=True)
download.add_argument("--allow-pattern", action="append", default=[])
download.add_argument("--stage", type=int, default=None)
download.add_argument("--stages", type=int, default=None)
download.add_argument("--stage-label", default=None)
background = commands.add_parser("remove-background")
background.add_argument("--input", required=True)