AnyImage/src/anyimage/panel.py

import bpy

from .runtime import runtime
from .properties import BEN2_MODEL_KEY, model_ready

class ServerPanel(bpy.types.Panel):
    bl_label = "Job Server"
    bl_idname = "ANYIMAGE_PT_server"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "AnyImage"
    bl_order = 0

    def draw(self, _context):
        layout = self.layout
        running = runtime.server_busy()
        if not runtime.environment_ready():
            layout.label(
                text="Environment is not installed",
                icon="ERROR",
            )
            setup_row = layout.row()
            setup_row.enabled = not running
            setup_row.operator(
                "anyimage.setup_ai_environment",
                text="Set Up AI Environment…",
                icon="IMPORT",
            )
            return

        status = runtime.server_status()
        state = status.get("state", "STOPPED")
        icons = {
            "READY": "RADIOBUT_ON",
            "BUSY": "TIME",
            "STARTING": "TIME",
            "ERROR": "ERROR",
            "UNAVAILABLE": "ERROR",
            "STOPPED": "PAUSE",
        }
        row = layout.row(align=True)
        if state in {"READY", "BUSY"}:
            row.label(
                text=(
                    "Running"
                    if state == "READY"
                    else "Running (Busy)"
                ),
                icon=icons[state],
            )
            if state == "READY":
                controls = row.row(align=True)
                controls.operator(
                    "anyimage.restart_server",
                    text="",
                    icon="FILE_REFRESH",
                )
                controls.operator(
                    "anyimage.stop_server",
                    text="",
                    icon="PAUSE",
                )
        elif state == "STOPPED":
            row.operator(
                "anyimage.start_server",
                text="Start",
                icon="PLAY",
            )
        else:
            row.label(
                text=status.get("message", state.title()),
                icon=icons.get(state, "QUESTION"),
            )
            if state == "ERROR":
                row.operator(
                    "anyimage.start_server",
                    text="Start",
                    icon="PLAY",
                )

        if state in {"READY", "BUSY"}:
            cached = (
                status.get("resources", {})
                .get("model_manager", {})
                .get("loaded", {})
            )
            cached_names = [name for name in cached.values() if name]
            cached_row = layout.row(align=True)
            cached_row.label(
                text=(
                    f"Cached: {', '.join(cached_names)}"
                    if cached_names
                    else "Cached: None"
                )
            )
            clear = cached_row.row(align=True)
            clear.enabled = state == "READY"
            clear.operator(
                "anyimage.clear_models",
                text="",
                icon="TRASH",
            )


class DebugPanel(bpy.types.Panel):
    bl_label = "Debug"
    bl_idname = "ANYIMAGE_PT_debug"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "AnyImage"
    bl_options = {"DEFAULT_CLOSED"}
    bl_order = 2

    def draw(self, context):
        layout = self.layout
        settings = context.scene.anyimage_settings

        layout.prop(settings, "device")


def draw_da3_input_settings(layout, settings, model_property):
    layout.prop(settings, model_property)
    layout.prop(settings, "input_path")
    input_path = bpy.path.abspath(settings.input_path)
    if input_path.lower().endswith(
        (
            ".mp4",
            ".mov",
            ".avi",
            ".mkv",
            ".webm",
            ".m4v",
            ".mts",
            ".mpg",
            ".mpeg",
        )
    ):
        layout.prop(settings, "video_frames")
    layout.prop(settings, "process_res")
    layout.prop(settings, "process_res_method")
    layout.prop(settings, "ref_view_strategy")


class Da3OutputDebugPanel(bpy.types.Panel):
    bl_label = "DA3"
    bl_idname = "ANYIMAGE_PT_debug_da3_output"
    bl_parent_id = DebugPanel.bl_idname
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "AnyImage"
    bl_options = {"DEFAULT_CLOSED"}
    bl_order = 1

    def draw(self, context):
        layout = self.layout
        settings = context.scene.anyimage_settings

        draw_da3_input_settings(layout, settings, "model")
        export_row = layout.row()
        export_row.enabled = (
            not runtime.server_busy()
            and model_ready(settings.model)
        )
        export_row.operator(
            "anyimage.debug_da3",
            text="Run DA3 API",
            icon="EXPORT",
        )


class Moge2ApiDebugPanel(bpy.types.Panel):
    bl_label = "MoGe-2"
    bl_idname = "ANYIMAGE_PT_debug_moge2_api"
    bl_parent_id = DebugPanel.bl_idname
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "AnyImage"
    bl_options = {"DEFAULT_CLOSED"}
    bl_order = 0

    def draw(self, context):
        layout = self.layout
        settings = context.scene.anyimage_settings
        layout.prop(settings, "input_path")
        layout.prop(settings, "moge_model")
        layout.prop(settings, "moge_resolution_level")
        run_row = layout.row()
        run_row.enabled = (
            not runtime.server_busy()
            and model_ready(settings.moge_model)
        )
        run_row.operator(
            "anyimage.debug_moge2",
            text="Run MoGe-2 API",
            icon="EXPORT",
        )


class Ben2RemovalDebugPanel(bpy.types.Panel):
    bl_label = "BEN2"
    bl_idname = "ANYIMAGE_PT_debug_ben2_removal"
    bl_parent_id = DebugPanel.bl_idname
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "AnyImage"
    bl_options = {"DEFAULT_CLOSED"}
    bl_order = 2

    def draw(self, context):
        layout = self.layout
        settings = context.scene.anyimage_settings

        layout.prop(settings, "background_input_path")
        background_row = layout.row()
        background_row.enabled = (
            not runtime.server_busy()
            and model_ready(BEN2_MODEL_KEY)
        )
        background_row.operator(
            "anyimage.debug_ben2",
            text="Remove Background",
            icon="IMAGE_ALPHA",
        )