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:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user