velai/nodes/image_data.py

from __future__ import annotations

from typing import Any

from nicegui import events, ui

from velai import image_utils
from velai.data_types import ImageType
from velai.dataflow.enums import NodeKind, PortDirection
from velai.dataflow.nodes import NodeType
from velai.dataflow.ports import PortSchema
from velai.nodes.actions.node_action_decorator import as_action_name, node_action
from velai.nodes.actions.node_action_models import AsyncNodeActionResultIterator, NodeActionArguments, NodeActionResult
from velai.nodes.base_node import BaseNode, BaseNodeData
from velai.nodes.base_node_renderable import T_BASE_NODE, BaseNodeRenderable
from velai.storage import storage_endpoint

ImageDataNodeType = NodeType(
    kind=NodeKind.IMAGE_DATA,
    display_name="Image Upload",
    inputs=[],
    outputs=[PortSchema(name="image", dtype=ImageType, direction=PortDirection.OUTPUT, tooltip="Image")],
)


class ImageDataNode(BaseNode[BaseNodeData]):
    """Node that exposes a constant image value via its output port.

    The text is stored on the "text" output port and edited through the UI.
    """


class ImageDataNodeRenderable(BaseNodeRenderable[ImageDataNode]):
    def get_header_buttons(self, node: T_BASE_NODE) -> list[dict[str, Any]]:
        return [
            *super().get_header_buttons(node),
            {
                "name": "upload_python",
                "icon": "upload",
                "tooltip": "Upload image",
                "action": as_action_name(self._on_upload_image),
                "disableWhileProcessing": True,
            },
        ]

    def get_fields(self, node: T_BASE_NODE) -> list[dict[str, Any]]:
        return [
            *super().get_fields(node),
            {
                "name": "image",
                "kind": "image",
                "label": "Image",
                "placeholder": "No image",
            },
        ]

    @node_action
    async def _on_upload_image(self, args: NodeActionArguments) -> AsyncNodeActionResultIterator:
        # create the dialog in the main page content context
        with ui.context.client.content:
            with ui.dialog() as dialog:
                with ui.card():
                    ui.label("Upload an image")

                    async def handle_upload(e: events.UploadEventArguments) -> None:
                        file = e.file
                        data = await file.read()

                        # resize image file to max side length, keeping aspect ratio
                        try:
                            image_data = await image_utils.prepare_image_bytes_for_storage_async(data)
                        except ValueError as ex:
                            ui.notification(f"Image upload did not work {ex}", type="negative")
                            dialog.submit(False)
                            return

                        # download and store mesh
                        image_url = await storage_endpoint.store_data(
                            image_data, original_name=file.name, mime_type=e.file.content_type
                        )

                        # update image url
                        args.node.outputs["image"].value = image_url
                        dialog.submit(True)

                    # upload definition
                    (
                        ui.upload(
                            label="Choose an image",
                            on_upload=handle_upload,
                            max_files=1,
                            multiple=False,
                            max_file_size=1024 * 1024 * 1024 * 5,  # max 5MB
                            auto_upload=True,
                        ).props("accept=.jpeg,.jpg,.gif,.png,.tiff,.webp")
                    )

                    ui.button("Cancel", on_click=dialog.close)

            # open and run dialog
            await dialog

        yield NodeActionResult.update_node()