bizydraft 0.2.31__py3-none-any.whl → 0.2.78.dev20251117024007__py3-none-any.whl

bizydraft/env.py

@@ -9,3 +9,6 @@ logger.info(f"{BIZYDRAFT_DOMAIN=} {BIZYDRAFT_SERVER=}")
 
 BIZYAIR_API_KEY = os.getenv("BIZYAIR_API_KEY")
 logger.info(f"{BIZYAIR_API_KEY=}")
+
+COMFYAGENT_NODE_CONFIG = os.getenv("COMFYAGENT_NODE_CONFIG", "")
+logger.info(f"{COMFYAGENT_NODE_CONFIG=}")

bizydraft/hijack_nodes.py

@@ -1,7 +1,10 @@
 import re
+from datetime import datetime
 
 from loguru import logger
 
+from bizydraft.env import COMFYAGENT_NODE_CONFIG
+
 try:
     from comfy_extras.nodes_video import LoadVideo
     from nodes import NODE_CLASS_MAPPINGS, LoadImage
@@ -17,56 +20,88 @@ class BizyDraftLoadVideo(LoadVideo):
         super().__init__(*args, **kwargs)
 
     @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {"file": (["choose your file"], {"video_upload": True})},
-        }
+    def INPUT_TYPES(cls, **kwargs):
+        # 调用父类方法，保持兼容性
+        return super().INPUT_TYPES(**kwargs)
 
     @classmethod
     def VALIDATE_INPUTS(s, *args, **kwargs):
         return True
 
+    @classmethod
+    def validate_inputs(s, *args, **kwargs):
+        # V3 API 使用小写的 validate_inputs
+        return True
+
 
 class BizyDraftLoadImage(LoadImage):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {"image": (["choose your file"], {"image_upload": True})},
-        }
+    def INPUT_TYPES(cls, **kwargs):
+        # 调用父类方法，保持兼容性
+        return super().INPUT_TYPES(**kwargs)
 
     @classmethod
     def VALIDATE_INPUTS(s, *args, **kwargs):
         return True
 
+    @classmethod
+    def validate_inputs(s, *args, **kwargs):
+        # V3 API 使用小写的 validate_inputs
+        return True
+
 
 CLASS_PATCHES = {
-    # "LoadImage": BizyDraftLoadImage,
-    # "LoadVideo": BizyDraftLoadVideo,
+    #     "LoadImage": BizyDraftLoadImage,
+    #     "LoadVideo": BizyDraftLoadVideo,
 }
 
+
+def get_data_load_classes_from_url(config_url):
+    import requests
+
+    # 获取当前时间，精确到分钟
+    current_time = datetime.now().strftime("%Y%m%d%H%M")
+
+    try:
+        config_url = config_url + "?t=" + current_time
+        response = requests.get(config_url)
+        response.raise_for_status()
+        data = response.json()
+        keys_list = []
+        if "weight_load_nodes" in data:
+            keys_list.extend(list(data["weight_load_nodes"].keys()))
+        if "media_load_nodes" in data:
+            keys_list.extend(list(data["media_load_nodes"].keys()))
+
+        return keys_list
+    except Exception as e:
+        logger.error(
+            f"Failed to fetch or comfyagent node config from {config_url}: {e}"
+        )
+        return []
+
+
 DATA_LOAD_CLASSES = [
     "LoadImage",
     "LoadVideo",
     "LoadImageMask",
     "LoadAudio",
     "Load3D",
-    "LayerMask: YoloV8Detect",
-    "Lora Loader Stack (rgthree)",
-    "easy loraNames",
-    "easy loraStack",
-    "Load Lora",
-    "Intrinsic_lora_sampling",
-    "ADE_LoadAnimateDiffModel",
-    "ADE_AnimateDiffLoRALoader",
-    "easy ultralyticsDetectorPipe",
-    "UltralyticsDetectorProvider",
-    "SAMLoader",
-    "easy samLoaderPipe",
+    "VHS_LoadAudioUpload",
+    "VHS_LoadVideo",
 ]
 
+if COMFYAGENT_NODE_CONFIG.startswith("http"):
+    fetched_classes = get_data_load_classes_from_url(COMFYAGENT_NODE_CONFIG)
+    if fetched_classes:
+        DATA_LOAD_CLASSES.extend(fetched_classes)
+        logger.info(f"Fetched additional data load classes: {fetched_classes}")
+    else:
+        logger.warning("No additional data load classes fetched from the URL.")
+
 
 def hijack_nodes():
     def _hijack_node(node_name, new_class):
@@ -82,6 +117,7 @@ def hijack_nodes():
 
     # 通用情况，正则匹配后，打通用patch、替换
     for node_name, base_class in NODE_CLASS_MAPPINGS.items():
+
         regex = r"^(?!BizyAir_)\w+.*Loader.*"
         match = re.match(regex, node_name, re.IGNORECASE)
         if (match and (node_name not in CLASS_PATCHES)) or (
@@ -94,7 +130,10 @@ def hijack_nodes():
 
 def create_patched_class(base_class, validate_inputs_func=None):
     class PatchedClass(base_class):
-        pass
+        @classmethod
+        def validate_inputs(cls, *args, **kwargs):
+            # V3 API
+            return True
 
     if validate_inputs_func:
         PatchedClass.VALIDATE_INPUTS = classmethod(validate_inputs_func)

bizydraft/hijack_routes.py

@@ -1,8 +1,8 @@
 from aiohttp import web
 from loguru import logger
 
-from bizydraft.oss_utils import upload_image
-from bizydraft.patch_handlers import post_prompt, view_image
+from bizydraft.oss_utils import upload_image, upload_mask
+from bizydraft.patch_handlers import post_prompt, view_image, view_video
 
 try:
     from server import PromptServer
@@ -23,10 +23,14 @@ def hijack_routes_pre_add_routes():
             ("/view", "GET"): view_image,
             ("/prompt", "POST"): post_prompt,
             ("/upload/image", "POST"): upload_image,
+            ("/upload/mask", "POST"): upload_mask,
             # /api alias
             ("/api/view", "GET"): view_image,
             ("/api/prompt", "POST"): post_prompt,
             ("/api/upload/image", "POST"): upload_image,
+            ("/api/upload/mask", "POST"): upload_mask,
+            # VHS plugin support
+            ("/api/vhs/viewvideo", "GET"): view_video,
         }
 
         async def middleware_handler(request):
@@ -52,6 +56,8 @@ def hijack_routes_pre_add_routes():
             "/prompt",
             "/view",
             "/upload/image",
+            "/upload/mask",
+            "/vhs/viewvideo",
             "/",
             "/ws",
             "/extensions",

bizydraft/oss_utils.py

@@ -1,7 +1,10 @@
 import base64
 import json
 import os
+import re
+import uuid
 from http.cookies import SimpleCookie
+from pathlib import Path
 from time import time
 from typing import Any, Dict
 
@@ -16,6 +19,8 @@ from werkzeug.utils import secure_filename
 
 from bizydraft.env import BIZYAIR_API_KEY, BIZYDRAFT_SERVER
 
+CLIPSPACE_TO_OSS_MAPPING = {}
+
 private_key_pem = os.getenv(
     "RSA_PRIVATE_KEY",
     """-----BEGIN RSA PRIVATE KEY-----
@@ -59,6 +64,9 @@ def decrypt(encrypted_message):
         if not encrypted_message or not isinstance(encrypted_message, str):
             raise ValueError("无效的加密消息")
 
+        if "v4.public" in encrypted_message:
+            return encrypted_message
+
         private_key = serialization.load_pem_private_key(
             private_key_pem.encode(), password=None, backend=default_backend()
         )
@@ -90,7 +98,13 @@ async def get_upload_token(
     filename: str,
     api_key: str,
 ) -> Dict[str, Any]:
-    url = f"{BIZYDRAFT_SERVER}/upload/token?file_name={filename}&file_type=inputs"
+    from urllib.parse import quote
+
+    # 对文件名进行URL编码，避免特殊字符导致问题
+    encoded_filename = quote(filename, safe="")
+    url = (
+        f"{BIZYDRAFT_SERVER}/upload/token?file_name={encoded_filename}&file_type=inputs"
+    )
 
     headers = {
         "Content-Type": "application/json",
@@ -173,10 +187,15 @@ async def upload_to_oss(post, api_key: str):
     if not (image and image.file):
         return web.Response(status=400)
 
+    original_frontend_filename = image.filename  # 保存前端发送的原始文件名
     filename = image.filename
     if not filename:
         return web.Response(status=400)
 
+    should_clean, filename = clean_filename(filename)
+    if should_clean:
+        filename = f"{uuid.uuid4()}.{filename}"
+
     oss_token = await get_upload_token(filename, api_key)
     result = await upload_filefield_to_oss(image, oss_token)
     if result["status"] != 200:
@@ -189,8 +208,23 @@ async def upload_to_oss(post, api_key: str):
     except Exception as e:
         logger.error(f"Commit file failed: {e}")
         return web.Response(status=500, text=str(e))
+
+    # 将 OSS URL 拆分成 filename 和 subfolder，以便前端正确构建 /api/view 请求
+    # 例如: https://bizyair-prod.oss-cn-shanghai.aliyuncs.com/inputs/20250930/file.png
+    oss_url = result["url"]
+    oss_filename = oss_url.split("/")[-1]  # 获取最后一部分作为文件名
+    oss_subfolder = "/".join(
+        oss_url.split("/")[:-1]
+    )  # 获取除文件名外的部分作为 subfolder
+
+    if original_frontend_filename:
+        CLIPSPACE_TO_OSS_MAPPING[original_frontend_filename] = oss_url
+        logger.info(
+            f"[OSS_MAPPING] Cached mapping: {original_frontend_filename} -> {oss_url}"
+        )
+
     return web.json_response(
-        {"name": result["url"], "subfolder": subfolder, "type": image_upload_type}
+        {"name": oss_filename, "subfolder": oss_subfolder, "type": image_upload_type}
     )
 
 
@@ -225,3 +259,190 @@ async def upload_image(request):
         return web.Response(status=403, text="No validated key found")
     post = await request.post()
     return await upload_to_oss(post, api_key)
+
+
+async def upload_mask(request):
+    """
+    处理 mask editor 上传，将带 alpha 通道的图片上传到 OSS
+    """
+    import io
+    import json
+    import tempfile
+
+    from PIL import Image
+    from PIL.PngImagePlugin import PngInfo
+
+    api_key = get_api_key(request)
+    if not api_key:
+        logger.error("[UPLOAD_MASK] No API key found")
+        return web.Response(status=403, text="No validated key found")
+
+    post = await request.post()
+
+    # 获取上传的 mask 图片
+    mask_image = post.get("image")
+    if not (mask_image and mask_image.file):
+        logger.error("[UPLOAD_MASK] No image provided in request")
+        return web.Response(status=400, text="No image provided")
+
+    # 保存前端发送的原始文件名，用于后续缓存映射
+    original_frontend_filename = mask_image.filename
+
+    # 获取原始图片引用
+    original_ref_str = post.get("original_ref")
+
+    if not original_ref_str:
+        # 如果没有 original_ref，直接上传 mask
+        return await upload_to_oss(post, api_key)
+
+    try:
+        from urllib.parse import unquote
+
+        original_ref = json.loads(original_ref_str)
+        original_filename = original_ref.get("filename")
+        original_subfolder = original_ref.get("subfolder", "")
+
+        if not original_filename:
+            logger.error("[UPLOAD_MASK] No filename in original_ref")
+            return web.Response(status=400, text="No filename in original_ref")
+
+        # 构建完整的 OSS URL（类似 view_image 的逻辑）
+        http_prefix_options = ("http:", "https:")
+
+        if "http" in original_subfolder:
+            # subfolder 中包含 URL 基础路径
+            original_subfolder = original_subfolder[original_subfolder.find("http") :]
+            original_subfolder = unquote(original_subfolder).replace(
+                "https:/", "https://"
+            )
+            original_url = f"{original_subfolder}/{original_filename}"
+        elif original_filename.startswith(http_prefix_options):
+            # filename 本身就是完整 URL
+            original_url = original_filename
+        elif (
+            original_subfolder == "clipspace"
+            and original_filename in CLIPSPACE_TO_OSS_MAPPING
+        ):
+            # 检查缓存：如果是 clipspace 文件且在缓存中，使用缓存的 OSS URL
+            original_url = CLIPSPACE_TO_OSS_MAPPING[original_filename]
+        else:
+            # 不是 OSS URL 格式且不在缓存中，直接上传 mask 图片
+            return await upload_to_oss(post, api_key)
+
+        async with aiohttp.ClientSession() as session:
+            async with session.get(original_url) as resp:
+                if resp.status != 200:
+                    logger.error(
+                        f"[UPLOAD_MASK] Failed to download original image: {resp.status}"
+                    )
+                    return web.Response(
+                        status=502,
+                        text=f"Failed to download original image: {resp.status}",
+                    )
+                original_image_data = await resp.read()
+
+        # 处理图片：应用 alpha 通道
+        with Image.open(io.BytesIO(original_image_data)) as original_pil:
+            # 保存元数据
+            metadata = PngInfo()
+            if hasattr(original_pil, "text"):
+                for key in original_pil.text:
+                    metadata.add_text(key, original_pil.text[key])
+
+            # 转换为 RGBA
+            original_pil = original_pil.convert("RGBA")
+
+            # 读取上传的 mask
+            mask_pil = Image.open(mask_image.file).convert("RGBA")
+
+            # alpha copy - 从 mask 提取 alpha 通道并应用到原图
+            new_alpha = mask_pil.getchannel("A")
+            original_pil.putalpha(new_alpha)
+
+            # 保存到临时文件
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+                tmp_filepath = tmp_file.name
+                original_pil.save(tmp_filepath, compress_level=4, pnginfo=metadata)
+
+        # 准备上传到 OSS
+        filename = f"clipspace-mask-{uuid.uuid4().hex[:8]}.png"
+        # subfolder = post.get("subfolder", "clipspace")
+        image_upload_type = post.get("type", "input")
+
+        try:
+            # 获取上传 token
+            oss_token = await get_upload_token(filename, api_key)
+
+            # 读取临时文件并上传
+            with open(tmp_filepath, "rb") as f:
+                # 创建一个类似 FileField 的对象
+                class FileFieldLike:
+                    def __init__(self, file_obj, filename, content_type):
+                        self.file = file_obj
+                        self.filename = filename
+                        self.content_type = content_type
+
+                file_field = FileFieldLike(f, filename, "image/png")
+                result = await upload_filefield_to_oss(file_field, oss_token)
+
+            if result["status"] != 200:
+                logger.error(f"[UPLOAD_MASK] Upload failed: {result.get('reason', '')}")
+                return web.Response(
+                    status=result["status"], text=result.get("reason", "")
+                )
+
+            # Commit file
+            object_key = oss_token["data"]["file"]["object_key"]
+            await commit_file(object_key, filename, api_key)
+
+            # 将 OSS URL 拆分成 filename 和 subfolder，以便前端正确构建 /api/view 请求
+            oss_url = result["url"]
+            oss_filename = oss_url.split("/")[-1]
+            oss_subfolder = "/".join(oss_url.split("/")[:-1])
+
+            if original_frontend_filename:
+                CLIPSPACE_TO_OSS_MAPPING[original_frontend_filename] = oss_url
+
+            response_data = {
+                "name": oss_filename,
+                "subfolder": oss_subfolder,
+                "type": image_upload_type,
+            }
+            return web.json_response(response_data)
+
+        finally:
+            # 清理临时文件
+            if os.path.exists(tmp_filepath):
+                os.remove(tmp_filepath)
+
+    except Exception as e:
+        logger.error(f"[UPLOAD_MASK] ERROR processing mask upload: {e}", exc_info=True)
+        return web.Response(status=500, text=f"Error processing mask: {str(e)}")
+
+
+def _should_clean(name: str) -> bool:
+    """True -> 包含非白名单字符；False -> 正常
+
+    使用白名单机制：只允许安全字符（中英文、数字、下划线、连字符、点、空格）
+    如果文件名包含白名单之外的字符，则需要清理
+    """
+    if not name:
+        return False
+
+    # 分离文件名和扩展名
+    if "." not in name:
+        return False
+
+    # 白名单：允许中英文、数字、下划线、连字符、点、空格、圆括号
+    safe_pattern = r"^[\w\u4e00-\u9fa5\s\-().]+$"
+
+    return not bool(re.match(safe_pattern, name))
+
+
+def clean_filename(bad: str) -> (bool, str):
+    """对乱码串提取最后扩展名；正常串直接返回原值"""
+    if not _should_clean(bad):
+        return False, bad
+    # 提取最后扩展名（含点）
+    ext = re.search(r"(\.[\w]+)$", bad)
+    return True, ext.group(1) if ext else bad