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

102
src/snapform/models.py
View File

@@ -11,6 +11,10 @@ class ModelSpec:
description: str
supports_gs: bool
download_patterns: tuple[str, ...] = ()
ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors")
enabled: bool = True
r2_directory: str = ""
r2_files: tuple[tuple[str, int, str], ...] = ()
MODELS = {
@@ -22,6 +26,11 @@ MODELS = {
license_name="Apache 2.0",
description="Fastest any-view model with Apache 2.0 weights",
supports_gs=False,
r2_directory="DA3Small",
r2_files=(
("config.json", 1202, "a486e29e82b7ab4a7d4cefc1ea4526cfe2ae438a572c8ca98917cfbcde7447d2"),
("model.safetensors", 137248940, "364492e38a3a06d221ac75da7f6621ada3f2361cd24fde11ba79091e9f40efcf"),
),
),
"BASE": ModelSpec(
key="BASE",
@@ -31,6 +40,12 @@ MODELS = {
license_name="Apache 2.0",
description="Balanced any-view model with Apache 2.0 weights",
supports_gs=False,
r2_directory="DA3Base",
r2_files=(
("config.json", 1205, "5e34115ebc17bd2d8d43033c5f72e9446ac8833fd61d3fa160b7e67e0bb5b7b5"),
("configuration.json", 74, "b684d553257b2ca379fd5dfdf96a6e4a633eec3b969c97119c9875465a7c07a8"),
("model.safetensors", 541518028, "e01067dc1659613083d9145a9a2547ccdbe6ccbbf83c4fe7b3e8a4e2bdae78b5"),
),
),
"LARGE_1_1": ModelSpec(
key="LARGE_1_1",
@@ -40,6 +55,7 @@ MODELS = {
license_name="CC BY-NC 4.0",
description="Higher quality non-commercial model with larger VRAM use",
supports_gs=False,
enabled=False,
),
"GIANT_1_1": ModelSpec(
key="GIANT_1_1",
@@ -49,6 +65,7 @@ MODELS = {
license_name="CC BY-NC 4.0",
description="Non-commercial Giant model with Gaussian head",
supports_gs=True,
enabled=False,
),
"NESTED_GIANT_LARGE_1_1": ModelSpec(
key="NESTED_GIANT_LARGE_1_1",
@@ -58,6 +75,7 @@ MODELS = {
license_name="CC BY-NC 4.0",
description="Non-commercial nested model with metric depth and Gaussian head",
supports_gs=True,
enabled=False,
),
}
@@ -70,8 +88,51 @@ BEN2_MODEL = ModelSpec(
description="Fast high-quality background removal model",
supports_gs=False,
download_patterns=("config.json", "model.safetensors"),
r2_directory="BEN2Base",
r2_files=(
("config.json", 124, "01ac876925577b52fa7c6af3f5b1d1420bc9b40016b6c25e2de89b38d8fff865"),
("model.safetensors", 380577976, "ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb"),
),
)
DOWNLOADABLE_MODELS = {**MODELS, BEN2_MODEL.key: BEN2_MODEL}
def moge2_model(key, label, directory_name, repository, description, r2_directory, size, checksum):
return ModelSpec(
key=key,
label=label,
directory_name=directory_name,
repository=repository,
license_name="MIT",
description=description,
supports_gs=False,
download_patterns=("model.pt",),
ready_patterns=("model.pt",),
r2_directory=r2_directory,
r2_files=(("model.pt", size, checksum),),
)
MOGE2_MODELS = {
"MOGE2_VITS_NORMAL": moge2_model(
"MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal", "moge2_vits_normal",
"Ruicheng/moge-2-vits-normal", "Fastest MoGe-2 model", "MoGe2ViTSNormal",
140550416, "79a16621928c2bf0ed04659218c55c01075e950507f40bb3332fb4c873d3e1dc",
),
"MOGE2_VITB_NORMAL": moge2_model(
"MOGE2_VITB_NORMAL", "MoGe-2 ViT-B Normal", "moge2_vitb_normal",
"Ruicheng/moge-2-vitb-normal", "Balanced MoGe-2 model", "MoGe2ViTBNormal",
419110160, "16b8110e86d5dc5a849db120ca96ef3a223fd30b0c9146d1d81db504073da5f6",
),
"MOGE2_VITL_NORMAL": moge2_model(
"MOGE2_VITL_NORMAL", "MoGe-2 ViT-L Normal", "moge2_vitl_normal",
"Ruicheng/moge-2-vitl-normal", "Highest-quality MoGe-2 model", "MoGe2ViTLNormal",
1323815904, "280741fd09bc3f403ccff9967784c2a391b52d2c0742ae3efdb21d9f90cc1a01",
),
}
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITB_NORMAL"]
DOWNLOADABLE_MODELS = {
**MODELS,
BEN2_MODEL.key: BEN2_MODEL,
**MOGE2_MODELS,
}
MODEL_NUMBERS = {
model_key: number for number, model_key in enumerate(MODELS)
}
@@ -88,7 +149,8 @@ def enum_items(supports_gs=None, show_download_status=False):
models = [
model
for model in MODELS.values()
if supports_gs is None or model.supports_gs == supports_gs
if model.enabled
and (supports_gs is None or model.supports_gs == supports_gs)
]
download_status = (
tuple(model_ready(model.key) for model in models)
@@ -130,6 +192,42 @@ def gs_model_enum_items(_owner, _context):
return enum_items(supports_gs=True, show_download_status=True)
def moge2_model_enum_items(_owner, _context):
from .client import model_ready
return [
(
model.key,
model.label if model_ready(model.key) else f"{model.label} (Missing)",
f"{model.description}; license: {model.license_name}",
index,
)
for index, model in enumerate(MOGE2_MODELS.values())
]
def depth_plane_model_enum_items(_owner, _context):
from .client import model_ready
models = [
*MOGE2_MODELS.values(),
*(model for model in MODELS.values() if model.enabled),
]
items = []
for index, model in enumerate(models):
label = model.label if model_ready(model.key) else f"{model.label} (Missing)"
items.append(
(
model.key,
label,
f"{model.description}; license: {model.license_name}",
index,
)
)
items.append(("ORIGINAL", "Original", "Use the original image as depth", len(items)))
return items
def selected_model_key(settings, export_format=None):
selected_format = export_format or settings.export_format
if selected_format == "gs_ply":