RiceRound-ComfyUI_CryptoCat/updown_workflow.py

import json
import os
import requests
from .url_config import CatUrlConfig, download_crypto_workflow, user_upload_workflow
from .utils import combine_files, get_machine_id, get_local_app_setting_path
from .auth_unit import AuthUnit
from server import PromptServer
from aiohttp import web
import time


class Cancelled(Exception):
    pass


class MessageHolder:
    stash = {}
    messages = {}
    cancelled = False

    @classmethod
    def addMessage(cls, id, message):
        if message == "__cancel__":
            cls.messages = {}
            cls.cancelled = True
        elif message == "__start__":
            cls.messages = {}
            cls.stash = {}
            cls.cancelled = False
        else:
            cls.messages[str(id)] = message

    @classmethod
    def waitForMessage(cls, id, period=0.1, timeout=60):
        sid = str(id)
        cls.messages.clear()
        start_time = time.time()
        while not sid in cls.messages and not "-1" in cls.messages:
            if cls.cancelled:
                cls.cancelled = False
                raise Cancelled()
            if time.time() - start_time > timeout:
                raise Cancelled("Operation timed out")
            time.sleep(period)
        if cls.cancelled:
            cls.cancelled = False
            raise Cancelled()
        message = cls.messages.pop(str(id), None) or cls.messages.pop("-1")
        return message.strip()


routes = PromptServer.instance.routes


@routes.post("/cryptocat/message")
async def message_handler(request):
    post = await request.post()
    MessageHolder.addMessage(post.get("id"), post.get("message"))
    return web.json_response({})


class UploadWorkflow:
    def __init__(self, user_token):
        self.user_token = user_token

    def check_workflow(self, template_id):
        url = CatUrlConfig().workflow_url
        headers = {"Authorization": f"Bearer {self.user_token}"}
        response = requests.get(
            url, params={"template_id": template_id, "action": "check"}, headers=headers
        )
        if response.status_code == 200:
            response_data = response.json()
            error_code = response_data.get("code")
            error_msg = response_data.get("message")
            return error_code, error_msg
        else:
            return -1, ""

    def upload_workflow(self, template_id, temp_dir):
        error_code, error_msg = self.check_workflow(template_id)
        if error_code == 1:
            auto_overwrite = UserWorkflowSetting().get_auto_overwrite()
            if not auto_overwrite:
                json_content = {
                    "title": "已经存在相同template_id的数据，是否覆盖？",
                    "icon": "info",
                    "confirmButtonText": "覆盖",
                    "cancelButtonText": "取消",
                    "showCancelButton": True,
                    "timer": 50000,
                }
                PromptServer.instance.send_sync(
                    "cryptocat_dialog",
                    {"json_content": json.dumps(json_content), "id": template_id},
                )
                msg_result = MessageHolder.waitForMessage(template_id, timeout=60000)
                try:
                    result_code = int(msg_result)
                except ValueError:
                    print("crypto cat upload cancel: Invalid response format")
                    return False
                if result_code != 1:
                    print("crypto cat upload cancel: User rejected overwrite")
                    return False
        elif error_code != 0:
            print(f"crypto cat upload failed: {error_msg}")
            PromptServer.instance.send_sync(
                "cryptocat_toast", {"content": f"异常情况，{error_msg}", "type": "error"}
            )
            return False
        combine_target_path = os.path.join(temp_dir, f"cryptocat_{template_id}.zip")
        combine_file_paths = []
        for file in [
            f"crypto_{template_id}.json",
            f"original_workflow_{template_id}.json",
            f"original_prompt_{template_id}.json",
        ]:
            src_path = os.path.join(temp_dir, file)
            combine_file_paths.append(src_path)
        combine_files(combine_file_paths, template_id, combine_target_path)
        for file_path in combine_file_paths:
            os.remove(file_path)
        success, message = user_upload_workflow(
            template_id, combine_target_path, self.user_token
        )
        if success:
            PromptServer.instance.send_sync(
                "cryptocat_toast", {"content": "上传成功", "type": "info", "duration": 5000}
            )
        else:
            PromptServer.instance.send_sync(
                "cryptocat_toast",
                {"content": f"上传失败: {message}", "type": "error", "duration": 5000},
            )
        return success

    def generate_serial_number(
        self, template_id, expire_date=None, use_days=None, count=1
    ):
        url = CatUrlConfig().serial_number_url
        headers = {"Authorization": f"Bearer {self.user_token}"}
        params = {"template_id": template_id, "count": count}
        if expire_date is not None:
            params["expire_date"] = expire_date
        if use_days is not None:
            params["use_days"] = use_days
        response = requests.get(url, params=params, headers=headers)
        if response.status_code == 200:
            response_data = response.json()
            error_code = response_data.get("code")
            error_msg = response_data.get("message", "")
            if error_code == 0:
                return response_data.get("data", [])
            else:
                raise ValueError(f"failed to generate serial number: {error_msg}")
        else:
            raise ValueError(
                f"failed to generate serial number: {response.status_code}"
            )


class DownloadWorkflow:
    _instance = None
    _cache = {}

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(DownloadWorkflow, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if not hasattr(self, "_initialized"):
            self._initialized = True

    def download_workflow(self, template_id, serial_number_token, user_token):
        if not serial_number_token:
            raise ValueError("需要输入序列号")
        cache_key = f"{template_id}_{serial_number_token}_{user_token}"
        if cache_key in self._cache:
            return self._cache[cache_key]
        hardware_id = get_machine_id()
        success, message = download_crypto_workflow(
            template_id, hardware_id, serial_number_token, user_token
        )
        if success:
            self._cache[cache_key] = message
            return message
        elif message is not None:
            if message == "need login" or message == "需要登录":
                AuthUnit().login_dialog("运行工作流，请先完成登录")
            else:
                raise ValueError(message)
        else:
            raise ValueError("failed to download workflow")


class UserWorkflowSetting:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(UserWorkflowSetting, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if not hasattr(self, "initialized"):
            local_app_path = get_local_app_setting_path()
            local_app_path.mkdir(parents=True, exist_ok=True)
            self.config_path = local_app_path / "config.ini"
            self.section_name = "Workflow"
            self.initialized = True

    def get_auto_overwrite(self):
        try:
            import configparser

            config = configparser.ConfigParser()
            config.read(self.config_path, encoding="utf-8")
            if not config.has_section(self.section_name):
                return False
            return config.getboolean(
                self.section_name, "auto_overwrite", fallback=False
            )
        except Exception as e:
            print(f"Error reading auto_overwrite setting: {str(e)}")
            return False

    def set_auto_overwrite(self, value):
        try:
            import configparser

            config = configparser.ConfigParser()
            config.read(self.config_path, encoding="utf-8")
            if not config.has_section(self.section_name):
                config.add_section(self.section_name)
            config.set(self.section_name, "auto_overwrite", str(value).lower())
            with open(self.config_path, "w", encoding="utf-8") as configfile:
                config.write(configfile)
            return True
        except Exception as e:
            print(f"Error saving auto_overwrite setting: {str(e)}")
            return False