AnyImage/src/snapform/properties.py

import platform
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 = {
    "cuda": ("CUDA", "NVIDIA CUDA"),
    "directml": ("DirectML", "Windows GPU (DirectML)"),
    "mps": ("MPS", "Apple Metal (MPS)"),
    "cpu": ("CPU", "CPU"),
}


def supported_devices():
    devices = runtime.environment_manifest().get("devices")
    if isinstance(devices, list) and devices:
        return [str(device).lower() for device in devices]
    if sys.platform == "win32":
        return ["directml", "cpu"]
    if sys.platform == "darwin" and platform.machine().lower() in {
        "arm64",
        "aarch64",
    }:
        return ["mps", "cpu"]
    return ["cpu"]


def device_enum_items(_owner, _context):
    return [
        (name.upper(), *DEVICE_LABELS[name], index)
        for index, name in enumerate(supported_devices())
    ]


def models_status(refresh=False):
    status = runtime.server_status().get("resources", {}).get("model_manager", {})
    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):
    return any(
        model.get("key") == model_key and model.get("ready")
        for model in models_status().get("catalog", ())
    )


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", ())


def model_record(model_key):
    return next(
        (model for model in model_catalog() if model.get("key") == model_key),
        None,
    )


def _model_items(models):
    return [
        (
            model["key"],
            model["label"] if model.get("ready") else f"{model['label']} (Missing)",
            f"{model['description']}; license: {model['license']}",
            index,
        )
        for index, model in enumerate(models)
    ]


def model_enum_items(_owner, _context):
    return _model_items(
        model
        for model in model_catalog()
        if model.get("family") == "da3" and model.get("enabled")
    )


def moge2_model_enum_items(_owner, _context):
    return _model_items(
        model for model in model_catalog() if model.get("family") == "moge2"
    )


def depth_plane_model_enum_items(_owner, _context):
    models = [
        model
        for model in model_catalog()
        if model.get("family") in {"moge2", "da3"} and model.get("enabled")
    ]
    models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
    items = _model_items(models)
    items.append(
        ("ORIGINAL", "Original", "Use the original image as depth", len(items))
    )
    return items


def surface_model_enum_items(_owner, _context):
    models = [
        model
        for model in model_catalog()
        if model.get("family") in {"moge2", "da3"}
        and model.get("enabled")
    ]
    models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
    return _model_items(models)


class SnapFormSettings(bpy.types.PropertyGroup):
    input_path: bpy.props.StringProperty(
        name="Image or Video",
        description="A single image or video file used for reconstruction",
        subtype="FILE_PATH",
    )
    background_input_path: bpy.props.StringProperty(
        name="Image",
        description="A single image used to test BEN2 background removal",
        subtype="FILE_PATH",
    )
    video_frames: bpy.props.IntProperty(
        name="Extract Frames",
        description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
        default=16,
        min=1,
        soft_max=64,
    )
    model: bpy.props.EnumProperty(
        name="Model",
        description="Depth Anything 3 model",
        items=model_enum_items,
    )
    device: bpy.props.EnumProperty(
        name="Debug Device",
        description=(
            "Inference device used by Debug tasks only; production tasks use "
            "the Device setting in the Add-on Preferences"
        ),
        items=device_enum_items,
        default=0,
    )
    process_res: bpy.props.IntProperty(
        name="Process Res",
        description="DA3 base processing resolution",
        default=504,
        min=1,
        soft_max=2048,
    )
    process_res_method: bpy.props.EnumProperty(
        name="Process Res Method",
        description="DA3 image resize method",
        items=(
            (
                "upper_bound_resize",
                "Upper Bound Resize",
                "Make Process Res the longer side",
            ),
            (
                "lower_bound_resize",
                "Lower Bound Resize",
                "Make Process Res the shorter side",
            ),
        ),
        default="upper_bound_resize",
    )
    moge_resolution_level: bpy.props.IntProperty(
        name="Resolution Level",
        description="MoGe-2 inference resolution level; higher is finer and slower",
        default=5,
        min=0,
        max=9,
    )
    moge_model: bpy.props.EnumProperty(
        name="Model",
        description="MoGe-2 model used for depth estimation",
        items=moge2_model_enum_items,
        default=0,
    )
    ref_view_strategy: bpy.props.EnumProperty(
        name="Ref View Strategy",
        description="DA3 reference view selection strategy",
        items=(
            ("first", "First", "Always use the first view"),
            ("middle", "Middle", "Use the middle view"),
            (
                "saddle_balanced",
                "Saddle Balanced",
                "Select a view with balanced feature metrics",
            ),
            (
                "saddle_sim_range",
                "Saddle Sim Range",
                "Select the view with the largest similarity range",
            ),
        ),
        default="saddle_balanced",
    )