ai_collections/self_created/Tools/comfyUI_Workflow_switch.py

"""
title: Open-WebUI Configuration Updater
description: Updates image generation settings in Open-WebUI via API
author: Pakobbix
author_url: zephyre.one
github:
funding_url:
version: 0.2.0
license: MIT
"""

import asyncio
import requests
import json
from typing import Callable, Any
from pydantic import BaseModel, Field


class EventEmitter:
    def __init__(self, event_emitter: Callable[[dict], Any] = None):
        self.event_emitter = event_emitter

    async def emit(self, description="Unknown State", status="in_progress", done=False):
        if self.event_emitter:
            await self.event_emitter(
                {
                    "type": "status",
                    "data": {
                        "status": status,
                        "description": description,
                        "done": done,
                    },
                }
            )


class Mechanics:
    """Handles API communication for Open-WebUI."""

    @staticmethod
    def get_current_config(open_webui_ip: str, api_key: str) -> dict:
        """Fetches the current configuration from Open-WebUI."""
        api_url = f"http://{open_webui_ip}/api/v1/images/config"
        headers = {"Authorization": f"Bearer {api_key}"}

        response = requests.get(api_url, headers=headers, timeout=10)
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(
                f"Failed to fetch configuration: {response.status_code} {response.text}"
            )

    @staticmethod
    def update_config(open_webui_ip: str, api_key: str, payload: dict) -> str:
        """Updates the configuration in Open-WebUI."""
        api_url = f"http://{open_webui_ip}/api/v1/images/config/update"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

        response = requests.post(api_url, json=payload, headers=headers, timeout=15)
        if response.status_code == 200:
            return "Configuration updated successfully!"
        else:
            raise Exception(
                f"Failed to update configuration: {response.status_code} {response.text}"
            )


class Tools:
    class Valves(BaseModel):
        open_webui_ip: str = Field(
            "localhost",
            description="IP address of Open-WebUI. Default is localhost",
        )
        open_webui_api_key: str = Field(
            "Not-Required", description="API key for Open-WebUI authentication."
        )
        comfy_url: str = Field(
            "http://localhost:8188",
            description="Base URL of ComfyUI. Default is http://localhost:8188",
        )
        comfy_key: str = Field("Not-Required", description="API key for ComfyUI.")
        StableDiffusion_1_5: str = Field(
            '""', description="StableDiffusion 1.5 configuration (JSON string)"
        )
        StableDiffusion_XL: str = Field(
            '""', description="StableDiffusion XL configuration (JSON string)"
        )
        FLUX: str = Field('""', description="FLUX configuration (JSON string)")
        LTXV: str = Field('""', description="LTXV configuration (JSON string)")
        HiDream: str = Field('""', description="HiDream configuration (JSON string)")
        WAN_2_1: str = Field('""', description="WAN2.1 configuration (JSON string)")
        FramePack: str = Field(
            '""', description="FramePack configuration (JSON string)"
        )
        SkyReelsV2: str = Field(
            '""', description="SkyReelsV2 configuration (JSON string)"
        )
        nodes: str = Field(
            "{}", description="ComfyUI workflow nodes configuration (JSON array)"
        )

    def __init__(self):
        self.valves = self.Valves()

    async def switch_workflow(
        self, workflow_name: str, __event_emitter__: Callable[[dict], Any] = None
    ):
        """Switches the workflow to the specified one."""
        event_emitter = EventEmitter(__event_emitter__)

        try:
            await event_emitter.emit("Fetching current configuration...")
            open_webui_ip = self.valves.open_webui_ip
            open_webui_api_key = self.valves.open_webui_api_key

            if not open_webui_api_key.strip():
                raise ValueError("Missing Open-WebUI API key in tool settings.")

            # Fetch current configuration
            current_config = Mechanics.get_current_config(
                open_webui_ip, open_webui_api_key
            )

            # Get the workflow configuration
            workflow_config = getattr(self.valves, workflow_name, "").strip()
            if not workflow_config:
                raise ValueError(
                    f"The workflow '{workflow_name}' is empty or not configured."
                )

            # Update the configuration payload
            payload = current_config
            payload["comfyui"]["COMFYUI_WORKFLOW"] = workflow_config
            payload["comfyui"]["COMFYUI_WORKFLOW_NODES"] = (
                json.loads(self.valves.nodes) if self.valves.nodes.strip() else [{}]
            )

            await event_emitter.emit("Updating configuration...")
            result = Mechanics.update_config(open_webui_ip, open_webui_api_key, payload)

            await event_emitter.emit(result, status="success", done=True)
            return result

        except Exception as e:
            error_message = f"Error: {str(e)}"
            await event_emitter.emit(error_message, status="failed", done=True)
            return error_message


if __name__ == "__main__":
    tools = Tools()
    workflow_name = "FLUX"  # Example workflow name
    print(asyncio.run(tools.switch_workflow(workflow_name)))