bizyengine 0.4.2__py3-none-any.whl

bizyengine/core/commands/base.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+class Command(ABC):
+    @abstractmethod
+    def execute(self):
+        raise NotImplementedError("Subclasses should implement this!")
+
+
+class Processor(ABC):
+    @abstractmethod
+    def process(
+        self,
+        url: str,
+        prompt,
+    ) -> Any:
+        pass
+
+    @abstractmethod
+    def validate_input(self, input: Any) -> bool:
+        pass
+
+    def __call__(self, *args: Any, **kwds: Any) -> Any:
+        if not self.validate_input(*args, **kwds):
+            raise ValueError(f"Invalid input {args=} {kwds=}")
+        return self.process(*args, **kwds)

bizyengine/core/commands/invoker.py

@@ -0,0 +1,4 @@
+from .processors.prompt_processor import PromptProcessor, SearchServiceRouter
+from .servers.prompt_server import PromptServer
+
+prompt_server = PromptServer(router=SearchServiceRouter(), processor=PromptProcessor())

bizyengine/core/commands/processors/prompt_processor.py

@@ -0,0 +1,123 @@
+import json
+from collections import deque
+from typing import Any, Dict, List
+
+from bizyengine.core.commands.base import Processor  # type: ignore
+from bizyengine.core.common import client, get_api_key
+from bizyengine.core.common.caching import BizyAirTaskCache, CacheConfig
+from bizyengine.core.common.env_var import (
+    BIZYAIR_DEBUG,
+    BIZYAIR_DEV_REQUEST_URL,
+    BIZYAIR_SERVER_ADDRESS,
+)
+from bizyengine.core.configs.conf import ModelRule, config_manager
+from bizyengine.core.path_utils import (
+    convert_prompt_label_path_to_real_path,
+    guess_url_from_node,
+)
+
+
+def is_link(obj):
+    if not isinstance(obj, list):
+        return False
+    if len(obj) != 2:
+        return False
+    if not isinstance(obj[0], str):
+        return False
+    if not isinstance(obj[1], int) and not isinstance(obj[1], float):
+        return False
+    return True
+
+
+from dataclasses import dataclass
+
+
+class SearchServiceRouter(Processor):
+    def process(self, prompt: Dict[str, Dict[str, Any]], last_node_ids: List[str]):
+        if BIZYAIR_DEV_REQUEST_URL:
+            return BIZYAIR_DEV_REQUEST_URL
+
+        # TODO Improve distribution logic
+        queue = deque(last_node_ids)
+        visited = {key: True for key in last_node_ids}
+        results: List[ModelRule] = []
+        class_type_table = {
+            node_data["class_type"]: True for node_data in prompt.values()
+        }
+
+        while queue:
+            vertex = queue.popleft()
+            if BIZYAIR_DEBUG:
+                print(vertex, end="->")
+
+            rules = guess_url_from_node(prompt[vertex], class_type_table)
+            if rules:
+                results.extend(rules)
+            for _, in_data in prompt[vertex].get("inputs", {}).items():
+                if is_link(in_data):
+                    neighbor = in_data[0]
+                    if neighbor not in visited:
+                        visited[neighbor] = True
+                        queue.append(neighbor)
+
+        base_model, out_route, out_score = None, None, 0
+        for rule in results[::-1]:
+            # TODO add to config models.yaml
+            if rule.mode_type in {"unet", "vae", "checkpoint", "upscale_models"}:
+                base_model = rule.base_model
+                out_route = rule.route
+                out_score = rule.score
+                break
+
+        for rule in results:
+            if base_model is None:
+                if rule.score > out_score:
+                    out_route, out_score = rule.route, rule.score
+            if rule.base_model == base_model:
+                if rule.score > out_score:
+                    out_route, out_score = rule.route, rule.score
+        assert (
+            out_route is not None
+        ), "Failed to find out_route, please check your prompt"
+        return f"{BIZYAIR_SERVER_ADDRESS}{out_route}"
+
+    def validate_input(
+        self, prompt: Dict[str, Dict[str, Any]], last_node_ids: List[str]
+    ):
+        assert len(last_node_ids) == 1
+        return True
+
+
+class PromptProcessor(Processor):
+    def _exec_info(self, prompt: Dict[str, Dict[str, Any]]):
+        exec_info = {
+            "model_version_ids": [],
+            "api_key": get_api_key(),
+        }
+
+        model_version_id_prefix = config_manager.get_model_version_id_prefix()
+        for node_id, node_data in prompt.items():
+            for k, v in node_data.get("inputs", {}).items():
+                if isinstance(v, str) and v.startswith(model_version_id_prefix):
+                    model_version_id = int(v[len(model_version_id_prefix) :])
+                    exec_info["model_version_ids"].append(model_version_id)
+        return exec_info
+
+    def process(
+        self, url: str, prompt: Dict[str, Dict[str, Any]], last_node_ids: List[str]
+    ):
+        return client.send_request(
+            url=url,
+            data=json.dumps(
+                {
+                    "prompt": prompt,
+                    "last_node_id": last_node_ids[0],
+                    "exec_info": self._exec_info(prompt),
+                }
+            ).encode("utf-8"),
+        )
+
+    def validate_input(
+        self, url: str, prompt: Dict[str, Dict[str, Any]], last_node_ids: List[str]
+    ):
+        return True

bizyengine/core/commands/servers/prompt_server.py

@@ -0,0 +1,234 @@
+import hashlib
+import json
+import pprint
+import time
+import traceback
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+import comfy
+from bizyengine.core.commands.base import Command, Processor  # type: ignore
+from bizyengine.core.common.caching import BizyAirTaskCache, CacheConfig
+from bizyengine.core.common.client import send_request
+from bizyengine.core.common.env_var import (
+    BIZYAIR_DEBUG,
+    BIZYAIR_DEV_GET_TASK_RESULT_SERVER,
+    BIZYAIR_SERVER_ADDRESS,
+)
+from bizyengine.core.common.utils import truncate_long_strings
+from bizyengine.core.configs.conf import config_manager
+from bizyengine.core.image_utils import decode_data, encode_data
+
+
+def get_task_result(task_id: str, offset: int = 0) -> dict:
+    """
+    Get the result of a task.
+    """
+    import requests
+
+    task_api = config_manager.get_task_api()
+    if BIZYAIR_DEV_GET_TASK_RESULT_SERVER:
+        url = f"{BIZYAIR_DEV_GET_TASK_RESULT_SERVER}{task_api.task_result_endpoint}/{task_id}"
+    else:
+        url = f"{BIZYAIR_SERVER_ADDRESS}{task_api.task_result_endpoint}/{task_id}"
+
+    if BIZYAIR_DEBUG:
+        print(f"Debug: get task result url: {url}")
+    response_json = send_request(
+        method="GET", url=url, data=json.dumps({"offset": offset}).encode("utf-8")
+    )
+    out = response_json
+    events = out.get("data", {}).get("events", [])
+    new_events = []
+    for event in events:
+        if (
+            "data" in event
+            and isinstance(event["data"], str)
+            and event["data"].startswith("https://")
+        ):
+            # event["data"] = requests.get(event["data"]).json()
+            event["data"] = send_request(method="GET", url=event["data"])
+        new_events.append(event)
+    out["data"]["events"] = new_events
+    return out
+
+
+@dataclass
+class BizyAirTask:
+    TASK_DATA_STATUS = ["PENDING", "PROCESSING", "COMPLETED"]
+    task_id: str
+    data_pool: list[dict] = field(default_factory=list)
+    data_status: str = None
+
+    @staticmethod
+    def check_inputs(inputs: dict) -> bool:
+        return (
+            inputs.get("code") == 20000
+            and inputs.get("status", False)
+            and "task_id" in inputs.get("data", {})
+        )
+
+    @classmethod
+    def from_data(cls, inputs: dict, check_inputs: bool = True) -> "BizyAirTask":
+        if check_inputs and not cls.check_inputs(inputs):
+            raise ValueError(f"Invalid inputs: {inputs}")
+        data = inputs.get("data", {})
+        task_id = data.get("task_id", "")
+        return cls(task_id=task_id, data_pool=[], data_status="started")
+
+    def is_finished(self) -> bool:
+        if not self.data_pool:
+            return False
+        if self.data_pool[-1].get("data_status") == self.TASK_DATA_STATUS[-1]:
+            return True
+        return False
+
+    def send_request(self, offset: int = 0) -> dict:
+        if offset >= len(self.data_pool):
+            return get_task_result(self.task_id, offset)
+        else:
+            return self.data_pool[offset]
+
+    def get_data(self, offset: int = 0) -> dict:
+        if offset >= len(self.data_pool):
+            return {}
+        return self.data_pool[offset]
+
+    @staticmethod
+    def _fetch_remote_data(url: str) -> dict:
+        import requests
+
+        return requests.get(url).json()
+
+    def get_last_data(self) -> dict:
+        return self.get_data(len(self.data_pool) - 1)
+
+    def do_task_until_completed(
+        self, *, timeout: int = 600, poll_interval: float = 1
+    ) -> list[dict]:
+        offset = 0
+        start_time = time.time()
+        pbar = None
+        while not self.is_finished():
+            try:
+                data = self.send_request(offset)
+                data_lst = self._extract_data_list(data)
+                self.data_pool.extend(data_lst)
+                offset += len(data_lst)
+                for data in data_lst:
+                    message = data.get("data", {}).get("message", {})
+                    if (
+                        isinstance(message, dict)
+                        and message.get("event", None) == "progress"
+                    ):
+                        value = message["data"]["value"]
+                        total = message["data"]["max"]
+                        if pbar is None:
+                            pbar = comfy.utils.ProgressBar(total)
+                        pbar.update_absolute(value + 1, total, None)
+            except Exception as e:
+                print(f"Exception: {e}")
+
+            if time.time() - start_time > timeout:
+                raise TimeoutError(f"Timeout waiting for task {self.task_id} to finish")
+
+            time.sleep(poll_interval)
+
+        return self.data_pool
+
+    def _extract_data_list(self, data):
+        data_lst = data.get("data", {}).get("events", [])
+        if not data_lst:
+            raise ValueError(f"No data found in task {self.task_id}")
+        return data_lst
+
+
+class PromptServer(Command):
+    cache_manager: BizyAirTaskCache = BizyAirTaskCache(
+        config=CacheConfig.from_config(config_manager.get_cache_config())
+    )
+
+    def __init__(self, router: Processor, processor: Processor):
+        self.router = router
+        self.processor = processor
+
+    def get_task_id(self, result: Dict[str, Any]) -> str:
+        return result.get("data", {}).get("task_id", "")
+
+    def is_async_task(self, result: Dict[str, Any]) -> str:
+        """Determine if the result indicates an asynchronous task."""
+        return (
+            result.get("code") == 20000
+            and result.get("status", False)
+            and "task_id" in result.get("data", {})
+        )
+
+    def _get_result(self, result: Dict[str, Any], *, cache_key: str = None):
+        try:
+            response_data = result["data"]
+            if BizyAirTask.check_inputs(result):
+                self.cache_manager.set(cache_key, result)
+                bz_task = BizyAirTask.from_data(result, check_inputs=False)
+                bz_task.do_task_until_completed(timeout=10 * 60)  # 10 minutes
+                last_data = bz_task.get_last_data()
+                response_data = last_data.get("data")
+            out = response_data["payload"]
+            assert out is not None, "Output payload should not be None"
+            self.cache_manager.set(cache_key, out, overwrite=True)
+            return out
+        except Exception as e:
+            self.cache_manager.delete(cache_key)
+            raise RuntimeError(f"Exception: {e}, response_data: {response_data}") from e
+
+    def execute(
+        self,
+        prompt: Dict[str, Dict[str, Any]],
+        last_node_ids: List[str],
+        *args,
+        **kwargs,
+    ):
+
+        prompt = encode_data(prompt)
+
+        if BIZYAIR_DEBUG:
+            debug_info = {
+                "prompt": truncate_long_strings(prompt, 50),
+                "last_node_ids": last_node_ids,
+            }
+            pprint.pprint(debug_info, indent=4)
+
+        url = self.router(prompt=prompt, last_node_ids=last_node_ids)
+
+        if BIZYAIR_DEBUG:
+            print(f"Generated URL: {url}")
+
+        start_time = time.time()
+        sh256 = hashlib.sha256(
+            json.dumps({"url": url, "prompt": prompt}).encode("utf-8")
+        ).hexdigest()
+        end_time = time.time()
+        if BIZYAIR_DEBUG:
+            print(
+                f"Time taken to generate sh256-{sh256}: {end_time - start_time} seconds"
+            )
+
+        cached_output = self.cache_manager.get(sh256)
+        if cached_output:
+            if BIZYAIR_DEBUG:
+                print(f"Cache hit for sh256-{sh256}")
+            out = cached_output
+        else:
+            result = self.processor(url, prompt=prompt, last_node_ids=last_node_ids)
+            out = self._get_result(result, cache_key=sh256)
+
+        if BIZYAIR_DEBUG:
+            pprint.pprint({"out": truncate_long_strings(out, 50)}, indent=4)
+
+        try:
+            real_out = decode_data(out)
+            return [x[0] for x in real_out]
+        except Exception as e:
+            print("Exception occurred while decoding data")
+            self.cache_manager.delete(sh256)
+            traceback.print_exc()
+            raise RuntimeError(f"Exception: {e=}") from e

bizyengine/core/common/__init__.py

@@ -0,0 +1,8 @@
+from .client import (
+    fetch_models_by_type,
+    get_api_key,
+    send_request,
+    set_api_key,
+    validate_api_key,
+)
+from .env_var import create_api_key_file, load_api_key

bizyengine/core/common/caching.py

@@ -0,0 +1,198 @@
+import glob
+import json
+import os
+import time
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, Dict
+
+
+@dataclass
+class CacheConfig:
+    max_size: int = 100
+    expiration: int = 300  # 300 seconds
+    cache_dir: str = "./cache"
+    file_prefix: str = "bizyair_cache_"
+    file_suffix: str = ".json"
+    use_cache: bool = True
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]):
+        return cls(
+            max_size=config.get("max_size", 100),
+            expiration=config.get("expiration", 300),
+            cache_dir=config.get("cache_dir", "./cache"),
+            file_prefix=config.get("file_prefix", "bizyair_cache_"),
+            file_suffix=config.get("file_suffix", ".json"),
+            use_cache=config.get("use_cache", True),
+        )
+
+
+class CacheManager(ABC):
+    @abstractmethod
+    def get(self, key):
+        pass
+
+    @abstractmethod
+    def set(self, key, value):
+        pass
+
+    @abstractmethod
+    def clear(self):
+        pass
+
+    @abstractmethod
+    def disable(self):
+        pass
+
+
+class BizyAirTaskCache(CacheManager):
+    def __init__(self, config: CacheConfig):
+        self.config = config
+        self.cache = OrderedDict()
+        self.cache_dir = config.cache_dir
+        self.ensure_directory_exists()
+        self.cache = self.load_cache() if config.use_cache else self.cache
+
+    def ensure_directory_exists(self):
+        if not os.path.exists(self.cache_dir):
+            os.makedirs(self.cache_dir)
+
+    def load_cache(self):
+        cache_v_files = glob.glob(
+            os.path.join(
+                self.cache_dir, f"{self.config.file_prefix}*{self.config.file_suffix}"
+            )
+        )
+        output = OrderedDict()
+        cache_datas = []
+        for cache_file in cache_v_files:
+            try:
+                file_name = os.path.basename(cache_file)[
+                    len(self.config.file_prefix) : -len(self.config.file_suffix)
+                ]
+                cache_key = file_name.split("-")[0]
+                cache_timestamp = file_name.split("-")[1]
+                if int(time.time()) - int(cache_timestamp) > self.config.expiration:
+                    self.delete_file(cache_file)
+                    continue
+                cache_datas.append(
+                    {
+                        "key": cache_key,
+                        "timestamp": int(cache_timestamp),
+                        "file_path": cache_file,
+                    }
+                )
+            except Exception as e:
+                print(
+                    f"Warning: Error loading cache file {cache_file}: because {e}, will delete it"
+                )
+        cache_datas = sorted(cache_datas, key=lambda x: x["timestamp"])
+        for cache_data in cache_datas:
+            output[cache_data["key"]] = (
+                cache_data["file_path"],
+                cache_data["timestamp"],
+            )
+        return output
+
+    def delete(self, key):
+        if key in self.cache:
+            self.delete_file(self.cache[key][0])
+            del self.cache[key]
+
+    def get(self, key):
+        if key not in self.cache:
+            return None
+
+        file_path, timestamp = self.cache[key]
+        if time.time() - timestamp >= self.config.expiration:
+            self._remove_expired_entry(file_path, key)
+            return None
+
+        cache_data = self._read_file(file_path)
+        if cache_data["cache_key"] == key:
+            return cache_data["result"]
+        else:
+            self._remove_expired_entry(file_path, key)
+            return None
+
+    def _read_file(self, file_path):
+        try:
+            with open(file_path, "r") as f:
+                cache_data = json.load(f)
+                return cache_data
+        except Exception as e:
+            print(f"Error reading file {file_path}: {e}")
+        return None
+
+    def _remove_expired_entry(self, file_path, key):
+        self.delete_file(file_path)
+        del self.cache[key]
+
+    def set(self, key, value, *, overwrite=False):
+        if not overwrite and key in self.cache:
+            raise ValueError(
+                f"Key '{key}' already exists in cache. Use overwrite=True to replace it."
+            )
+        assert isinstance(key, str), "Key must be a string"
+
+        if len(self.cache) >= self.config.max_size:
+            self._evict_oldest()
+
+        timestamp = int(time.time())
+        file_path = os.path.join(
+            self.cache_dir, f"{self.config.file_prefix}{key}-{timestamp}.json"
+        )
+        self.write_file(key, value, file_path, timestamp)
+
+    def _evict_oldest(self):
+        oldest_key, (oldest_file_path, _) = self.cache.popitem(last=False)
+        self.delete_file(oldest_file_path)
+
+    def write_file(self, key: str, value: Any, file_path: str, timestamp: int):
+        try:
+            with open(file_path, "w") as f:
+                json.dump(
+                    {"result": value, "cache_key": key, "timestamp": timestamp}, f
+                )
+            self.cache[key] = (file_path, timestamp)
+        except Exception as e:
+            print(f"Error writing file for key '{key}': {e}")
+
+    def delete_file(self, file_path):
+        if os.path.exists(file_path):
+            try:
+                os.remove(file_path)
+            except Exception as e:
+                print(f"Error deleting file '{file_path}': {e}")
+
+    def clear(self):
+        for file_path, _ in self.cache.values():
+            self.delete_file(file_path)
+        self.cache.clear()
+
+    def disable(self):
+        self.clear()
+
+
+# Example usage
+if __name__ == "__main__":
+    cache_config = CacheConfig(max_size=12, expiration=10, cache_dir="./cache")
+    cache = BizyAirTaskCache(cache_config)
+
+    # Set some cache values
+    cache.set("key1", "This is the value for key1")
+    cache.set("key2", "This is the value for key2")
+
+    # Retrieve values from cache
+    print(cache.get("key1"))  # Output: This is the value for key1
+    print(cache.get("key2"))  # Output: This is the value for key2
+
+    # Wait for expiration
+    time.sleep(9)
+    print(cache.get("key1"))  # Output: None (expired)
+
+    # Clear cache
+    cache.clear()
+    print(cache.get("key2"))  # Output: None (cache cleared)