Add model catalog with specifications for DA3 and MoGe-2 models

- Introduced ModelSpec dataclass to define model attributes.
- Added multiple models including DA3 Small, Base, Giant 1.1, Nested Giant Large 1.1, and MoGe-2 variants.
- Implemented functions to retrieve model details and metadata.
- Included BEN2 model specifications for background removal.
This commit is contained in:
2026-08-17 14:53:25 +08:00
parent c5b77c5c85
commit 1c98b21235
44 changed files with 637 additions and 2709 deletions

View File

@@ -4,11 +4,13 @@ import sys
import bpy
from .job_runtime import runtime
from . import model_catalog as shared_model_catalog
BEN2_MODEL_KEY = "BEN2_BASE"
DEFAULT_DA3_MODEL_KEY = "BASE"
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
REQUIRED_MODEL_KEYS = (DEFAULT_MOGE2_MODEL_KEY, BEN2_MODEL_KEY)
DEVICE_LABELS = {
@@ -17,7 +19,6 @@ DEVICE_LABELS = {
"mps": ("MPS", "Apple Metal (MPS)"),
"cpu": ("CPU", "CPU"),
}
_models_cache = {}
def supported_devices():
@@ -42,22 +43,25 @@ def device_enum_items(_owner, _context):
def models_status(refresh=False):
global _models_cache
status = runtime.server_status().get("resources", {}).get("model_manager", {})
if status:
_models_cache = status
if status and not refresh:
return status
if _models_cache and not refresh:
return _models_cache
if not runtime.environment_ready():
return {}
try:
status = runtime.resource("model_manager")
_models_cache = status
return status
except RuntimeError:
return _models_cache
return {
"catalog": tuple(
shared_model_catalog.model_record(
model,
_model_files_ready(model),
)
for model in shared_model_catalog.DOWNLOADABLE_MODELS.values()
),
"loaded": status.get("loaded", {}),
}
def _model_files_ready(model):
directory = runtime.storage_root() / "models" / model.directory_name
return all(
any(directory.glob(pattern))
for pattern in model.ready_patterns
)
def model_ready(model_key):
@@ -67,6 +71,10 @@ def model_ready(model_key):
)
def required_models_ready():
return all(model_ready(model_key) for model_key in REQUIRED_MODEL_KEYS)
def model_catalog():
return models_status().get("catalog", ())
@@ -243,11 +251,6 @@ class SnapFormSettings(bpy.types.PropertyGroup):
description="Include DA3 camera wireframes in the GLB export",
default=True,
)
use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
description="Legacy task compatibility; ONNX uses the camera decoder",
default=False,
)
ref_view_strategy: bpy.props.EnumProperty(
name="Ref View Strategy",
description="DA3 reference view selection strategy",