velai/ui/node_progress.py

from __future__ import annotations

from typing import TYPE_CHECKING

from velai.dataflow.enums import NodeState
from velai.nodes.base_node import BaseNode

if TYPE_CHECKING:
    from velai.dataflow.nodes import NodeInstance
    from velai.ui.vueflow_canvas import VueFlowCanvas


class ProgressStateUpdater:
    def __init__(self, canvas: VueFlowCanvas) -> None:
        self.canvas = canvas

    async def update(self, node: NodeInstance) -> None:
        """Update the progress UI based on the current state of the node."""
        node_id = node.node_id
        state = node.state

        if state == NodeState.QUEUED:
            await self.canvas.set_node_processing(node_id, True)
            await self.canvas.update_node_progress(node_id, None, None)

        elif state == NodeState.PROCESSING:
            await self.canvas.set_node_processing(node_id, True)
            if isinstance(node, BaseNode):
                # Only show message if it's not None or empty
                msg = node.data.progress_message
                if not msg:
                    msg = None

                await self.canvas.update_node_progress(
                    node_id,
                    node.data.progress_value,
                    msg,
                )
            else:
                await self.canvas.update_node_progress(node_id, None, None)

        elif state == NodeState.IDLE:
            await self.canvas.set_node_processing(node_id, False)
            await self.canvas.update_node_progress(node_id, None, None)