bizydraft 0.1.30__py3-none-any.whl → 0.2.0__py3-none-any.whl

bizydraft/patch_handlers.py

@@ -0,0 +1,150 @@
+import asyncio
+import math
+import mimetypes
+import os
+import uuid
+from urllib.parse import unquote
+
+from aiohttp import ClientSession, ClientTimeout, web
+from loguru import logger
+
+try:
+    import execution
+    from server import PromptServer
+
+    comfy_server = PromptServer.instance
+except ImportError:
+    logger.error(
+        "failed to import ComfyUI modules, ensure PYTHONPATH is set correctly. (export PYTHONPATH=$PYTHONPATH:/path/to/ComfyUI)"
+    )
+    exit(1)
+
+BIZYDRAFT_MAX_FILE_SIZE = int(
+    os.getenv("BIZYDRAFT_MAX_FILE_SIZE", 100 * 1024 * 1024)
+)  # 100MB
+BIZYDRAFT_REQUEST_TIMEOUT = int(
+    os.getenv("BIZYDRAFT_REQUEST_TIMEOUT", 20 * 60)
+)  # 20分钟
+BIZYDRAFT_CHUNK_SIZE = int(os.getenv("BIZYDRAFT_CHUNK_SIZE", 1024 * 16))  # 16KB
+
+
+async def view_image(request):
+    logger.debug(f"Received request for /view with query: {request.rel_url.query}")
+    if "filename" not in request.rel_url.query:
+        logger.warning("'filename' not provided in query string, returning 404")
+        return web.Response(status=404, text="'filename' not provided in query string")
+
+    filename = request.rel_url.query["filename"]
+    subfolder = request.rel_url.query.get("subfolder", "")
+
+    http_prefix_options = ("http:", "https:")
+
+    if not filename.startswith(http_prefix_options) and "http" not in subfolder:
+        logger.warning(
+            f"Invalid filename format: {filename=}, {subfolder=} only URLs are supported"
+        )
+        return web.Response(
+            status=400, text="Invalid filename format(only url supported)"
+        )
+
+    try:
+        if "http" in subfolder:
+            subfolder = subfolder[subfolder.find("http") :]
+        subfolder = unquote(subfolder)
+        filename = (
+            f"{subfolder}/{filename}"
+            if not filename.startswith(http_prefix_options)
+            else filename
+        )  # preview 3d request: https://host:port/api/view?filename=filename.glb&type=output&subfolder=https://bizyair-dev.oss-cn-shanghai.aliyuncs.com/outputs&rand=0.5763957215362988
+
+        content_type, _ = mimetypes.guess_type(filename)
+        if content_type and any(x in content_type for x in ("image", "video")):
+            return web.HTTPFound(filename)
+
+        timeout = ClientTimeout(total=BIZYDRAFT_REQUEST_TIMEOUT)
+        async with ClientSession(timeout=timeout) as session:
+            async with session.get(filename) as resp:
+                resp.raise_for_status()
+                content_length = int(resp.headers.get("Content-Length", 0))
+                if content_length > BIZYDRAFT_MAX_FILE_SIZE:
+                    logger.warning(
+                        f"File size {human_readable_size(content_length)} exceeds limit {human_readable_size(BIZYDRAFT_MAX_FILE_SIZE)}"
+                    )
+                    return web.Response(
+                        status=413,
+                        text=f"File size exceeds limit ({human_readable_size(BIZYDRAFT_MAX_FILE_SIZE)})",
+                    )
+
+                headers = {
+                    "Content-Disposition": f'attachment; filename="{uuid.uuid4()}"',
+                    "Content-Type": "application/octet-stream",
+                }
+
+                proxy_response = web.StreamResponse(headers=headers)
+                await proxy_response.prepare(request)
+
+                total_bytes = 0
+                async for chunk in resp.content.iter_chunked(BIZYDRAFT_CHUNK_SIZE):
+                    total_bytes += len(chunk)
+                    if total_bytes > BIZYDRAFT_MAX_FILE_SIZE:
+                        await proxy_response.write(b"")
+                        return web.Response(
+                            status=413,
+                            text=f"File size exceeds limit during streaming ({human_readable_size(BIZYDRAFT_MAX_FILE_SIZE)})",
+                        )
+                    await proxy_response.write(chunk)
+
+                return proxy_response
+
+    except asyncio.TimeoutError:
+        return web.Response(
+            status=504,
+            text=f"Request timed out (max {BIZYDRAFT_REQUEST_TIMEOUT//60} minutes)",
+        )
+    except Exception as e:
+        return web.Response(
+            status=502, text=f"Failed to fetch remote resource: {str(e)}"
+        )
+
+
+def human_readable_size(size_bytes):
+    if size_bytes == 0:
+        return "0B"
+    size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
+    i = int(math.floor(math.log(size_bytes, 1024)))
+    p = math.pow(1024, i)
+    s = round(size_bytes / p, 2)
+    return f"{s} {size_name[i]}"
+
+
+async def post_prompt(request):
+    json_data = await request.json()
+
+    json_data = comfy_server.trigger_on_prompt(json_data)
+    if "prompt" in json_data:
+        prompt = json_data["prompt"]
+        valid = execution.validate_prompt(prompt)
+        if valid[0]:
+            response = {
+                "prompt_id": None,
+                "number": None,
+                "node_errors": valid[3],
+            }
+            logger.debug(f"Received POST request to /prompt with valid prompt")
+            return web.json_response(response)
+        else:
+            logger.debug(
+                f"Fail to validate prompt: {valid[1]=}, node_errors: {valid[3]=}"
+            )
+            return web.json_response(
+                {"error": valid[1], "node_errors": valid[3]}, status=400
+            )
+    else:
+        error = {
+            "type": "no_prompt",
+            "message": "No prompt provided",
+            "details": "No prompt provided",
+            "extra_info": {},
+        }
+        logger.debug(f"Received POST request to /prompt with no prompt: {error}")
+        return web.json_response({"error": error, "node_errors": {}}, status=400)

bizydraft/postload.py

@@ -4,21 +4,39 @@ from loguru import logger
 def lazy_hook():
     try:
         import app.database.db
+        import comfy.utils
 
-        origin_init_db = app.database.db.init_db
+        origin_fun_post_add_routes = comfy.utils.set_progress_bar_global_hook
+        origin_fun_pre_add_routes = app.database.db.init_db
 
-        def hijack_all():
+        def hijack_all_pre_add_routes():
+            from bizydraft.hijack_routes import hijack_routes_pre_add_routes
+
+            hijack_routes_pre_add_routes()
+
+        def new_fun_pre_add_routes(*args, **kwargs):
+            hijack_all_pre_add_routes()
+            origin_fun_pre_add_routes(*args, **kwargs)
+
+        def hijack_all_post_add_routes():
             from bizydraft.hijack_nodes import hijack_nodes
-            from bizydraft.hijack_routes import hijack_routes
 
             hijack_nodes()
-            hijack_routes()
 
-        def new_init_db():
-            hijack_all()
-            origin_init_db()
+            from bizydraft.hijack_routes import hijack_routes_post_add_routes
+
+            hijack_routes_post_add_routes()
+
+            from bizydraft.block_nodes import remove_blacklisted_nodes
+
+            remove_blacklisted_nodes()
+
+        def new_fun_post_add_routes(*args, **kwargs):
+            hijack_all_post_add_routes()
+            origin_fun_post_add_routes(*args, **kwargs)
 
-        app.database.db.init_db = new_init_db
+        comfy.utils.set_progress_bar_global_hook = new_fun_post_add_routes
+        app.database.db.init_db = new_fun_pre_add_routes
 
     except Exception as e:
         logger.error(f"failed to lazyhook: {e}")