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

66
tests/test_addon.py
View File

@@ -138,6 +138,15 @@ class AddonTest(unittest.TestCase):
self.addon.BuildDepthPlane.bl_options,
{"INTERNAL"},
)
self.assertIn(
"snapform_color_frames_dir",
self.addon.BuildDepthPlane.__annotations__,
)
self.assertIn(
"snapform_depth_frames_dir",
self.addon.BuildDepthPlane.__annotations__,
)
self.assertIn("execute", self.addon.BuildDepthPlane.__dict__)
self.assertFalse(hasattr(self.addon, "ImportDepthPlane"))
def test_registers_perspective_mesh_operators(self):
@@ -553,7 +562,7 @@ class AddonTest(unittest.TestCase):
def test_original_depth_background_starts_independent_job(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ORIGINAL"
operator.snapform_model = "ORIGINAL"
operator.snapform_unshaded = False
operator.snapform_remove_background = True
operator.snapform_keep_original = True
@@ -586,7 +595,7 @@ class AddonTest(unittest.TestCase):
def test_original_depth_animation_extracts_frames_and_background(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ORIGINAL"
operator.snapform_model = "ORIGINAL"
operator.snapform_unshaded = False
operator.snapform_remove_background = True
operator.snapform_keep_original = True
@@ -630,7 +639,7 @@ class AddonTest(unittest.TestCase):
def test_original_depth_without_background_uses_frame_extraction(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ORIGINAL"
operator.snapform_model = "ORIGINAL"
operator.snapform_unshaded = False
operator.snapform_remove_background = False
operator.snapform_keep_original = False
@@ -664,10 +673,7 @@ class AddonTest(unittest.TestCase):
def test_conversion_dialog_exposes_requested_settings(self):
operator = self.addon.ConvertImageToDepthPlane
self.assertNotIn(
"default",
operator.__annotations__["snapform_model"],
)
self.assertEqual(operator.__annotations__["snapform_model"]["default"], 1)
self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"])
self.assertEqual(
operator.__annotations__["snapform_process_res_factor"]["default"],
@@ -693,23 +699,16 @@ class AddonTest(unittest.TestCase):
self.assertEqual(
set(operator.__annotations__),
{
"snapform_depth_source",
"snapform_unshaded",
"snapform_model",
"snapform_use_ray_pose",
"snapform_process_res_factor",
"snapform_moge_resolution_level",
"snapform_remove_background",
"snapform_keep_original",
},
)
depth_sources = operator.__annotations__["snapform_depth_source"]
self.assertEqual(depth_sources["default"], "ORIGINAL")
self.assertEqual(
[item[0] for item in depth_sources["items"]],
["ESTIMATE", "ORIGINAL"],
)
self.assertEqual(depth_sources["items"][0][1], "AI Generated")
self.assertEqual(depth_sources["name"], "Depth from")
self.assertEqual(operator.__annotations__["snapform_model"]["name"], "Model")
self.assertEqual(
operator.__annotations__["snapform_unshaded"]["name"],
"Unshaded",
@@ -724,11 +723,12 @@ class AddonTest(unittest.TestCase):
self.assertEqual(
[item[0] for item in model_items],
[
"MOGE2_VITS_NORMAL",
"MOGE2_VITB_NORMAL",
"MOGE2_VITL_NORMAL",
"SMALL",
"BASE",
"LARGE_1_1",
"GIANT_1_1",
"NESTED_GIANT_LARGE_1_1",
"ORIGINAL",
],
)
@@ -758,7 +758,7 @@ class AddonTest(unittest.TestCase):
result = operator.invoke(context, None)
self.assertIs(result, operator)
self.assertEqual(operator.snapform_model, "GIANT_1_1")
self.assertNotEqual(getattr(operator, "snapform_model", ""), "GIANT_1_1")
def test_remove_background_defaults_to_empty_image_frames(self):
source_object = SimpleNamespace(
@@ -808,7 +808,6 @@ class AddonTest(unittest.TestCase):
)
context = SimpleNamespace(object=source_object, scene=scene)
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ESTIMATE"
operator.snapform_model = "BASE"
operator.snapform_use_ray_pose = True
operator.snapform_process_res_factor = 50.0
@@ -839,7 +838,7 @@ class AddonTest(unittest.TestCase):
self.assertTrue(calls[0]["snapform_remove_background"])
self.assertNotIn("snapform_background_depth_distance", calls[0])
def test_conversion_dialog_only_lists_installed_models(self):
def test_conversion_dialog_marks_missing_models(self):
items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[
"snapform_model"
]["items"]
@@ -850,7 +849,20 @@ class AddonTest(unittest.TestCase):
):
items = items_callback(None, None)
self.assertEqual([item[0] for item in items], ["SMALL", "BASE"])
self.assertEqual(
[item[0] for item in items],
[
"MOGE2_VITS_NORMAL",
"MOGE2_VITB_NORMAL",
"MOGE2_VITL_NORMAL",
"SMALL",
"BASE",
"ORIGINAL",
],
)
labels = {item[0]: item[1] for item in items}
self.assertEqual(labels["BASE"], "DA3 Base")
self.assertEqual(labels["MOGE2_VITB_NORMAL"], "MoGe-2 ViT-B Normal (Missing)")
def test_update_status_formats_stage_progress(self):
operator = self.addon.RunJob()
@@ -936,7 +948,6 @@ class AddonTest(unittest.TestCase):
def test_packed_image_is_exported_without_unpacking(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ESTIMATE"
operator.snapform_model = "GIANT_1_1"
operator.snapform_use_ray_pose = True
operator.snapform_process_res_factor = 50.0
@@ -1869,7 +1880,7 @@ class AddonTest(unittest.TestCase):
self.assertEqual(
identifiers,
["GIANT_1_1", "NESTED_GIANT_LARGE_1_1"],
[],
)
def test_device_enums_list_only_supported_devices(self):
@@ -1966,7 +1977,7 @@ class AddonTest(unittest.TestCase):
"model_ready",
return_value=False,
):
self.assertFalse(
self.assertTrue(
self.addon.ConvertImageToDepthPlane.poll(context)
)
@@ -1977,6 +1988,7 @@ class AddonTest(unittest.TestCase):
)
panels = (
self.addon.Da3OutputDebugPanel,
self.addon.Moge2ApiDebugPanel,
self.addon.Ben2RemovalDebugPanel,
)
for panel in panels:
@@ -1984,7 +1996,7 @@ class AddonTest(unittest.TestCase):
self.assertIn("DEFAULT_CLOSED", panel.bl_options)
self.assertEqual(
[panel.bl_label for panel in panels],
["DA3 Output Test", "BEN2 Removal Test"],
["DA3 Output Test", "MoGe-2 API Test", "BEN2 Removal Test"],
)
def test_service_panel_exposes_shared_service_state_and_controls(self):

2
tests/test_download_progress.py
View File

@@ -168,7 +168,7 @@ class DownloadProgressTest(unittest.TestCase):
self.assertFalse(status_path.exists())
def test_download_endpoints_keeps_direct_endpoint(self):
def test_download_endpoints_uses_direct_endpoint(self):
with patch.dict(os.environ, {"HF_ENDPOINT": ""}):
endpoints = self.worker.download_endpoints()

70
tests/test_install_env.py
View File

@@ -10,7 +10,7 @@ import unittest
import zipfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
from unittest.mock import Mock, patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
@@ -32,6 +32,8 @@ class InstallEnvTest(unittest.TestCase):
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
(vendor / "moge.tar.gz").write_bytes(b"moge")
return vendor
def test_installs_da3_from_archives_without_gsplat(self):
@@ -155,19 +157,13 @@ class InstallEnvTest(unittest.TestCase):
if any(
repository in command
for repository in (
"depth-anything/DA3-SMALL",
"depth-anything/DA3-BASE",
"Ruicheng/moge-2-vitb-normal",
"PramaLLC/BEN2",
)
)
]
self.assertEqual(len(default_downloads), 2)
self.assertFalse(
any(
"depth-anything/DA3-BASE" in map(str, command)
for command in default_downloads
)
)
self.assertTrue(any("Ruicheng/moge-2-vitb-normal" in map(str, command) for command in default_downloads))
ben2_download = next(
command
for command in default_downloads
@@ -663,6 +659,32 @@ class InstallEnvTest(unittest.TestCase):
env["PIP_INDEX_URL"],
"https://pypi.tuna.tsinghua.edu.cn/simple",
)
self.assertEqual(
env["PYTHONIOENCODING"],
"utf-8:backslashreplace",
)
self.assertEqual(env["PYTHONUTF8"], "1")
def test_standard_streams_are_reconfigured_for_unicode_logs(self):
bootstrap = load_install_env()
stdout = Mock()
stderr = Mock()
with patch.object(bootstrap.sys, "stdout", stdout), patch.object(
bootstrap.sys,
"stderr",
stderr,
):
bootstrap.configure_standard_streams()
stdout.reconfigure.assert_called_once_with(
encoding="utf-8",
errors="backslashreplace",
)
stderr.reconfigure.assert_called_once_with(
encoding="utf-8",
errors="backslashreplace",
)
def test_ready_json_records_installed_packages(self):
bootstrap = load_install_env()
@@ -706,6 +728,10 @@ class InstallEnvTest(unittest.TestCase):
["numpy==1.26.4", "torch==2.9.1+cu128"],
)
self.assertEqual(ready["devices"], ["cuda", "cpu"])
self.assertEqual(
ready["manifest_revision"],
bootstrap.MANIFEST_REVISION,
)
def test_broken_env_is_cleaned_before_reinstall(self):
bootstrap = load_install_env()
@@ -752,7 +778,9 @@ class InstallEnvTest(unittest.TestCase):
(venv_dir / "Scripts").mkdir()
(venv_dir / "Scripts" / "python.exe").write_text("python")
(root / "manifest.json").write_text(
json.dumps({}),
json.dumps(
{"manifest_revision": bootstrap.MANIFEST_REVISION}
),
encoding="utf-8",
)
(venv_dir / "marker.txt").write_text("keep")
@@ -786,7 +814,7 @@ class InstallEnvTest(unittest.TestCase):
def fake_run(command, environment=None):
commands.append(command)
if any("DA3-SMALL" in str(value) for value in command):
if any("moge-2-vitb-normal" in str(value) for value in command):
raise subprocess.CalledProcessError(1, command)
with tempfile.TemporaryDirectory() as directory:
@@ -880,6 +908,8 @@ class InstallEnvTest(unittest.TestCase):
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
(bundled / "moge.tar.gz").write_bytes(b"moge")
target, archive_name = bootstrap.uv_download_target()
archive = bundled / f"uv-{target}.{archive_name}"
if archive_name.endswith("zip"):
@@ -1025,6 +1055,8 @@ class InstallEnvTest(unittest.TestCase):
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
(bundled / "moge.tar.gz").write_bytes(b"moge")
vendor = root / "vendor"
bootstrap.copy_vendor_inputs(bundled, vendor)
@@ -1046,6 +1078,8 @@ class InstallEnvTest(unittest.TestCase):
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"new")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
(bundled / "moge.tar.gz").write_bytes(b"moge")
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"old")
@@ -1086,6 +1120,8 @@ class InstallEnvTest(unittest.TestCase):
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
(vendor / "moge.tar.gz").write_bytes(b"moge")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
@@ -1126,6 +1162,8 @@ class InstallEnvTest(unittest.TestCase):
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
(vendor / "moge.tar.gz").write_bytes(b"moge")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
@@ -1150,9 +1188,15 @@ class InstallEnvTest(unittest.TestCase):
for label, command, _environment, kind in stages
if kind == "download"
}
self.assertIn("da3_small", " ".join(map(str, downloads["DA3 Small"])))
self.assertIn(
"moge2_vitb_normal",
" ".join(map(str, downloads["MoGe-2 ViT-B Normal"])),
)
self.assertIn("ben2_base", " ".join(map(str, downloads["BEN2 Base"])))
self.assertNotIn("DA3 Base", downloads)
self.assertEqual(
set(downloads),
{"MoGe-2 ViT-B Normal", "BEN2 Base"},
)
def test_install_does_not_verify_env(self):
bootstrap = load_install_env()

26
tests/test_models.py
View File

@@ -34,11 +34,18 @@ class ModelTest(unittest.TestCase):
model = self.models.get_model("LARGE_1_1")
self.assertEqual(model.license_name, "CC BY-NC 4.0")
self.assertFalse(model.enabled)
def test_giant_models_expose_gaussian_head(self):
for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"):
with self.subTest(model_key=model_key):
self.assertTrue(self.models.get_model(model_key).supports_gs)
self.assertFalse(self.models.get_model(model_key).enabled)
def test_only_uploaded_da3_models_are_enabled(self):
enabled = [key for key, model in self.models.MODELS.items() if model.enabled]
self.assertEqual(enabled, ["SMALL", "BASE"])
def test_unknown_model_is_rejected(self):
with self.assertRaisesRegex(ValueError, "Unknown DA3 model"):
@@ -72,3 +79,22 @@ class ModelTest(unittest.TestCase):
self.models.model_directory_name("BEN2_BASE"),
"ben2_base",
)
def test_moge2_uses_single_checkpoint_and_mit_license(self):
model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL")
self.assertEqual(model.repository, "Ruicheng/moge-2-vitb-normal")
self.assertEqual(model.license_name, "MIT")
self.assertEqual(model.download_patterns, ("model.pt",))
self.assertEqual(model.ready_patterns, ("model.pt",))
self.assertFalse(model.supports_gs)
def test_moge2_models_match_uploaded_r2_variants(self):
self.assertEqual(
list(self.models.MOGE2_MODELS),
["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE2_VITL_NORMAL"],
)
self.assertEqual(
[model.r2_directory for model in self.models.MOGE2_MODELS.values()],
["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"],
)

3
tests/test_service.py
View File

@@ -131,6 +131,9 @@ class ServiceTest(unittest.TestCase):
state.reserve("other", "second", "reconstruct", "other.json")
)
def test_service_accepts_moge_command(self):
self.assertIn("moge", self.service.JOB_COMMANDS)
def test_cancelling_job_keeps_service_and_model_cache_alive(self):
fake_models = SimpleNamespace(
status=lambda: {"ben2": "BEN2_BASE", "da3": ""},

45
tests/test_worker.py
View File

@@ -436,7 +436,13 @@ class WorkerTest(unittest.TestCase):
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
self.worker.prepare_input(input_path, 16, directory)
def test_download_model_skips_unreachable_endpoint(self):
def test_moge_runtime_disables_unsupported_xformers_attention(self):
with patch.dict(os.environ, {}, clear=True):
self.worker.prepare_moge_runtime()
self.assertEqual(os.environ["XFORMERS_DISABLED"], "1")
def test_download_model_uses_actual_request_without_preflight(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: calls.append(
@@ -457,13 +463,6 @@ class WorkerTest(unittest.TestCase):
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"hf_endpoint_reachable",
return_value=False,
), patch.object(
self.worker,
"download_from_modelscope",
) as mocked_modelscope, patch.object(
self.worker,
"_verify_model_complete",
return_value=None,
@@ -472,8 +471,7 @@ class WorkerTest(unittest.TestCase):
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.assertEqual(calls, [])
self.assertEqual(mocked_modelscope.call_count, 1)
self.assertEqual(len(calls), 1)
def test_download_model_refuses_success_when_files_missing(self):
fake_hugging_face = ModuleType("huggingface_hub")
@@ -492,10 +490,7 @@ class WorkerTest(unittest.TestCase):
self.worker,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"download_from_modelscope",
):
), patch.dict(os.environ, {"HF_ENDPOINT": ""}):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
with self.assertRaisesRegex(
@@ -527,39 +522,35 @@ class WorkerTest(unittest.TestCase):
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"download_from_modelscope",
"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, "modelscope.cn"):
with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
self.worker.download_model(args)
def test_download_from_modelscope_downloads_selected_files(self):
def test_download_from_r2_downloads_declared_files(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"],
)
downloaded = []
with tempfile.TemporaryDirectory() as directory, patch.object(
self.worker,
"modelscope_file_list",
return_value=[
("config.json", 100),
("model.safetensors", 200),
(".gitattributes", 10),
],
), patch.object(
self.worker,
"_download_modelscope_file",
"_download_r2_file",
side_effect=lambda *args, **kwargs: downloaded.append(args[1]),
), patch.object(self.worker, "write_status"):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_from_modelscope(args, args.model_dir)
self.worker.download_from_r2(args, args.model_dir)
self.assertEqual(downloaded, ["config.json"])