nan
caeb6d706c
- Introduced new enum items for depth plane and surface models with caching. - Added MOGE2_MESH_JOB to operators for converting images to surfaces. - Implemented surface mesh building and exporting in RunJob operator. - Enhanced UI to include export format options for MoGe-2. - Created validation tools for MoGe-2 surface assets and depth planes. - Updated geometry nodes to support new surface generation features. - Refactored existing tests to accommodate new functionality and ensure stability.
277 lines
8.6 KiB
Python
277 lines
8.6 KiB
Python
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModelSpec:
|
|
key: str
|
|
label: str
|
|
directory_name: str
|
|
repository: str
|
|
license_name: str
|
|
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 = {
|
|
"SMALL": ModelSpec(
|
|
key="SMALL",
|
|
label="DA3 Small",
|
|
directory_name="da3_small",
|
|
repository="depth-anything/DA3-SMALL",
|
|
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",
|
|
label="DA3 Base",
|
|
directory_name="da3_base",
|
|
repository="depth-anything/DA3-BASE",
|
|
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",
|
|
label="DA3 Large 1.1",
|
|
directory_name="da3_large_1_1",
|
|
repository="depth-anything/DA3-LARGE-1.1",
|
|
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",
|
|
label="DA3 Giant 1.1",
|
|
directory_name="da3_giant_1_1",
|
|
repository="depth-anything/DA3-GIANT-1.1",
|
|
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",
|
|
label="DA3 Nested Giant Large 1.1",
|
|
directory_name="da3_nested_giant_large_1_1",
|
|
repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1",
|
|
license_name="CC BY-NC 4.0",
|
|
description="Non-commercial nested model with metric depth and Gaussian head",
|
|
supports_gs=True,
|
|
enabled=False,
|
|
),
|
|
}
|
|
|
|
BEN2_MODEL = ModelSpec(
|
|
key="BEN2_BASE",
|
|
label="BEN2 Base",
|
|
directory_name="ben2_base",
|
|
repository="PramaLLC/BEN2",
|
|
license_name="MIT",
|
|
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"),
|
|
),
|
|
)
|
|
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)
|
|
}
|
|
|
|
|
|
_ENUM_ITEMS_CACHE = {}
|
|
_DEPTH_PLANE_ENUM_ITEMS_CACHE = {}
|
|
_SURFACE_ENUM_ITEMS_CACHE = {}
|
|
|
|
|
|
def enum_items(supports_gs=None, show_download_status=False):
|
|
model_ready = None
|
|
if show_download_status:
|
|
from .client import model_ready
|
|
|
|
models = [
|
|
model
|
|
for model in MODELS.values()
|
|
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)
|
|
if model_ready is not None
|
|
else ()
|
|
)
|
|
cache_key = (
|
|
supports_gs,
|
|
show_download_status,
|
|
download_status,
|
|
)
|
|
cached_items = _ENUM_ITEMS_CACHE.get(cache_key)
|
|
if cached_items is not None:
|
|
return cached_items
|
|
|
|
items = []
|
|
for index, model in enumerate(models):
|
|
label = model.label
|
|
if download_status and not download_status[index]:
|
|
label = f"{label} (Missing)"
|
|
items.append(
|
|
(
|
|
model.key,
|
|
label,
|
|
f"{model.description}; license: {model.license_name}",
|
|
MODEL_NUMBERS[model.key],
|
|
)
|
|
)
|
|
|
|
_ENUM_ITEMS_CACHE[cache_key] = items
|
|
return items
|
|
|
|
|
|
def model_enum_items(_owner, _context):
|
|
return enum_items(show_download_status=True)
|
|
|
|
|
|
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),
|
|
]
|
|
ready = tuple(model_ready(model.key) for model in models)
|
|
cached = _DEPTH_PLANE_ENUM_ITEMS_CACHE.get(ready)
|
|
if cached is not None:
|
|
return cached
|
|
items = []
|
|
for index, model in enumerate(models):
|
|
label = model.label if ready[index] 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)))
|
|
_DEPTH_PLANE_ENUM_ITEMS_CACHE[ready] = items
|
|
return items
|
|
|
|
|
|
def surface_model_enum_items(owner, context):
|
|
depth_items = depth_plane_model_enum_items(owner, context)
|
|
cache_key = id(depth_items)
|
|
cached = _SURFACE_ENUM_ITEMS_CACHE.get(cache_key)
|
|
if cached is not None:
|
|
return cached
|
|
items = [
|
|
item
|
|
for item in depth_items
|
|
if item[0] != "ORIGINAL"
|
|
]
|
|
_SURFACE_ENUM_ITEMS_CACHE[cache_key] = items
|
|
return items
|
|
|
|
|
|
def selected_model_key(settings, export_format=None):
|
|
selected_format = export_format or settings.export_format
|
|
if selected_format == "gs_ply":
|
|
return settings.gs_model
|
|
return settings.model
|
|
|
|
|
|
def get_model(model_key):
|
|
try:
|
|
return MODELS[model_key]
|
|
except KeyError as error:
|
|
raise ValueError(f"Unknown DA3 model: {model_key}") from error
|
|
|
|
|
|
def get_downloadable_model(model_key):
|
|
try:
|
|
return DOWNLOADABLE_MODELS[model_key]
|
|
except KeyError as error:
|
|
raise ValueError(f"Unknown model: {model_key}") from error
|
|
|
|
|
|
def model_directory_name(model_key):
|
|
"""Return the model's folder name inside the shared models directory."""
|
|
return get_downloadable_model(model_key).directory_name
|