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

feat: Add MoGe-2 support and enhance depth plane functionality

Browse Source 
- Introduced new MoGe-2 depth plane group with geometry and material displacement options.
- Updated existing depth plane group to improve background depth handling and subdivision settings.
- Enhanced tests to validate MoGe-2 models and ensure proper integration with existing systems.
- Modified installation environment to include new dependencies for MoGe-2 and utils3d.
- Updated vendor lock file to include new artifacts for MoGe and utils3d.
- Refactored download and model handling to accommodate new MoGe-2 models.
- Improved validation scripts for depth plane assets to ensure compliance with new specifications.
This commit is contained in:
马南
2026-08-13 15:55:58 +08:00
parent 92d2cea1d1
commit a6370d678d
25 changed files with 1770 additions and 258 deletions
Download Patch File Download Diff File Expand all files Collapse all files

345
src/snapform/worker.py
View File

@@ -1,5 +1,5 @@
import argparse
import fnmatch
import hashlib
import json
import os
import shutil
@@ -7,7 +7,7 @@ import tempfile
import threading
import time
from pathlib import Path
from urllib.parse import quote, urlencode
from urllib.parse import quote
from urllib.request import Request, urlopen
@@ -33,9 +33,8 @@ EXPORT_FORMATS = (
"mesh_data",
)
HF_ENDPOINT = "https://huggingface.co"
MODELSCOPE_ENDPOINT = "https://modelscope.cn"
MODELSCOPE_REVISION = "master"
MODELSCOPE_TIMEOUT = 60.0
R2_ENDPOINT = "https://models.omoolab.xyz"
DOWNLOAD_TIMEOUT = 60.0
class TaskCancelled(RuntimeError):
@@ -925,21 +924,23 @@ def download_model(args):
"stages": getattr(args, "stages", None),
"stage_label": getattr(args, "stage_label", None),
}
from huggingface_hub import snapshot_download
progress_class = create_download_progress_class(
args.status,
args.repository,
start=start,
span=span,
label=label,
**stage_kwargs,
)
failures = []
for endpoint in download_endpoints():
if not hf_endpoint_reachable(endpoint):
failures.append((endpoint, "endpoint unreachable"))
continue
try:
from huggingface_hub import snapshot_download
except Exception as error:
snapshot_download = None
failures.append((HF_ENDPOINT, error))
progress_class = None
if snapshot_download:
progress_class = create_download_progress_class(
args.status,
args.repository,
start=start,
span=span,
label=label,
**stage_kwargs,
)
for endpoint in download_endpoints() if snapshot_download else ():
os.environ["HF_ENDPOINT"] = endpoint
message = f"Downloading {label}"
if endpoint != HF_ENDPOINT:
@@ -953,7 +954,10 @@ def download_model(args):
allow_patterns=args.allow_pattern or None,
endpoint=endpoint,
)
_verify_model_complete(model_dir)
_verify_model_complete(
model_dir,
getattr(args, "required_pattern", None),
)
except Exception as error:
failures.append((endpoint, error))
continue
@@ -964,15 +968,22 @@ def download_model(args):
**stage_kwargs,
)
return
# Hugging Face 官方源与已配置镜像均失败时，回退 ModelScope（仓库名与 Hugging Face 一致）
# 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} | ModelScope mirror",
f"Downloading {label} | OmooLab mirror",
**stage_kwargs,
)
download_from_modelscope(
download_from_r2(
args,
model_dir,
start=start,
@@ -980,9 +991,12 @@ def download_model(args):
label=label,
stage_kwargs=stage_kwargs,
)
_verify_model_complete(model_dir)
_verify_model_complete(
model_dir,
getattr(args, "required_pattern", None),
)
except Exception as error:
failures.append((MODELSCOPE_ENDPOINT, error))
failures.append((R2_ENDPOINT, error))
else:
write_status(
args.status,
@@ -998,33 +1012,23 @@ def download_model(args):
def download_endpoints():
"""Hugging Face download candidates: configured mirror first, official endpoint otherwise."""
"""Return the configured Hugging Face endpoint or official source."""
endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT
return [endpoint.rstrip("/")]
def hf_endpoint_reachable(endpoint, timeout=10.0):
"""Quick connectivity probe so slow or blocked Hugging Face endpoints fail fast."""
try:
request = Request(
f"{endpoint.rstrip('/')}/api/models",
method="HEAD",
headers={"User-Agent": "SnapForm"},
)
with urlopen(request, timeout=timeout):
return True
except Exception:
return False
def _verify_model_complete(model_dir):
def _verify_model_complete(model_dir, required_patterns=None):
"""Refuse to report success when the required model files are missing."""
model_dir = Path(model_dir)
missing = []
if not (model_dir / "config.json").is_file():
missing.append("config.json")
if not any(model_dir.glob("*.safetensors")):
missing.append("*.safetensors")
required_patterns = required_patterns or (
"config.json",
"*.safetensors",
)
missing = [
pattern
for pattern in required_patterns
if not any(model_dir.glob(pattern))
]
if missing:
raise RuntimeError(
"Model download reported success but files are missing: "
@@ -1032,59 +1036,42 @@ def _verify_model_complete(model_dir):
)
def _hidden_repo_path(path):
return any(part.startswith(".") for part in Path(path).parts)
def _model_file_allowed(path, patterns):
if not patterns:
return True
return any(fnmatch.fnmatch(path, pattern) for pattern in patterns)
def modelscope_file_list(repository, revision=MODELSCOPE_REVISION):
"""List (path, size) file entries of a ModelScope repository."""
url = (
f"{MODELSCOPE_ENDPOINT}/api/v1/models/{repository}/repo/files?"
+ urlencode({"Revision": revision, "Recursive": "true"})
)
request = Request(url, headers={"User-Agent": "SnapForm"})
with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response:
payload = json.loads(response.read().decode("utf-8"))
entries = []
for entry in payload.get("Data", {}).get("Files", []):
path = entry.get("Path")
if path and entry.get("Type") == "blob" and not _hidden_repo_path(path):
entries.append((path, int(entry.get("Size") or 0)))
return entries
def _download_modelscope_file(
repository,
path,
def _download_r2_file(
directory,
filename,
destination,
expected_size,
expected_checksum,
progress=None,
retries=2,
):
"""Download one ModelScope file atomically with a per-chunk progress callback."""
url = (
f"{MODELSCOPE_ENDPOINT}/models/{repository}/resolve/"
f"{MODELSCOPE_REVISION}/{quote(path)}"
)
"""Download and verify one OmooLab R2 model file atomically."""
url = f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}"
request = Request(url, headers={"User-Agent": "SnapForm"})
last_error = None
for _attempt in range(retries):
temporary = destination.with_name(f"{destination.name}.part")
try:
with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response:
checksum = hashlib.sha256()
downloaded = 0
with urlopen(request, timeout=DOWNLOAD_TIMEOUT) as response:
with open(temporary, "wb") as output:
while True:
chunk = response.read(1024 * 1024)
if not chunk:
break
output.write(chunk)
checksum.update(chunk)
downloaded += len(chunk)
if progress is not None:
progress(len(chunk))
if downloaded != expected_size:
raise RuntimeError(
f"Size mismatch for {filename}: "
f"{downloaded} != {expected_size}"
)
if checksum.hexdigest().lower() != expected_checksum.lower():
raise RuntimeError(f"SHA-256 mismatch for {filename}")
temporary.replace(destination)
return
except Exception as error:
@@ -1096,7 +1083,7 @@ def _download_modelscope_file(
raise last_error
def download_from_modelscope(
def download_from_r2(
args,
model_dir,
start=0.05,
@@ -1104,19 +1091,12 @@ def download_from_modelscope(
label=None,
stage_kwargs=None,
):
"""Download a model from ModelScope using only standard library requests."""
"""Download the declared model files from OmooLab R2."""
repository = args.repository
files = modelscope_file_list(repository)
patterns = args.allow_pattern or None
selected = [
(path, size)
for path, size in files
if _model_file_allowed(path, patterns)
]
if not selected:
raise RuntimeError(
f"No matching files in ModelScope repository: {repository}"
)
selected = []
for value in args.r2_file:
filename, size, checksum = value.split("|", 2)
selected.append((filename, int(size), checksum))
reporter = DownloadProgressReporter(
args.status,
repository,
@@ -1125,7 +1105,7 @@ def download_from_modelscope(
label=label or repository,
**(stage_kwargs or {}),
)
total_size = sum(size for _path, size in selected)
total_size = sum(size for _path, size, _checksum in selected)
downloaded = 0
def on_chunk(count):
@@ -1133,13 +1113,15 @@ def download_from_modelscope(
downloaded += count
reporter.report(downloaded, total_size)
for path, _size in selected:
for path, size, checksum in selected:
destination = Path(model_dir) / path
destination.parent.mkdir(parents=True, exist_ok=True)
_download_modelscope_file(
repository,
_download_r2_file(
args.r2_directory,
path,
destination,
size,
checksum,
progress=on_chunk,
)
@@ -1500,6 +1482,168 @@ def flatten_exported_paths(
return flattened_paths
def _save_moge_preview(output_directory, index, depth, normal, mask):
import numpy as np
from PIL import Image
suffix = "" if index is None else f".{index + 1:04d}"
valid = mask & np.isfinite(depth) & (depth > 0)
preview = np.zeros(depth.shape, dtype=np.uint8)
if np.any(valid):
low, high = np.percentile(depth[valid], (2.0, 98.0))
span = max(float(high - low), 1e-6)
preview[valid] = np.clip(
(depth[valid] - low) / span * 255.0,
0.0,
255.0,
).astype(np.uint8)
Image.fromarray(preview, mode="L").save(
output_directory / f"depth{suffix}.png"
)
normal_rgb = np.clip(
(normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0)
* 127.5,
0.0,
255.0,
).astype(np.uint8)
Image.fromarray(normal_rgb, mode="RGB").save(
output_directory / f"normal{suffix}.png"
)
Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(
output_directory / f"mask{suffix}.png"
)
def run_moge(args, model_cache=None, cancel_check=None):
import numpy as np
from PIL import Image
prepare_moge_runtime()
write_status(args.status, 0.02, "Starting MoGe-2")
if model_cache is None:
import torch
from moge.model.v2 import MoGeModel
else:
torch = model_cache.torch
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input(
args.input,
args.video_frames,
output_directory,
start_frame=getattr(args, "start_frame", 1),
cancel_check=cancel_check,
)
try:
device = select_device(args.device, torch)
write_status(args.status, 0.08, f"Loading MoGe-2 on {device.type}")
if model_cache is None:
checkpoint = Path(args.model_dir) / "model.pt"
model = MoGeModel.from_pretrained(checkpoint)
model = model.to(device=device)
model.eval()
else:
model = model_cache.get_moge(args.model_dir, device)
depths = []
normals = []
masks = []
points = []
intrinsics = []
for index, image_path in enumerate(images):
check_cancelled(cancel_check)
write_status(
args.status,
0.15 + 0.68 * index / len(images),
f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})",
)
with Image.open(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device)
with torch.inference_mode():
prediction = model.infer(
image_tensor,
resolution_level=int(args.resolution_level),
force_projection=True,
apply_mask=False,
use_fp16=device.type == "cuda",
)
depth = prediction["depth"].detach().float().cpu().numpy()
normal = prediction["normal"].detach().float().cpu().numpy()
mask = prediction["mask"].detach().cpu().numpy().astype(bool)
point = prediction["points"].detach().float().cpu().numpy()
intrinsic = prediction["intrinsics"].detach().float().cpu().numpy()
depths.append(depth)
normals.append(normal)
masks.append(mask)
intrinsics.append(intrinsic)
if bool(getattr(args, "include_points", False)):
points.append(point)
_save_moge_preview(
output_directory,
None if len(images) == 1 else index,
depth,
normal,
mask,
)
shapes = {array.shape for array in depths}
if len(shapes) != 1:
raise RuntimeError("MoGe-2 frame sizes must match")
result = {
"depth": np.stack(depths).astype(np.float32),
"normal": np.stack(normals).astype(np.float32),
"mask": np.stack(masks),
"intrinsics": np.stack(intrinsics).astype(np.float32),
}
if points:
result["points"] = np.stack(points).astype(np.float32)
np.savez(output_directory / "results.npz", **result)
remove_background_model_dir = getattr(
args, "remove_background_model_dir", ""
)
if remove_background_model_dir:
write_status(args.status, 0.87, "Removing image background")
if len(images) == 1:
foreground, _, _ = infer_foreground(
images[0],
remove_background_model_dir,
device,
torch,
model_cache=model_cache,
cancel_check=cancel_check,
)
foreground.save(output_directory / "foreground.png")
else:
infer_foreground_frames(
images,
remove_background_model_dir,
device,
torch,
output_directory,
args.status,
0.87,
0.08,
model_cache=model_cache,
cancel_check=cancel_check,
)
copy_frames(images, output_directory)
write_status(
args.status,
1.0,
f"Exported MoGe-2 data for {len(images)} frame(s)",
)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
def prepare_moge_runtime():
"""Use PyTorch SDPA because xFormers does not support every CUDA GPU."""
os.environ["XFORMERS_DISABLED"] = "1"
def reconstruct(args, model_cache=None, cancel_check=None):
write_status(args.status, 0.02, "Starting Environment")
check_cancelled(cancel_check)
@@ -1686,6 +1830,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("--required-pattern", action="append", default=[])
download.add_argument("--r2-directory", default="")
download.add_argument("--r2-file", 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)