system prompt

Update extension-node-map.json

.github/workflows/publish.yml

@@ -0,0 +1,21 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+    paths:
+      - "pyproject.toml"
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          ## Add your own personal access token to your Github Repository secrets and reference it here.
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -3,4 +3,6 @@ https/
 nodes/config.json
 workflow/my_workflow.json
 workflow/my_workflow_app.json
-app/*
+workflow/prompt_result.json
+app/*
+workflow/prompt_result.json

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 shadow
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,28 +1,53 @@
-> 适配了最新版comfyui的py3.11 ，torch 2.1.2+cu121
+![](https://img.shields.io/github/release/shadowcz007/comfyui-mixlab-nodes)
 
+> 适配了最新版 comfyui 的 py3.11 ，torch 2.1.2+cu121
 > [Mixlab nodes discord](https://discord.gg/cXs9vZSqeK)
 
-#### 
+
+##### `最新`：
+
+ChatGPT 节点支持 Local LLM（llama.cpp），Phi3、llama3 都可以直接一个节点运行了。模型下载后，放置到 `models/llamafile/`
+
+- 右键菜单支持 text-to-text，方便对 prompt 词补全
+
+强烈推荐：
+[Phi-3-mini-4k-instruct-function-calling-GGUF](https://huggingface.co/nold/Phi-3-mini-4k-instruct-function-calling-GGUF)
+
+[Phi-3-mini-4k-instruct-GGUF](https://huggingface.co/lmstudio-community/Phi-3-mini-4k-instruct-GGUF/tree/main)，备选：[llama3_if_ai_sdpromptmkr_q2k](https://hf-mirror.com/impactframes/llama3_if_ai_sdpromptmkr_q2k/tree/main)
+
+- 右键菜单支持 image-to-text，使用多模态模型，多模态使用 [llava-phi-3-mini-gguf](https://huggingface.co/xtuner/llava-phi-3-mini-gguf/tree/main)，注意需要把llava-phi-3-mini-mmproj-f16.gguf也下载
+
+![](./assets/prompt_ai_setup.png)
+![](./assets/prompt-ai.png)
+
+
+#### `相关插件推荐`
+
+<!-- [comfyui-sd-prompt-mixlab](https://github.com/shadowcz007/comfyui-sd-prompt-mixlab) -->
+
+[comfyui-Image-reward](https://github.com/shadowcz007/comfyui-Image-reward)
+
 [comfyui-ultralytics-yolo](https://github.com/shadowcz007/comfyui-ultralytics-yolo)
 
 [comfyui-moondream](https://github.com/shadowcz007/comfyui-moondream)
 
-[comfyui-CLIPSeg](https://github.com/shadowcz007/comfyui-CLIPSeg)
+<!-- [comfyui-CLIPSeg](https://github.com/shadowcz007/comfyui-CLIPSeg) -->
 
+## 🚀🚗🚚🏃 Workflow-to-APP
 
-## 🚀🚗🚚🏃 Workflow-to-APP 
-- 新增AppInfo节点，可以通过简单的配置，把workflow转变为一个Web APP。
-- 支持多个web app 切换
-- 发布为app的workflow，可以在右键里再次编辑了
-- web app可以设置分类，在comfyui右键菜单可以编辑更新web app
+- 新增 AppInfo 节点，可以通过简单的配置，把 workflow 转变为一个 Web APP。
+- 支持多个 web app 切换
+- 发布为 app 的 workflow，可以在右键里再次编辑了
+- web app 可以设置分类，在 comfyui 右键菜单可以编辑更新 web app
+- 支持动态提示
 
+![](./assets/微信图片_20240421205440.png)
 
 - Support multiple web app switching.
 - Add the AppInfo node, which allows you to transform the workflow into a web app by simple configuration.
 - The workflow, which is now released as an app, can also be edited again by right-clicking.
 - The web app can be configured with categories, and the web app can be edited and updated in the right-click menu of ComfyUI.
 
-
 ![](./assets/0-m-app.png)
 
 ![](./assets/appinfo-readme.png)
@@ -30,54 +55,91 @@
 ![](./assets/appinfo-2.png)
 
 Example:
+
 - workflow
-![APP info](./workflow/appinfo-workflow.svg)
-[text-to-image](./workflow/Text-to-Image-app.json)
+  ![APP info](./workflow/appinfo-workflow.svg)
+  [text-to-image](./workflow/Text-to-Image-app.json)
 
 APP-JSON:
+
 - [text-to-image](./example/Text-to-Image_3.json)
 - [image-to-image](./example/Image-to-Image_2.json)
 - text-to-text
 
-> 暂时支持 9 种节点作为界面上的输入节点：Load Image、VHS_LoadVideo、CLIPTextEncode、PromptSlide、TextInput_、Color、FloatSlider、IntNumber、CheckpointLoaderSimple、LoraLoader
+> 暂时支持 9 种节点作为界面上的输入节点：Load Image、VHS*LoadVideo、CLIPTextEncode、PromptSlide、TextInput*、Color、FloatSlider、IntNumber、CheckpointLoaderSimple、LoraLoader
 
 > 输出节点：PreviewImage 、SaveImage、ShowTextForGPT、VHS_VideoCombine、PromptImage
 
-> seed统一输入控件，支持：SamplerCustom、KSampler
+> seed 统一输入控件，支持：SamplerCustom、KSampler
 
-> [ps插件](https://github.com/shadowcz007/comfyui-ps-plugin)
+> 配套[ps 插件](https://github.com/shadowcz007/comfyui-ps-plugin)
 
+> 如果遇到上传图片不成功，请检查下：局域网或者是云服务，请使用 https，端口 8189 这个服务（ 感谢 @Damien 反馈问题）
+
+> If you encounter difficulties in uploading images, please check the following: for local network or cloud services, please use HTTPS and the service on port 8189. (Thanks to @Damien for reporting the issue.)
+
+## 🏃🚗🚚🚀 Real-time Design
 
-## 🏃🚗🚚🚀  Real-time Design
 > ScreenShareNode & FloatingVideoNode. Now comfyui supports capturing screen pixel streams from any software and can be used for LCM-Lora integration. Let's get started with implementation and design! 💻🌐
 
 ![screenshare](./assets/screenshare.png)
 
 https://github.com/shadowcz007/comfyui-mixlab-nodes/assets/12645064/e7e77f90-e43e-410a-ab3a-1952b7b4e7da
 
-
 <!-- [ScreenShareNode](./workflow/2-screeshare.json) -->
+
 [ScreenShareNode & FloatingVideoNode](./workflow/3-FloatVideo-workflow.json)
 
 !! Please use the address with HTTPS (https://127.0.0.1).
 
-
 ### SpeechRecognition & SpeechSynthesis
+
 ![f](./assets/audio-workflow.svg)
 
 [Voice + Real-time Face Swap Workflow](./workflow/语音+实时换脸workflow.json)
 
 ### GPT
-> Support for calling multiple GPTs.ChatGPT、ChatGLM3 , Some code provided by rui. If you are using OpenAI's service, fill in https://api.openai.com/v1 . If you are using a local LLM service, fill in http://127.0.0.1:xxxx/v1 .  Azure OpenAI:https://xxxx.openai.azure.com 
+
+> Support for calling multiple GPTs.Local LLM（llama.cpp）、 ChatGPT、ChatGLM3 、ChatGLM4 , Some code provided by rui. If you are using OpenAI's service, fill in https://api.openai.com/v1 . If you are using a local LLM service, fill in http://127.0.0.1:xxxx/v1 . Azure OpenAI:https://xxxx.openai.azure.com
 
 ![gpt-workflow.svg](./assets/gpt-workflow.svg)
 
 [workflow-5](./workflow/5-gpt-workflow.json)
 
+最新：ChatGPT 节点支持 Local LLM（llama.cpp），Phi3、llama3 都可以直接一个节点运行了。
+
+Model download,move to :`models/llamafile/`
+
+强烈推荐：[Phi-3-mini-4k-instruct-GGUF](https://huggingface.co/lmstudio-community/Phi-3-mini-4k-instruct-GGUF/tree/main)
+
+备选：[llama3_if_ai_sdpromptmkr_q2k](https://hf-mirror.com/impactframes/llama3_if_ai_sdpromptmkr_q2k/tree/main)
+
+> 如果碰到安装失败，可以尝试手动安装
+
+```
+../../../python_embeded/python.exe -s -m pip install llama-cpp-python --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121
+
+../../../python_embeded/python.exe -s -m pip install llama-cpp-python[server]
+
+```
+
+> [Mac](https://llama-cpp-python.readthedocs.io/en/latest/install/macos/)
+
+```
+pip uninstall llama-cpp-python -y
+CMAKE_ARGS="-DLLAMA_METAL=on" pip install -U llama-cpp-python --no-cache-dir
+pip install 'llama-cpp-python[server]'
+```
+
+```
+pip install llama-cpp-python \
+  --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/metal
+```
 
 ## Prompt
+
 > PromptSlide
-![](./assets/prompt_weight.png)
+> ![](./assets/prompt_weight.png)
 
 <!-- ![](./workflow/promptslide-appinfo-workflow.svg) -->
 
@@ -91,81 +153,93 @@ https://github.com/shadowcz007/comfyui-mixlab-nodes/assets/12645064/e7e77f90-e43
 
 > PromptImage & PromptSimplification,Assist in simplifying prompt words, comparing images and prompt word nodes.
 
-> ChinesePrompt && PromptGenerate，中文prompt节点，直接用中文书写你的prompt
+> ChinesePrompt && PromptGenerate，中文 prompt 节点，直接用中文书写你的 prompt
 
 ![](./assets/ChinesePrompt_workflow.svg)
 
-
 ### Layers
+
 > A new layer class node has been added, allowing you to separate the image into layers. After merging the images, you can input the controlnet for further processing.
 
 ![layers](./assets/layers-workflow.svg)
 
 ![poster](./assets/poster-workflow.svg)
 
-
 ### 3D
+
+![](./assets/3d-workflow.png)
+![](./assets/3d_app.png)
+[workflow](./assets/Image-to-3D_1.json)
+
 ![](./assets/3dimage.png)
 [workflow](./workflow/3D-workflow.json)
 
+### Image
 
-### LoadImagesFromLocal
-> Monitor changes to images in a local folder, and trigger real-time execution of workflows, supporting common image formats, especially PSD format, in conjunction with Photoshop. 
+#### LoadImagesToBatch
+
+> Upload multiple images for batch input into the IP adapter.
+
+#### LoadImagesFromLocal
+
+> Monitor changes to images in a local folder, and trigger real-time execution of workflows, supporting common image formats, especially PSD format, in conjunction with Photoshop.
 
 ![watch](./assets/4-loadfromlocal-watcher-workflow.svg)
 
 [workflow-4](./workflow/4-loadfromlocal-watcher-workflow.json)
 
-### LoadImagesFromURL
+#### LoadImagesFromURL
+
 > Conveniently load images from a fixed address on the internet to ensure that default images in the workflow can be executed.
 
+### Style
+
+> Apply VisualStyle Prompting , Modified from [ComfyUI_VisualStylePrompting](https://github.com/ExponentialML/ComfyUI_VisualStylePrompting)
+
+![](./assets/VisualStylePrompting.png)
+
+> StyleAligned , Modified from [style_aligned_comfy](https://github.com/brianfitzgerald/style_aligned_comfy)
+
+### Utils
 
-## Utils
 > The Color node provides a color picker for easy color selection, the Font node offers built-in font selection for use with TextImage to generate text images, and the DynamicDelayByText node allows delayed execution based on the length of the input text.
 
-- [添加了DynamicDelayByText功能，可以根据输入文本的长度进行延迟执行。](./workflow/audio-chatgpt-workflow.json)
+- [添加了 DynamicDelayByText 功能，可以根据输入文本的长度进行延迟执行。](./workflow/audio-chatgpt-workflow.json)
 
 - [Added DynamicDelayByText, enabling delayed execution based on input text length.](./workflow/audio-chatgpt-workflow.json)
 
-- [使用CkptNames 对比不同的模型效果](./workflow/ckpts-image-workflow.json)
+- [使用 CkptNames 对比不同的模型效果](./workflow/ckpts-image-workflow.json)
 
 - [CkptNames compare the effects of different models.](./workflow/ckpts-image-workflow.json)
 
-
-
-## Other Nodes
+### Other Nodes
 
 ![main](./assets/all-workflow.svg)
 ![main2](./assets/detect-face-all.png)
 
 [workflow-1](./workflow/1-workflow.json)
 
-
-
 > TransparentImage
 
 ![TransparentImage](./assets/TransparentImage.png)
 
-
 > FeatheredMask、SmoothMask
 
 Add edges to an image.
 
 ![FeatheredMask](./assets/FlVou_Y6kaGWYoEj1Tn0aTd4AjMI.jpg)
 
-
 > LaMaInpainting
 
 from [simple-lama-inpainting](https://github.com/enesmsahin/simple-lama-inpainting)
 
-
 > rembgNode
 
 "briarmbg","u2net","u2netp","u2net_human_seg","u2net_cloth_seg","silueta","isnet-general-use","isnet-anime"
 
+**_ briarmbg _** model was developed by BRlA Al and can be used as an open-source model for non-commercial purposes
 
-
-### Improvement 
+### Improvement
 
 - Add "help" option to the context menu for each node.
 - Add "Nodes Map" option to the global context menu.
@@ -176,18 +250,21 @@ An improvement has been made to directly redirect to GitHub to search for missin
 
 ![node-not-found](./assets/node-not-found.png)
 
-
 ### Models
 
-[Download rembg Models](https://github.com/danielgatis/rembg/tree/main#Models),move to:models/rembg
+- [Download TripoSR](https://huggingface.co/stabilityai/TripoSR/blob/main/model.ckpt) and place it in `models/triposr`
 
-[Download lama](https://github.com/enesmsahin/simple-lama-inpainting/releases/download/v0.1.0/big-lama.pt), move to : models/lama
+- [Download facebook/dino-vitb16](https://huggingface.co/facebook/dino-vitb16/tree/main) and place it in `models/triposr/facebook/dino-vitb16`
 
-[Download Salesforce/blip-image-captioning-base](https://huggingface.co/Salesforce/blip-image-captioning-base), move to : models/clip_interrogator/Salesforce/blip-image-captioning-base
+[Download rembg Models](https://github.com/danielgatis/rembg/tree/main#Models),move to:`models/rembg`
 
-[Download succinctly/text2image-prompt-generator](https://huggingface.co/succinctly/text2image-prompt-generator/tree/main),move to:prompt_generator/text2image-prompt-generator
+[Download lama](https://github.com/enesmsahin/simple-lama-inpainting/releases/download/v0.1.0/big-lama.pt), move to : `models/lama`
 
-[Download Helsinki-NLP/opus-mt-zh-en](https://huggingface.co/Helsinki-NLP/opus-mt-zh-en/tree/main),move to:prompt_generator/opus-mt-zh-en
+[Download Salesforce/blip-image-captioning-base](https://huggingface.co/Salesforce/blip-image-captioning-base), move to :`models/clip_interrogator/Salesforce/blip-image-captioning-base`
+
+[Download succinctly/text2image-prompt-generator](https://huggingface.co/succinctly/text2image-prompt-generator/tree/main),move to:`models/prompt_generator/text2image-prompt-generator`
+
+[Download Helsinki-NLP/opus-mt-zh-en](https://huggingface.co/Helsinki-NLP/opus-mt-zh-en/tree/main),move to:`models/prompt_generator/opus-mt-zh-en`
 
 ## Installation
 
@@ -203,40 +280,35 @@ git clone https://github.com/shadowcz007/comfyui-mixlab-nodes.git
 Install the requirements:
 
 run directly:
+
 ```
 cd ComfyUI/custom_nodes/comfyui-mixlab-nodes
 install.bat
 ```
 
 or install the requirements using:
+
 ```
 ../../../python_embeded/python.exe -s -m pip install -r requirements.txt
 ```
 
 If you are using a venv, make sure you have it activated before installation and use:
+
 ```
 pip3 install -r requirements.txt
 ```
 
-
-
-
-
 #### Chinese community
-访问 [www.mixcomfy.com](https://www.mixcomfy.com)，获得更多内测功能，关注微信公众号：Mixlab无界社区
 
+访问 [www.mixcomfy.com](https://www.mixcomfy.com)，获得更多内测功能，关注微信公众号：Mixlab 无界社区
 
+####
 
-
-#### 
 File / LoadImagesFromPath SaveImageToLocal LoadImagesFromURL
 
-
-
-
 #### discussions:
-[discussions](https://github.com/shadowcz007/comfyui-mixlab-nodes/discussions)
 
+[discussions](https://github.com/shadowcz007/comfyui-mixlab-nodes/discussions)
 
 <picture>
   <source
@@ -256,4 +328,3 @@ File / LoadImagesFromPath SaveImageToLocal LoadImagesFromURL
     src="https://api.star-history.com/svg?repos=shadowcz007/comfyui-mixlab-nodes&type=Date"
   />
 </picture>
-

__init__.py

@@ -1,4 +1,3 @@
-#
 import os
 import subprocess
 import importlib.util
@@ -6,10 +5,35 @@ import sys,json
 import urllib
 import hashlib
 import datetime
-
-
+import folder_paths
+import logging
+import base64,io,re
+from PIL import Image
+from comfy.cli_args import args
 python = sys.executable
 
+# print("sys.path", sys.path)
+
+#修复 sys.stdout.isatty()  object has no attribute 'isatty'
+try:
+    sys.stdout.isatty()
+except:
+    # print('#fix sys.stdout.isatty')
+    sys.stdout.isatty = lambda: False
+
+llama_port=None
+llama_model=""
+llama_chat_format=""
+
+try:
+    from .nodes.ChatGPT import get_llama_models,get_llama_model_path,llama_cpp_client
+    llama_cpp_client("")
+
+except:
+    print("##nodes.ChatGPT ImportError")
+
+
+from .nodes.RembgNode import get_rembg_models,U2NET_HOME,run_briarmbg,run_rembg
 
 from server import PromptServer
 
@@ -42,7 +66,7 @@ def is_installed(package, package_overwrite=None):
             print(f"Couldn't install\nCommand: {command}\nError code: {result.returncode}")
     else:
         print(package+'## OK')
-        
+     
 try:
     import OpenSSL
 except ImportError:
@@ -79,6 +103,26 @@ install_openai()
 current_path = os.path.abspath(os.path.dirname(__file__))
 
 
+def remove_base64_prefix(base64_str):
+  """
+  去除 base64 字符串中的 data:image/*;base64, 前缀
+
+  Args:
+    base64_str: base64 编码的字符串
+
+  Returns:
+    去除前缀后的 base64 字符串
+  """
+
+  # 使用正则表达式匹配常见的前缀
+  pattern = r'^data:image\/(.*);base64,(.+)$'
+  match = re.match(pattern, base64_str)
+  if match:
+    # 如果匹配到常见的前缀，则去除前缀并返回
+    return match.group(2)
+  else:
+    # 如果不匹配到常见的前缀，则直接返回
+    return base64_str
 
 def calculate_md5(string):
     encoded_string = string.encode()
@@ -252,7 +296,8 @@ def get_my_workflow_for_app(filename="my_workflow_app.json",category="",is_all=F
                                 "name":x['app']['name'],
                                 "version":x['app']['version'],
                                 "input":input,
-                                "output":output
+                                "output":output,
+                                "id":x['app']['id']
                             }
                         },
                         "date":item["date"]
@@ -262,7 +307,7 @@ def get_my_workflow_for_app(filename="my_workflow_app.json",category="",is_all=F
                 print("发生异常：", str(e))
     else:
         app_workflow_path=os.path.join(category_path, filename)
-        # print('app_workflow_path: ',app_workflow_path)
+        print('app_workflow_path: ',app_workflow_path)
         try:
             with open(app_workflow_path) as json_file:
                 apps = [{
@@ -272,7 +317,8 @@ def get_my_workflow_for_app(filename="my_workflow_app.json",category="",is_all=F
         except Exception as e:
             print("发生异常：", str(e))
         
-        if len(apps)==1 and category!='' and category!=None:
+        # 这个代码不需要
+        # if len(apps)==1 and category!='' and category!=None:
             data=read_workflow_json_files(category_path)
             
             for item in data:
@@ -300,7 +346,8 @@ def get_my_workflow_for_app(filename="my_workflow_app.json",category="",is_all=F
                                 "name":x['app']['name'],
                                 "version":x['app']['version'],
                                 "input":input,
-                                "output":output
+                                "output":output,
+                                "id":x['app']['id']
                             }
                         },
                         "date":item["date"]
@@ -308,6 +355,31 @@ def get_my_workflow_for_app(filename="my_workflow_app.json",category="",is_all=F
 
     return apps
 
+# 历史记录
+def save_prompt_result(id,data):
+    prompt_result_path=os.path.join(current_path, "workflow/prompt_result.json")
+    prompt_result={}
+    if os.path.exists(prompt_result_path):
+        with open(prompt_result_path) as json_file:
+            prompt_result = json.load(json_file)
+    
+    prompt_result[id]=data
+    
+    with open(prompt_result_path, 'w') as file:
+        json.dump(prompt_result, file)
+    return prompt_result_path
+
+def get_prompt_result():
+    prompt_result_path=os.path.join(current_path, "workflow/prompt_result.json")
+    prompt_result={}
+    if os.path.exists(prompt_result_path):
+        with open(prompt_result_path) as json_file:
+            prompt_result = json.load(json_file)
+    res=list(prompt_result.values())
+    # print(res)
+    return res
+
+
 def save_workflow_json(data):
     workflow_path=os.path.join(current_path, "workflow/my_workflow.json")
     with open(workflow_path, 'w') as file:
@@ -386,37 +458,86 @@ async def new_start(self, address, port, verbose=True, call_on_start=None):
         runner = web.AppRunner(self.app, access_log=None)
         await runner.setup()
 
-        if not await check_port_available(address, port):
-            raise RuntimeError(f"Port {port} is already in use.")
+        # if not await check_port_available(address, port):
+        #     raise RuntimeError(f"Port {port} is already in use.")
+        
+        http_success = False
+        http_port=port
+        for i in range(11):  # 尝试最多11次
+            if await check_port_available(address, port + i):
+                http_port = port + i
+                site = web.TCPSite(runner, address, http_port)
+                await site.start()
+                http_success = True
+                break
 
-        site = web.TCPSite(runner, address, port)
-        await site.start()
+        if not http_success:
+            raise RuntimeError(f"Ports {port} to {port + 10} are all in use.")
+        
+
+        # site = web.TCPSite(runner, address, port)
+        # await site.start()
+
+        ssl_context = None
+        scheme = "http"
+        try:
+            # 跟着本体修改
+            if args.tls_keyfile and args.tls_certfile:
+                scheme = "https"
+                ssl_context = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_SERVER, verify_mode=ssl.CERT_NONE)
+                ssl_context.load_cert_chain(certfile=args.tls_certfile,
+                                    keyfile=args.tls_keyfile)
+            else:
+                # 如果没传，则自动创建
+                import ssl
+                crt, key = create_for_https()
+                ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+                ssl_context.load_cert_chain(crt, key)
+        except:
+            import ssl
+            crt, key = create_for_https()
+            ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+            ssl_context.load_cert_chain(crt, key)
 
-        import ssl
-        crt, key = create_for_https()
-        ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
-        ssl_context.load_cert_chain(crt, key)
 
         success = False
-        for i in range(10):  # 尝试最多10次
-            if await check_port_available(address, port + 1 + i):
-                https_port = port + 1 + i
+        for i in range(11):  # 尝试最多11次
+            if await check_port_available(address, http_port + 1 + i):
+                https_port = http_port + 1 + i
                 site2 = web.TCPSite(runner, address, https_port, ssl_context=ssl_context)
                 await site2.start()
                 success = True
                 break
 
         if not success:
-            raise RuntimeError(f"Ports {port + 1} to {port + 10} are all in use.")
+            raise RuntimeError(f"Ports {http_port + 1} to {http_port + 10} are all in use.")
 
         if address == '':
-            address = '0.0.0.0'
+            address = '127.0.0.1'
+        if address=='0.0.0.0':
+            address = '127.0.0.1'
+            
         if verbose:
-            print("\033[93mStarting server\n")
-            print("\033[93mTo see the GUI go to: http://{}:{}".format(address, port))
-            print("\033[93mTo see the GUI go to: https://{}:{}\033[0m".format(address, https_port))
+
+            logging.info("\n")
+            logging.info("\n\nStarting server")
+
+            # print("\033[93mStarting server\n")
+            logging.info("\033[93mTo see the GUI go to: http://{}:{}".format(address, http_port))
+            logging.info("\033[93mTo see the GUI go to: https://{}:{}\033[0m".format(address, https_port))
+ 
+            # print("\033[93mTo see the GUI go to: http://{}:{}".format(address, http_port))
+            # print("\033[93mTo see the GUI go to: https://{}:{}\033[0m".format(address, https_port))
+
         if call_on_start is not None:
-            call_on_start(address, port)
+            try:
+                if scheme=='https':
+                    call_on_start(scheme,address, https_port)
+                else:
+                    call_on_start(scheme,address, http_port)
+            except:
+                call_on_start(address,http_port)
+            
 
     except Exception as e:
         print(f"Error starting the server: {e}")
@@ -427,7 +548,6 @@ async def new_start(self, address, port, verbose=True, call_on_start=None):
         # webbrowser.open(f"https://{address}")
         # webbrowser.open(f"http://{address}:{port}")
 
-
 PromptServer.start=new_start
 
 # 创建路由表
@@ -456,7 +576,7 @@ async def mixlab_app_handler(request):
             return web.Response(text=html_data, content_type='text/html')
     else:
         return web.Response(text="HTML file not found", status=404)
-    
+
 
 @routes.post('/mixlab/workflow')
 async def mixlab_workflow_hander(request):
@@ -518,27 +638,301 @@ async def nodes_map_hander(request):
     return web.json_response(result)
 
 
-# 扩展api接口
-# from server import PromptServer
-# from aiohttp import web
+@routes.post("/mixlab/folder_paths")
+async def get_checkpoints(request):
+    data = await request.json()
+    t="checkpoints"
+    names=[]
+    try:
+        t=data['type']
+        names = folder_paths.get_filename_list(t)
+    except Exception as e:
+        print('/mixlab/folder_paths',False,e)
+
+    try:
+        if data['type']=='llamafile':
+            names=get_llama_models()
+    except:
+        print("llamafile none")
+
+    try:
+        if data['type']=='rembg':
+            names=get_rembg_models(U2NET_HOME)
+    except:
+        print("rembg none")
+    
+    return web.json_response({"names":names,"types":list(folder_paths.folder_names_and_paths.keys())})
+
+
+@routes.post('/mixlab/rembg')
+async def rembg_hander(request):
+    data = await request.json()
+    model=data['model']
+    result={}
+
+    data_base64=remove_base64_prefix(data['base64'])
+    image_data = base64.b64decode(data_base64)
+    
+    # 创建一个BytesIO对象
+    image_stream = io.BytesIO(image_data)
+
+    # 使用PIL Image模块读取图像
+    image = Image.open(image_stream)
+
+    if model=='briarmbg':
+        _,rgba_images,_=run_briarmbg([image])
+    else:
+        _,rgba_images,_=run_rembg(model,[image])
+
+    with io.BytesIO() as buf:
+        rgba_images[0].save(buf, format='PNG')
+        img_bytes = buf.getvalue()
+    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
+    
+    try: 
+        result={
+            'data':img_base64,
+            'model':model,
+            'status':'success',
+            }
+    except Exception as e:
+            print(e)
+
+    return web.json_response(result)
+
+
+@routes.post("/mixlab/prompt_result")
+async def post_prompt_result(request):
+    data = await request.json()
+    res=None
+    # print(data)
+    try:
+        action=data['action']
+        if action=='save':
+            result=data['data']
+            res=save_prompt_result(result['prompt_id'],result)
+        elif action=='all':
+            res=get_prompt_result()
+    except Exception as e:
+        print('/mixlab/prompt_result',False,e)
+    
+    return web.json_response({"result":res})
+
+
+
+def start_local_live_thread(data):
+    import asyncio
+    from VoiceStreamAI.server import Server
+    from VoiceStreamAI.asr.asr_factory import ASRFactory
+    from VoiceStreamAI.vad.vad_factory import VADFactory
+
+    model="large-v3"
+    if "model" in data:
+        model=data['model']
+
+    vad_pipeline = VADFactory.create_vad_pipeline("pyannote")
+    #device 
+    asr_pipeline = ASRFactory.create_asr_pipeline("faster_whisper", **{"model_size":model})
+
+    port=8765
+    if 'port' in data:
+        port=data['port']
+
+    llm_port=9000
+    if 'llm_port' in data:
+        llm_port=data['llm_port']
+
+    server = Server(vad_pipeline, 
+                    asr_pipeline, 
+                    host="127.0.0.1", 
+                    port=port, 
+                    sampling_rate=16000, 
+                    samples_width=2,
+                    llm_port=llm_port
+                    )
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(server.start())
+    loop.run_forever()
+
+
+async def start_local_llm(data):
+    global llama_port,llama_model,llama_chat_format
+    if llama_port and llama_model and llama_chat_format:
+        return {"port":llama_port,"model":llama_model,"chat_format":llama_chat_format}
+
+    import threading
+    import uvicorn
+    from llama_cpp.server.app import create_app
+    from llama_cpp.server.settings import (
+                Settings,
+                ServerSettings,
+                ModelSettings,
+                ConfigFileSettings,
+            )
+    
+    if not "model" in data and "model_path" in data:
+        data['model']= os.path.basename(data["model_path"])
+        model=data["model_path"]
+
+    elif "model" in data:
+        model=get_llama_model_path(data['model'])
+
+    n_gpu_layers=-1
+
+    if "n_gpu_layers" in data:
+        n_gpu_layers=data['n_gpu_layers']
+
+
+    chat_format="chatml"
+    if "model" in data and "function-calling" in data['model']:
+        chat_format="functionary-v2"
+
+    model_alias=os.path.basename(model)
+
+    # 多模态
+    clip_model_path=None
+    
+    prefix = "llava-phi-3-mini"
+    file_name = prefix+"-mmproj-"
+    if model_alias.startswith(prefix):
+        for file in os.listdir(os.path.dirname(model)):
+            if file.startswith(file_name):
+                clip_model_path=os.path.join(os.path.dirname(model),file)
+                chat_format='llava-1-5'
+    print('#clip_model_path',chat_format,clip_model_path)
+
+
+    address="127.0.0.1"
+    port=9090
+    success = False
+    for i in range(11):  # 尝试最多11次
+        if await check_port_available(address, port + i):
+            port = port + i
+            success = True
+            break
+
+    if success == False:
+        return {"port":None,"model":""}
+    
+    
+    server_settings=ServerSettings(host=address,port=port)
+
+    name, ext = os.path.splitext(os.path.basename(model))
+    print('#model',name)
+    app = create_app(
+                server_settings=server_settings,
+                model_settings=[
+                    ModelSettings(
+                    model=model,
+                    model_alias=name,
+                    n_gpu_layers=n_gpu_layers,
+                    n_ctx=4098,
+                    chat_format=chat_format,
+                    embedding=False,
+                    clip_model_path=clip_model_path
+                    )])
+
+    def run_uvicorn():
+        uvicorn.run(
+                app,
+                host=os.getenv("HOST", server_settings.host),
+                port=int(os.getenv("PORT", server_settings.port)),
+                ssl_keyfile=server_settings.ssl_keyfile,
+                ssl_certfile=server_settings.ssl_certfile,
+            )
+
+    # 创建一个子线程
+    thread = threading.Thread(target=run_uvicorn)
+
+    # 启动子线程
+    thread.start()
+
+    llama_port=port
+    llama_model=data['model']
+    llama_chat_format=chat_format
+
+    return  {"port":llama_port,"model":llama_model,"chat_format":llama_chat_format}
+
+# llam服务的开启
+@routes.post('/mixlab/start_llama')
+async def my_hander_method(request):
+    data =await request.json()
+    # print(data)
+    if llama_port and llama_model and llama_chat_format:
+        return web.json_response({"port":llama_port,"model":llama_model,"chat_format":llama_chat_format} )
+    try:
+        result=await start_local_llm(data)
+    except:
+        result= {"port":None,"model":"","llama_cpp_error":True}
+        print('start_local_llm error')
+
+    return web.json_response(result)
+
+
+@routes.post('/mixlab/start_live')
+async def mixlab_live_start_handler(request):
+    import threading
+    llm=await start_local_llm({
+        "model":"Phi-3-mini-4k-instruct-Q5_K_S.gguf",
+        "n_gpu_layers":2
+    })
+
+    # {"port":llama_port,"model":llama_model,"chat_format":llama_chat_format}
+
+    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' #hf_hub_download 里的下载地址修改
+    os.environ['PYANNOTE_AUTH_TOKEN'] = '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'
+
+    # Create and start the thread
+    data = {
+        "llm_port":llm['port'],
+        "port":8725,
+        "model":"large-v3"
+    }  # Replace with your actual data if needed
+    thread = threading.Thread(target=start_local_live_thread, args=(data,))
+    thread.start()
+
+    return web.json_response(data)
+
+
+@routes.get('/mixlab/live')
+async def mixlab_live_handler(request):
+    html_file = os.path.join(current_path, "web/live.html")
+    if os.path.exists(html_file):
+        with open(html_file, 'r', encoding='utf-8', errors='ignore') as f:
+            html_data = f.read()
+            return web.Response(text=html_data, content_type='text/html')
+    else:
+        return web.Response(text="HTML file not found", status=404)
+
+# 重启服务
+@routes.post('/mixlab/re_start')
+def re_start(request):
+    try:
+        sys.stdout.close_log()
+    except Exception as e:
+        pass
+    return os.execv(sys.executable, [sys.executable] + sys.argv)
 
-# @routes.post('/ws_image')
-# async def my_hander_method(request):
-#     post = await request.post()
-#     x = post.get("something")
-#     return web.json_response({})
 
 
 # 导入节点
-from .nodes.PromptNode import EmbeddingPrompt,RandomPrompt,PromptSlide,PromptSimplification,PromptImage,JoinWithDelimiter
-from .nodes.ImageNode import SaveImageToLocal,SplitImage,GridOutput,GetImageSize_,MirroredImage,ImageColorTransfer,NoiseImage,TransparentImage,GradientImage,LoadImagesFromPath,LoadImagesFromURL,ResizeImage,TextImage,SvgImage,Image3D,ShowLayer,NewLayer,MergeLayers,CenterImage,AreaToMask,SmoothMask,SplitLongMask,ImageCropByAlpha,EnhanceImage,FaceToMask
+from .nodes.PromptNode import GLIGENTextBoxApply_Advanced,EmbeddingPrompt,RandomPrompt,PromptSlide,PromptSimplification,PromptImage,JoinWithDelimiter
+from .nodes.ImageNode import ComparingTwoFrames,LoadImages_,CompositeImages,GridDisplayAndSave,GridInput,ImagesPrompt,SaveImageAndMetadata,SaveImageToLocal,SplitImage,GridOutput,GetImageSize_,MirroredImage,ImageColorTransfer,NoiseImage,TransparentImage,GradientImage,LoadImagesFromPath,LoadImagesFromURL,ResizeImage,TextImage,SvgImage,Image3D,ShowLayer,NewLayer,MergeLayers,CenterImage,AreaToMask,SmoothMask,SplitLongMask,ImageCropByAlpha,EnhanceImage,FaceToMask
 # from .nodes.Vae import VAELoader,VAEDecode
 from .nodes.ScreenShareNode import ScreenShareNode,FloatingVideo
 
 from .nodes.ChatGPT import ChatGPTNode,ShowTextForGPT,CharacterInText,TextSplitByDelimiter
 from .nodes.Audio import GamePal,SpeechRecognition,SpeechSynthesis
-from .nodes.Utils import CreateLoraNames,CreateSampler_names,CreateCkptNames,CreateSeedNode,TESTNODE_,TESTNODE_TOKEN,AppInfo,IntNumber,FloatSlider,TextInput,ColorInput,FontInput,TextToNumber,DynamicDelayProcessor,LimitNumber,SwitchByIndex,MultiplicationNode
-from .nodes.Mask import OutlineMask,FeatheredMask
+from .nodes.Utils import IncrementingListNode,ListSplit,CreateLoraNames,CreateSampler_names,CreateCkptNames,CreateSeedNode,TESTNODE_,TESTNODE_TOKEN,AppInfo,IntNumber,FloatSlider,TextInput,ColorInput,FontInput,TextToNumber,DynamicDelayProcessor,LimitNumber,SwitchByIndex,MultiplicationNode
+from .nodes.Mask import PreviewMask_,MaskListReplace,MaskListMerge,OutlineMask,FeatheredMask
+
+from .nodes.Style import ApplyVisualStylePrompting,StyleAlignedReferenceSampler,StyleAlignedBatchAlign,StyleAlignedSampleReferenceLatents
+
+from .nodes.Video import VideoCombine_Adv,LoadVideoAndSegment,ImageListReplace,VAEEncodeForInpaint_Frames
+
+from .nodes.TripoSR import LoadTripoSRModel,TripoSRSampler,SaveTripoSRMesh
 
 
 # 要导出的所有节点及其名称的字典
@@ -551,6 +945,7 @@ NODE_CLASS_MAPPINGS = {
     # "LoraPrompt":LoraPrompt,
     "EmbeddingPrompt":EmbeddingPrompt,
     "PromptSlide":PromptSlide,
+    "GLIGENTextBoxApply_Advanced":GLIGENTextBoxApply_Advanced,
     "PromptSimplification":PromptSimplification,
     "PromptImage":PromptImage,
     "MirroredImage":MirroredImage,
@@ -560,6 +955,7 @@ NODE_CLASS_MAPPINGS = {
     "ResizeImageMixlab":ResizeImage,
     "LoadImagesFromPath":LoadImagesFromPath,
     "LoadImagesFromURL":LoadImagesFromURL,
+    "LoadImagesToBatch":LoadImages_,
     "TextImage":TextImage,
     "EnhanceImage":EnhanceImage,
     "SvgImage":SvgImage,
@@ -567,9 +963,12 @@ NODE_CLASS_MAPPINGS = {
     "ImageColorTransfer":ImageColorTransfer,
     "ShowLayer":ShowLayer,
     "NewLayer":NewLayer,
+    "CompositeImages_":CompositeImages,
     "SplitImage":SplitImage,
     "CenterImage":CenterImage,
     "GridOutput":GridOutput,
+    "GridDisplayAndSave":GridDisplayAndSave,
+    "GridInput":GridInput,
     "MergeLayers":MergeLayers,
     "SplitLongMask":SplitLongMask,
     "FeatheredMask":FeatheredMask,
@@ -577,8 +976,11 @@ NODE_CLASS_MAPPINGS = {
     "FaceToMask":FaceToMask,
     "AreaToMask":AreaToMask,
     "ImageCropByAlpha":ImageCropByAlpha,
+    "ImagesPrompt_":ImagesPrompt,
     # "VAELoaderConsistencyDecoder":VAELoader,
     "SaveImageToLocal":SaveImageToLocal,
+    "SaveImageAndMetadata_":SaveImageAndMetadata,
+    "ComparingTwoFrames_":ComparingTwoFrames,
     # "VAEDecodeConsistencyDecoder":VAEDecode,
     "ScreenShare":ScreenShareNode,
     "FloatingVideo":FloatingVideo,
@@ -600,80 +1002,129 @@ NODE_CLASS_MAPPINGS = {
     "SwitchByIndex":SwitchByIndex,
     "LimitNumber":LimitNumber, 
     "OutlineMask":OutlineMask,
+    "MaskListMerge_":MaskListMerge,
     "JoinWithDelimiter":JoinWithDelimiter,
     "Seed_":CreateSeedNode,
     "CkptNames_":CreateCkptNames,
     "SamplerNames_":CreateSampler_names,
-    "LoraNames_":CreateLoraNames
-    # "LaMaInpainting":LaMaInpainting
+    "LoraNames_":CreateLoraNames,
+    "ApplyVisualStylePrompting_":ApplyVisualStylePrompting,
+    "StyleAlignedReferenceSampler_": StyleAlignedReferenceSampler,
+    "StyleAlignedSampleReferenceLatents_": StyleAlignedSampleReferenceLatents,
+    "StyleAlignedBatchAlign_": StyleAlignedBatchAlign,
+    "LoadVideoAndSegment_":LoadVideoAndSegment,
+    "VideoCombine_Adv":VideoCombine_Adv,
+    "ListSplit_":ListSplit,
+    "MaskListReplace_":MaskListReplace,
+    "ImageListReplace_":ImageListReplace,
+    "VAEEncodeForInpaint_Frames":VAEEncodeForInpaint_Frames,
+    "IncrementingListNode_":IncrementingListNode,
+    "PreviewMask_":PreviewMask_,
+     "LoadTripoSRModel_": LoadTripoSRModel,
+    "TripoSRSampler_": TripoSRSampler,
+    "SaveTripoSRMesh": SaveTripoSRMesh
     # "GamePal":GamePal
 }
 
 # 一个包含节点友好/可读的标题的字典
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "AppInfo":"AppInfo ♾️Mixlab",
-    "ResizeImageMixlab":"ResizeImage ♾️Mixlab",
+    "AppInfo":"App Info ♾️MixlabApp",
+    "Color":"Color Input ♾️MixlabApp",
+    "TextInput_":"Text Input ♾️MixlabApp",
+    "FloatSlider":"Float Slider Input ♾️MixlabApp",
+    "IntNumber":"Int Input ♾️MixlabApp",
+    "ImagesPrompt_":"Images Input ♾️MixlabApp",
+    "SaveImageAndMetadata_":"Save Image Output ♾️MixlabApp",
+    "ComparingTwoFrames_":"Comparing Two Frames ♾️MixlabApp",
+    "ResizeImageMixlab":"Resize Image ♾️Mixlab",
     "RandomPrompt": "Random Prompt ♾️Mixlab",
+    "PromptImage":"Output Prompt and Image ♾️Mixlab",
     "SplitLongMask":"Splitting a long image into sections",
     "VAELoaderConsistencyDecoder":"Consistency Decoder Loader",
     "VAEDecodeConsistencyDecoder":"Consistency Decoder Decode",
-    "ScreenShare":"ScreenShare ♾️Mixlab",
+    "ScreenShare":"Screen Share ♾️Mixlab",
     "FloatingVideo":"FloatingVideo ♾️Mixlab",
-    "ChatGPTOpenAI":"ChatGPT ♾️Mixlab",
-    "ShowTextForGPT":"ShowTextForGPT ♾️Mixlab",
-    "MergeLayers":"MergeLayers ♾️Mixlab",
+    "ChatGPTOpenAI":"ChatGPT & Local LLM ♾️Mixlab",
+    "ShowTextForGPT":"Show Text ♾️MixlabApp",
+    "MergeLayers":"Merge Layers ♾️Mixlab",
     "SpeechSynthesis":"SpeechSynthesis ♾️Mixlab",
     "SpeechRecognition":"SpeechRecognition ♾️Mixlab",
     "3DImage":"3DImage ♾️Mixlab",
+    "CompositeImages_":"Composite Images ♾️Mixlab",
     "DynamicDelayProcessor":"DynamicDelayByText ♾️Mixlab",
     "LaMaInpainting":"LaMaInpainting ♾️Mixlab",
-    "PromptSlide":"PromptSlide ♾️Mixlab",
-    "PromptGenerate_Mix":"PromptGenerate ♾️Mixlab",
-    "ChinesePrompt_Mix":"ChinesePrompt ♾️Mixlab",
+    "PromptSlide":"Prompt Slide ♾️Mixlab",
+    "PromptGenerate_Mix":"Prompt Generate ♾️Mixlab",
+    "ChinesePrompt_Mix":"Chinese Prompt ♾️Mixlab",
     "GamePal":"GamePal ♾️Mixlab",
-    "RembgNode_Mix":"Removebg",
-    "LoraNames_":"LoraName_TriggerWords.safetensors"
+    "RembgNode_Mix":"Remove Background ♾️Mixlab",
+    "LoraNames_":"LoraName ♾️Mixlab",
+    "ApplyVisualStylePrompting_":"Apply VisualStyle Prompting ♾️Mixlab",
+    "StyleAlignedReferenceSampler_": "StyleAligned Reference Sampler ♾️Mixlab",
+    "StyleAlignedSampleReferenceLatents_": "StyleAligned Sample Reference Latents ♾️Mixlab",
+    "StyleAlignedBatchAlign_": "StyleAligned Batch Align ♾️Mixlab",
+    "LoadVideoAndSegment_":"Load Video And Segment ♾️Mixlab",
+    "VideoCombine_Adv":"Video Combine ♾️Mixlab",
+    "MaskListMerge_":"MaskList to Mask ♾️Mixlab",
+    "ListSplit_":"Split List ♾️Mixlab",
+    "MaskListReplace_":"MaskList Replace ♾️Mixlab",
+    "ImageListReplace_":"ImageList Replace ♾️Mixlab",
+    "SwitchByIndex":"List Switch By Index ♾️Mixlab",
+    "GLIGENTextBoxApply_Advanced":"GLIGEN TextBox Apply ♾️Mixlab",
+    "GridDisplayAndSave":"Grid Display And Save ♾️Mixlab",
+    "GridInput":"Grid Input ♾️Mixlab",
+    "GridOutput":"Grid Output ♾️Mixlab",
+    "GetImageSize_":"Get Image Size ♾️Mixlab",
+    "VAEEncodeForInpaint_Frames":"VAE Encode For Inpaint Frames ♾️Mixlab",
+    "IncrementingListNode_":"Create Incrementing Number List ♾️Mixlab",
+    "LoadImagesToBatch":"Load Images(base64) ♾️Mixlab",
+    "PreviewMask_":"Preview Mask",
+    "LoadTripoSRModel_": "Load TripoSR Model",
+    "TripoSRSampler_": "TripoSR Sampler",
+    "SaveTripoSRMesh": "Save TripoSR Mesh"
 }
 
 # web ui的节点功能
 WEB_DIRECTORY = "./web"
 
-print('--------------')
-print('\033[91m ### Mixlab Nodes: \033[93mLoaded')
+
+logging.info('--------------')
+logging.info('\033[91m ### Mixlab Nodes: \033[93mLoaded')
+# print('\033[91m ### Mixlab Nodes: \033[93mLoaded')
 
 try:
     from .nodes.Lama import LaMaInpainting
-    print('LaMaInpainting.available',LaMaInpainting.available)
+    logging.info('LaMaInpainting.available {}'.format(LaMaInpainting.available))
     if LaMaInpainting.available:
         NODE_CLASS_MAPPINGS['LaMaInpainting']=LaMaInpainting
 except Exception as e:
-    print('LaMaInpainting.available',False,e)
+    logging.info('LaMaInpainting.available False')
 
 try:
     from .nodes.ClipInterrogator import ClipInterrogator
-    print('ClipInterrogator.available',ClipInterrogator.available)
+    logging.info('ClipInterrogator.available {}'.format(ClipInterrogator.available))
     if ClipInterrogator.available:
         NODE_CLASS_MAPPINGS['ClipInterrogator']=ClipInterrogator
 except Exception as e:
-    print('ClipInterrogator.available',False,e)
+    logging.info('ClipInterrogator.available False')
 
 try:
     from .nodes.TextGenerateNode import PromptGenerate,ChinesePrompt
-    print('PromptGenerate.available',PromptGenerate.available)
+    logging.info('PromptGenerate.available {}'.format(PromptGenerate.available))
     if PromptGenerate.available:
         NODE_CLASS_MAPPINGS['PromptGenerate_Mix']=PromptGenerate
-    print('ChinesePrompt.available',ChinesePrompt.available)
+    logging.info('ChinesePrompt.available {}'.format(ChinesePrompt.available))
     if ChinesePrompt.available:
         NODE_CLASS_MAPPINGS['ChinesePrompt_Mix']=ChinesePrompt
 except Exception as e:
-    print('TextGenerateNode.available',False,e)
+    logging.info('TextGenerateNode.available False')
 
 try:
     from .nodes.RembgNode import RembgNode_
-    print('RembgNode_.available',RembgNode_.available)
+    logging.info('RembgNode_.available {}'.format(RembgNode_.available))
     if RembgNode_.available:
         NODE_CLASS_MAPPINGS['RembgNode_Mix']=RembgNode_
 except Exception as e:
-    print('RembgNode_.available',False,e)
+    logging.info('RembgNode_.available False' )
 
-print('\033[93m -------------- \033[0m')
+logging.info('\033[93m -------------- \033[0m')

data/extension-node-map.json

@@ -5201,7 +5201,7 @@
             "title_aux": "ComfyUI Stable Video Diffusion"
         }
     ],
-    "https://github.com/thedyze/save-image-extended-comfyui": [
+    "https://github.com/audioscavenger/save-image-extended-comfyui": [
         [
             "SaveImageExtended"
         ],
@@ -5705,4 +5705,4 @@
             "title_aux": "SDXLCustomAspectRatio"
         }
     ]
-}
+}

data/image_prompt_example.json

@@ -0,0 +1,10 @@
+[
+    {
+        "keyword":"Dog",
+        "imgurl":"http://127.0.0.1:8188/view?filename=1709966910233.png&type=input&subfolder=&rand=0.2734446552394221"
+    },
+    {
+        "keyword":"x",
+        "imgurl":"http://127.0.0.1:8188/view?filename=image%20(33).png&type=input&subfolder=pasted&rand=0.6984318219852814"
+    }
+]

install.bat

@@ -10,6 +10,12 @@ if exist "%python_exec%" (
     for /f "delims=" %%i in (%requirements_txt%) do (
         %python_exec% -s -m pip install "%%i" -i https://pypi.tuna.tsinghua.edu.cn/simple
     )
+
+    %python_exec% -s -m pip install --upgrade --force llama-cpp-python --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121
+
+    %python_exec% -s -m pip install --upgrade --force llama-cpp-python[server]
+    
+
 ) else (
     echo Installing with system Python
     for /f "delims=" %%i in (%requirements_txt%) do (

nodes/ChatGPT.py

@@ -4,6 +4,17 @@ import urllib.error
 import re,json,os,string,random
 import folder_paths
 import hashlib
+import codecs,sys
+import importlib.util
+
+
+def is_installed(package):
+    try:
+        spec = importlib.util.find_spec(package)
+    except ModuleNotFoundError:
+        return False
+    return spec is not None
+
 
 def get_unique_hash(string):
     hash_object = hashlib.sha1(string.encode())
@@ -47,6 +58,108 @@ def openai_client(key,url):
     )
     return client
 
+def ZhipuAI_client(key):
+
+    try:
+        if is_installed('zhipuai')==False:
+            import subprocess
+
+            # 安装
+            print('#pip install zhipuai')
+
+            result = subprocess.run([sys.executable, '-s', '-m', 'pip', 'install', 'zhipuai'], capture_output=True, text=True)
+
+            #检查命令执行结果
+            if result.returncode == 0:
+                print("#install success")
+                from zhipuai import ZhipuAI
+            else:
+                print("#install error")
+            
+        else:
+            from zhipuai import ZhipuAI
+    except:
+        print("#install zhipuai error")
+
+    client = ZhipuAI(
+        api_key=key, # 填写您的 APIKey
+    ) 
+    return client
+
+
+# 优先使用phi
+def phi_sort(lst):
+    return sorted(lst, key=lambda x: x.lower().count('phi'), reverse=True)
+
+def get_llama_path():
+    try:
+        return folder_paths.get_folder_paths('llamafile')[0]
+    except:
+        return os.path.join(folder_paths.models_dir, "llamafile")
+
+def get_llama_models():
+    res=[]
+
+    model_path=get_llama_path()
+    if os.path.exists(model_path):
+        files = os.listdir(model_path)
+        for file in files:
+            if os.path.isfile(os.path.join(model_path, file)):
+                res.append(file)
+        res=phi_sort(res)
+    return res
+
+llama_modes_list=get_llama_models()
+
+def get_llama_model_path(file_name):
+    model_path=get_llama_path()
+    mp=os.path.join(model_path,file_name)
+    return mp
+
+def llama_cpp_client(file_name):
+    try:
+        if is_installed('llama_cpp')==False:
+            import subprocess
+
+            # 安装
+            print('#pip install llama-cpp-python')
+ 
+            result = subprocess.run([sys.executable, '-s', '-m', 'pip', 
+                                     'install', 
+                                     'llama-cpp-python',
+                                     '--extra-index-url',
+                                     'https://abetlen.github.io/llama-cpp-python/whl/cu121'
+                                     ], capture_output=True, text=True)
+
+            #检查命令执行结果
+            if result.returncode == 0:
+                print("#install success")
+                from llama_cpp import Llama
+
+                subprocess.run([sys.executable, '-s', '-m', 'pip', 
+                                     'install', 
+                                     'llama-cpp-python[server]'
+                                     ], capture_output=True, text=True)
+
+            else:
+                print("#install error")
+            
+        else:
+            from llama_cpp import Llama
+    except:
+        print("#install llama-cpp-python error")
+
+    if file_name:
+        mp=get_llama_model_path(file_name)
+        # file_name=get_llama_models()[0]
+        # model_path=os.path.join(folder_paths.models_dir, "llamafile")
+        # mp=os.path.join(model_path,file_name)
+
+        llm = Llama(model_path=mp, chat_format="chatml",n_gpu_layers=-1,n_ctx=512)
+
+        return llm
+
+    
 
 
 def chat(client, model_name,messages ):
@@ -55,10 +168,21 @@ def chat(client, model_name,messages ):
         while True:
             try_count += 1
             try:
-                response = client.chat.completions.create(
-                    model=model_name,
-                    messages=messages
-                )
+                if hasattr(client, "chat"):
+                    response = client.chat.completions.create(
+                        model=model_name,
+                        messages=messages
+                    )
+                else:
+                    # 是llama的
+                    response = client.create_chat_completion_openai_v1(
+                        messages=messages,
+                        # response_format={
+                        #     "type": "json_object",
+                        # },
+                        # temperature=0.7,
+                    )
+
                 break
             except openai.AuthenticationError as ex:
                 raise ex
@@ -67,7 +191,8 @@ def chat(client, model_name,messages ):
                     raise ex
                 time.sleep(3)
                 continue
-
+        
+        # print(response.keys())
         finish_reason = response.choices[0].finish_reason
         if finish_reason != "stop":
             raise RuntimeError("API finished with unexpected reason: " + finish_reason)
@@ -90,6 +215,16 @@ class ChatGPTNode:
 
     @classmethod
     def INPUT_TYPES(cls):
+        model_list=llama_modes_list+[
+                    "gpt-3.5-turbo",
+                    "gpt-3.5-turbo-0125",
+                    "gpt-35-turbo",
+                    "gpt-3.5-turbo-16k", 
+                    "gpt-3.5-turbo-16k-0613", 
+                    "gpt-4-0613",
+                    "gpt-4-1106-preview",
+                    "glm-4"
+                    ]
         return {
             "required": {
                 "api_key":("KEY", {"default": "", "multiline": True,"dynamicPrompts": False}),
@@ -100,8 +235,8 @@ class ChatGPTNode:
                                        "default": "You are ChatGPT, a large language model trained by OpenAI. Answer as concisely as possible.", 
                                        "multiline": True,"dynamicPrompts": False
                                        }),
-                "model": (["gpt-3.5-turbo","gpt-35-turbo","gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613", "gpt-4-0613","gpt-4-1106-preview"], 
-                          {"default": "gpt-3.5-turbo"}),
+                "model": ( model_list, 
+                    {"default": model_list[0]}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "step": 1}),
                 "context_size":("INT", {"default": 1, "min": 0, "max":30, "step": 1}),
             },
@@ -124,8 +259,8 @@ class ChatGPTNode:
                                  api_url, 
                                  prompt, 
                                  system_content,
-                                   model, 
-                                   seed,context_size,unique_id = None, extra_pnginfo=None):
+                                model, 
+                                seed,context_size,unique_id = None, extra_pnginfo=None):
         # print(api_key!='',api_url,prompt,system_content,model,seed)
         # 可以选择保留会话历史以维持上下文记忆
         # 或者在此处清除会话历史 self.session_history.clear()
@@ -143,8 +278,16 @@ class ChatGPTNode:
         if is_azure_url(api_url):
             client=azure_client(api_key,api_url)
         else:
-            client=openai_client(api_key,api_url)
-            print('openai url')
+            # 根据用户选择的模型，设置相应的接口和模型名称
+            if model == "glm-4" :
+                client = ZhipuAI_client(api_key)  # 使用 Zhipuai 的接口
+                print('using Zhipuai interface')
+            elif model in llama_modes_list:
+                #
+                client=llama_cpp_client(model)
+            else :
+                client = openai_client(api_key,api_url)  # 使用 ChatGPT  的接口
+                print('using ChatGPT interface')
 
         # 把用户的提示添加到会话历史中
         # 调用API时传递整个会话历史
@@ -199,7 +342,7 @@ class ShowTextForGPT:
     OUTPUT_NODE = True
     OUTPUT_IS_LIST = (True,)
 
-    CATEGORY = "♾️Mixlab/GPT"
+    CATEGORY = "♾️Mixlab/Text"
 
     def run(self, text,output_dir=[""]):
         
@@ -283,7 +426,7 @@ class CharacterInText:
     # OUTPUT_NODE = True
     OUTPUT_IS_LIST = (False,)
 
-    CATEGORY = "♾️Mixlab/GPT"
+    CATEGORY = "♾️Mixlab/Text"
 
     def run(self, text,character,start_index):
         # print(text,character,start_index)
@@ -296,8 +439,8 @@ class TextSplitByDelimiter:
     def INPUT_TYPES(s):
         return {
             "required": {
-                 "text": ("STRING", {"multiline": True,"dynamicPrompts": False}),
-                "delimiter":(["newline","comma"],),
+                "text": ("STRING", {"multiline": True,"dynamicPrompts": False}),
+                "delimiter":("STRING", {"multiline": False,"default":",","dynamicPrompts": False}),
                 "start_index": ("INT", {
                     "default": 0,
                     "min": 0, #Minimum value
@@ -328,17 +471,16 @@ class TextSplitByDelimiter:
     # OUTPUT_NODE = True
     OUTPUT_IS_LIST = (True,)
 
-    CATEGORY = "♾️Mixlab/GPT"
+    CATEGORY = "♾️Mixlab/Text"
 
     def run(self, text,delimiter,start_index,skip_every,max_count):
-        arr=[]
-        if delimiter=='newline':
-            arr = [line for line in text.split('\n') if line.strip()]
-        elif delimiter=='comma':
-            arr = [line for line in text.split(',') if line.strip()]
-        
+         
+        if delimiter=="":
+            arr=[text.strip()]
+        else:
+            delimiter=codecs.decode(delimiter, 'unicode_escape')
+            arr= [line for line in text.split(delimiter) if line.strip()]
+
         arr= arr[start_index:start_index + max_count * (skip_every+1):(skip_every+1)]
 
         return (arr,)
-
-

nodes/ClipInterrogator.py

@@ -70,14 +70,20 @@ def load_caption_model(model_path,config,t='blip-base'):
     return (caption_model,caption_processor)
     
 
+def get_clip_interrogator_path():
+    try:
+        return folder_paths.get_folder_paths('clip_interrogator')[0]
+    except:
+        return os.path.join(folder_paths.models_dir, "clip_interrogator")
 
-caption_model_path=os.path.join(folder_paths.models_dir, "clip_interrogator/Salesforce/blip-image-captioning-base")
+
+cache_path=get_clip_interrogator_path()
+
+caption_model_path=os.path.join(cache_path, "Salesforce/blip-image-captioning-base")
 if not os.path.exists(caption_model_path):
     print(f"## clip_interrogator_model not found: {caption_model_path}, pls download from https://huggingface.co/Salesforce/blip-image-captioning-base")
     caption_model_path='Salesforce/blip-image-captioning-base'
 
-cache_path=os.path.join(folder_paths.models_dir, "clip_interrogator")
-
 
 # Tensor to PIL
 def tensor2pil(image):

nodes/Lama.py

@@ -42,8 +42,13 @@ else:
     _available=True
 
 
- 
-llma_model_path=os.path.join(folder_paths.models_dir, "lama/big-lama.pt")
+def get_lama_path():
+    try:
+        return folder_paths.get_folder_paths('lama')[0]
+    except:
+        return os.path.join(folder_paths.models_dir, "lama")
+
+llma_model_path=os.path.join(get_lama_path(), "big-lama.pt")
 if not os.path.exists(llma_model_path):
     os.environ['LAMA_MODEL']=''
     print(f"## lama torchscript model not found: {llma_model_path},pls download from https://github.com/enesmsahin/simple-lama-inpainting/releases/download/v0.1.0/big-lama.pt")

nodes/Mask.py

@@ -2,16 +2,14 @@
 import scipy.ndimage
 import torch
 
-from nodes import MAX_RESOLUTION
-
 import numpy as np 
 # from PIL import Image, ImageDraw
 from PIL import Image, ImageOps 
 
 from comfy.cli_args import args
-import cv2
-
-
+import cv2,os
+from nodes import MAX_RESOLUTION, SaveImage, common_ksampler
+import folder_paths,random
 
 # Tensor to PIL
 def tensor2pil(image):
@@ -22,6 +20,19 @@ def pil2tensor(image):
     return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
 
 
+def add_masks(mask1, mask2):
+    mask1 = mask1.cpu()
+    mask2 = mask2.cpu()
+    cv2_mask1 = np.array(mask1) * 255
+    cv2_mask2 = np.array(mask2) * 255
+
+    if cv2_mask1.shape == cv2_mask2.shape:
+        cv2_mask = cv2.add(cv2_mask1, cv2_mask2)
+        return torch.clamp(torch.from_numpy(cv2_mask) / 255.0, min=0, max=1)
+    else:
+        return mask1
+    
+
 def grow(mask, expand, tapered_corners):
     c = 0 if tapered_corners else 1
     kernel = np.array([[c, 1, c],
@@ -58,6 +69,35 @@ def combine(destination, source, x, y):
 
     return output
 
+
+class PreviewMask_(SaveImage):
+    def __init__(self):
+        self.output_dir = folder_paths.get_temp_directory()
+        self.type = "temp"
+        self.prefix_append =''.join(random.choice("abcdehijklmnopqrstupvxyzfg") for x in range(5))
+        self.compress_level = 4
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+                "required": {
+                                "mask": ("MASK",), 
+                            }
+            }
+    
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/Mask"
+
+    # 运行的函数
+    def run(self, mask ):
+        img=tensor2pil(mask)
+        img=img.convert('RGB')
+        img=pil2tensor(img) 
+        return self.save_images(img, 'temp_', None,  None)
+
+
 class OutlineMask:
 
     @classmethod
@@ -87,6 +127,69 @@ class OutlineMask:
         return (m3,)
     
 
+class MaskListReplace:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "masks": ("MASK",),
+            "mask_replace": ("MASK",),
+            "start_index":("INT", {"default": 0, "min": 0, "step": 1}),
+            "end_index":("INT", {"default": 0, "min": 0, "step": 1}),
+            "invert": ("BOOLEAN", {"default": False}),
+        }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/Video"
+
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True,)
+
+    def run(self, masks,mask_replace,start_index,end_index,invert):
+        mask_replace=mask_replace[0]
+        start_index=start_index[0]
+        end_index=end_index[0]
+        invert=invert[0]
+        
+        new_masks=[]
+        for i in range(len(masks)):
+            if i>=start_index and i<=end_index:
+                if invert:
+                    new_masks.append(masks[i])
+                else:
+                    new_masks.append(mask_replace)
+            else:
+                if invert:
+                    new_masks.append(mask_replace)
+                else:
+                    new_masks.append(masks[i])
+
+        return (new_masks,)
+
+
+class MaskListMerge:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "masks": ("MASK",),
+        }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/Mask"
+
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (False,)
+
+    def run(self, masks):
+        mask=masks[0]
+        if isinstance(masks, list):
+            for m in masks:
+                # print(m.shape)
+                mask = add_masks(mask, m)
+        return (mask,)
 
 
 class FeatheredMask:

nodes/PromptNode.py

@@ -21,6 +21,7 @@ import json
 embeddings_path=os.path.join(folder_paths.models_dir, "embeddings")
 
 def get_files_with_extension(directory, extension):
+    
     file_list = []
     for root, dirs, files in os.walk(directory):
         for file in files:
@@ -188,7 +189,7 @@ class PromptImage:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Prompt"
+    CATEGORY = "♾️Mixlab/Output"
 
     # 运行的函数
     def run(self,prompts,images,save_to_image):
@@ -517,9 +518,10 @@ class RandomPrompt:
 class EmbeddingPrompt:
     @classmethod
     def INPUT_TYPES(s):
+
         return {
             "required": {
-                "embedding":(get_files_with_extension(embeddings_path,'.pt'),),
+                "embedding":(folder_paths.get_filename_list("embeddings"),),
                 "weight": ("FLOAT", {"default": 1, "min": -2, "max": 2,"step":0.01 ,"display": "slider"}),
                 },
                 
@@ -544,14 +546,119 @@ class EmbeddingPrompt:
         # return (new_prompt)
         return (prompt,)
 
-RETURN_TYPES = (any_type,)
+# RETURN_TYPES = (any_type,)
+    
+# conditioning ：提示，正向or负向
+# clip：clip模型
+# gligen_textbox_model：gligen模型
+# grids：矩形框的集合
+# labels：每个矩形框对应的标签的集合
+# index：选取第几个矩形框作为gligen的box
+
+class GLIGENTextBoxApply_Advanced:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"conditioning": ("CONDITIONING", ),
+                              "clip": ("CLIP", ),
+                              "gligen_textbox_model": ("GLIGEN", ),
+                              "grids": ("_GRID",),
+                              "labels": ("STRING", 
+                                        {
+                                            "multiline": True, 
+                                            "default": "",
+                                            "forceInput": True
+                                        }),
+                              "index": ("INT", {"default": -1, "min": -1, "max": 300, "step": 1}),
+                              "max_size": ("INT", {"default": 8, "min": 1, "max": 300, "step": 1}),
+                              "random_shuffle":(["on","off"],),
+                             },
+                 "optional":{
+                    "seed": (any_type,  {"default": 0, "min": 0, "max": 0xffffffffffffffff,"step": 1}),
+                } 
+                            }
+    RETURN_TYPES = ("CONDITIONING","STRING",)
+    RETURN_NAMES = ("CONDITIONING","label",)
+
+    FUNCTION = "run"
+    # INPUT_IS_LIST = True
+    CATEGORY = "♾️Mixlab/Prompt"
+
+    def run(self, conditioning, clip, gligen_textbox_model, grids, labels, index,max_size,random_shuffle,seed=0):
+        # print('grids',grids)
+        # conditioning=conditioning[0]
+        # clip=clip[0]
+        # gligen_textbox_model=gligen_textbox_model[0]
+        # index=index[0]
+        # max_size=max_size[0]
+        # random_shuffle=random_shuffle[0]
+
+        texts=labels
+        
+        if index>-1:
+            texts=[labels[index]]
+            grids=[grids[index]]
+
+        if random_shuffle=='on':
+            sss=[[texts[i],grids[i]] for i in range(len(texts))]
+            random.shuffle(sss)
+            texts=[s[0] for s in sss]
+            grids=[s[1] for s in sss]
+
+        if len(texts) > max_size:
+            texts = texts[:max_size]
+
+        c = []
+        
+        for t in conditioning:
+            n = [t[0], t[1].copy()]
+
+
+            # 多个
+            position_params=[]
+            for i in range(len(texts)):
+                text=texts[i]
+                grid=grids[i]
+                x,y,width,height=grid
+                # print(text)
+                cond, cond_pooled = clip.encode_from_tokens(clip.tokenize(text), return_pooled=True)
+                position_params =position_params+ [(cond_pooled, height // 8, width // 8, y // 8, x // 8)]
+
+            # 前一个
+            prev = []
+            if "gligen" in n[1]:
+                prev = n[1]['gligen'][2]
+
+            n[1]['gligen'] = ("position", gligen_textbox_model, prev + position_params)
+            # print('gligen',n)
+            c.append(n)
+        
+        #  下面这个写法有bug
+        # for i in range(len(texts)):
+        #     text=texts[i]
+        #     grid=grids[i]
+        #     x,y,width,height=grid
+
+        #     cond, cond_pooled = clip.encode_from_tokens(clip.tokenize(text), return_pooled=True)
+        #     for t in conditioning:
+        #         n = [t[0], t[1].copy()]
+        #         position_params = [(cond_pooled, height // 8, width // 8, y // 8, x // 8)]
+        #         prev = []
+        #         if "gligen" in n[1]:
+        #             prev = n[1]['gligen'][2]
+
+        #         n[1]['gligen'] = ("position", gligen_textbox_model, prev + position_params)
+        #         c.append(n)
+
+
+        return (c,texts, )
+    
 
 class JoinWithDelimiter:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
                     "text_list": (any_type,),
-                    "delimiter":(["newline","comma"],),
+                    "delimiter":(["newline","comma","backslash","space"],),
                              },
                 }
     
@@ -559,7 +666,7 @@ class JoinWithDelimiter:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Prompt"
+    CATEGORY = "♾️Mixlab/Text"
 
     INPUT_IS_LIST = True # 当true的时候，输入时list，当false的时候，如果输入是list，则会自动包一层for循环调用
     OUTPUT_IS_LIST = (False,)
@@ -570,6 +677,10 @@ class JoinWithDelimiter:
             delimiter='\n'
         elif delimiter=='comma':
             delimiter=','
+        elif delimiter=='backslash':
+            delimiter='\\'
+        elif delimiter=='space':
+            delimiter=' '
         t=''
         if isinstance(text_list, list):
             t=join_with_(text_list,delimiter)

nodes/RembgNode.py

@@ -467,15 +467,37 @@ class BriaRMBG(nn.Module):
 
 
 
+def get_U2NET_model_path():
+    try:
+        return folder_paths.get_folder_paths('rembg')[0]
+    except:
+        return os.path.join(folder_paths.models_dir, "rembg")
+    
 
-
-
-U2NET_HOME=os.path.join(folder_paths.models_dir, "rembg")
+U2NET_HOME=get_U2NET_model_path()
 os.environ["U2NET_HOME"] = U2NET_HOME
 
 global _available
 _available=False
 
+
+def get_rembg_models(path):
+    """从目录中获取文件并提取文件名
+    Args:
+        path: 目录路径
+    Returns:
+        文件名列表
+    """
+    filenames = []
+    for root, _, files in os.walk(path):
+        for filename in files:
+        # 过滤隐藏文件
+            if not filename.startswith('.'):
+                name, ext = os.path.splitext(os.path.basename(filename))
+                filenames.append(name)
+    return filenames
+
+
 def is_installed(package):
     try:
         spec = importlib.util.find_spec(package)
@@ -509,8 +531,8 @@ except:
     _available=False
 
 
-def briarmbg_run(images=[]):
-    mroot=os.path.join(folder_paths.models_dir, "rembg")
+def run_briarmbg(images=[]):
+    mroot=U2NET_HOME
     m=os.path.join(mroot,'briarmbg.pth')
     if os.path.exists(m)==False:
         # 下载
@@ -573,14 +595,15 @@ def briarmbg_run(images=[]):
     return (masks,rgba_images,rgb_images)   
 
 
-def run_bg(model_name= "unet",images=[]):
+def run_rembg(model_name= "unet",images=[],callback=None):
     # model_name = "unet" # "isnet-general-use"
+    # print('#run_rembg',model_name)
     rembg_session = new_session(model_name)
     masks=[]
     rgba_images=[]
     rgb_images=[]
     # 进度条
-    pbar = comfy.utils.ProgressBar(len(images) )
+    pbar=callback 
     for img in images:
         # use the post_process_mask argument to post process the mask to get better results.
         mask = remove(img, session=rembg_session,only_mask=True,post_process_mask=True)
@@ -620,8 +643,9 @@ def run_bg(model_name= "unet",images=[]):
             rgb_image = Image.new("RGB", image_rgba.size, (0, 0, 0))
             rgb_image.paste(image_rgba, mask=image_rgba.split()[3])
             rgb_images.append(rgb_image)
-            
-        pbar.update(1)
+
+        if pbar:
+            pbar.update(1)
     return (masks,rgba_images,rgb_images)
 
 
@@ -643,17 +667,7 @@ class RembgNode_:
     def INPUT_TYPES(s):
         return {"required": {
             "image": ("IMAGE",),
-            "model_name": ([
-                            "briarmbg",
-                            "u2net",
-                            "u2netp",
-                            "u2net_human_seg",
-                            "u2net_cloth_seg",
-                            "silueta",
-                            "isnet-general-use",
-                            "isnet-anime",
-                            
-                            ],),
+            "model_name": (get_rembg_models(U2NET_HOME),),
             
                              },
                 }
@@ -681,9 +695,9 @@ class RembgNode_:
                 images.append(im)
 
         if model_name=='briarmbg':
-            masks,rgba_images,rgb_images=briarmbg_run(images)
+            masks,rgba_images,rgb_images=run_briarmbg(images)
         else:
-            masks,rgba_images,rgb_images=run_bg(model_name,images)
+            masks,rgba_images,rgb_images=run_rembg(model_name,images, comfy.utils.ProgressBar(len(images) ))
 
         masks=[pil2tensor(m) for m in masks]
 

nodes/ScreenShareNode.py

@@ -93,7 +93,7 @@ class ScreenShareNode:
     RETURN_NAMES = ("IMAGE","PROMPT","FLOAT","INT")
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Image"
+    CATEGORY = "♾️Mixlab/Screen"
 
     # INPUT_IS_LIST = True
     OUTPUT_IS_LIST = (False,False,False,False)
@@ -118,7 +118,7 @@ class FloatingVideo:
     OUTPUT_NODE = True
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Image"
+    CATEGORY = "♾️Mixlab/Screen"
 
     # INPUT_IS_LIST = True
     # OUTPUT_IS_LIST = (False,False,)

nodes/Style.py

@@ -0,0 +1,503 @@
+import comfy
+import torch
+
+from dataclasses import dataclass
+import torch.nn as nn
+from comfy.model_patcher import ModelPatcher
+import comfy.ops
+from typing import Union
+import comfy.sample
+import latent_preview
+import comfy.utils
+
+T = torch.Tensor
+
+
+from .VisualStylePrompting.attention_functions import VisualStyleProcessor
+
+class ApplyVisualStylePrompting:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "reference_image": ("IMAGE",),
+                "reference_image_text": ("STRING", {"multiline": True}),
+                "model": ("MODEL",),
+                "clip": ("CLIP", ),
+                "vae": ("VAE", ),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING", ),
+                "enabled": ("BOOLEAN", {"default": True}),
+                "denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096,"step":2})
+            }
+        }
+
+    RETURN_TYPES = ("MODEL", "CONDITIONING","CONDITIONING", "LATENT")
+    RETURN_NAMES = ("model", "positive", "negative", "latents")
+    
+    CATEGORY = "♾️Mixlab/Style"
+
+    FUNCTION = "run"
+
+    def run(
+        self, 
+        reference_image, 
+        reference_image_text,
+        model: comfy.model_patcher.ModelPatcher, 
+        clip,
+        vae,
+        positive,
+        negative, 
+        enabled,
+        denoise,
+        batch_size=1
+    ):
+        
+        tokens = clip.tokenize(reference_image_text)
+        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
+        reference_image_prompt=[[cond, {"pooled_output": pooled}]]
+   
+        reference_image = reference_image.repeat(((batch_size+1)//2, 1,1,1))
+
+        self.model = model
+        reference_latent = vae.encode(reference_image[:,:,:,:3])
+        
+        for n, m in model.model.diffusion_model.named_modules():
+            if m.__class__.__name__  == "CrossAttention":
+                processor = VisualStyleProcessor(m, enabled=enabled)
+                setattr(m, 'forward', processor.visual_style_forward)
+
+        conditioning_prompt = reference_image_prompt + positive
+        negative_prompt = negative * 2 
+
+        latents = torch.zeros_like(reference_latent) 
+        latents = torch.cat([latents] * 2)
+
+        if denoise < 1.0:
+            latents[::1] = reference_latent[:1]
+        else:
+            latents[::2] = reference_latent
+
+        denoise_mask = torch.ones_like(latents)[:, :1, ...] * denoise
+
+        denoise_mask[0] = 0.
+
+        return (model, conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask})    
+
+
+
+
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d
+
+
+class StyleAlignedArgs:
+    def __init__(self, share_attn: str) -> None:
+        self.adain_keys = "k" in share_attn
+        self.adain_values = "v" in share_attn
+        self.adain_queries = "q" in share_attn
+
+    share_attention: bool = True
+    adain_queries: bool = True
+    adain_keys: bool = True
+    adain_values: bool = True
+
+
+def expand_first(
+    feat: T,
+    scale=1.0,
+) -> T:
+    """
+    Expand the first element so it has the same shape as the rest of the batch.
+    """
+    b = feat.shape[0]
+    feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1)
+    if scale == 1:
+        feat_style = feat_style.expand(2, b // 2, *feat.shape[1:])
+    else:
+        feat_style = feat_style.repeat(1, b // 2, 1, 1, 1)
+        feat_style = torch.cat([feat_style[:, :1], scale * feat_style[:, 1:]], dim=1)
+    return feat_style.reshape(*feat.shape)
+
+
+def concat_first(feat: T, dim=2, scale=1.0) -> T:
+    """
+    concat the the feature and the style feature expanded above
+    """
+    feat_style = expand_first(feat, scale=scale)
+    return torch.cat((feat, feat_style), dim=dim)
+
+
+def calc_mean_std(feat, eps: float = 1e-5) -> "tuple[T, T]":
+    feat_std = (feat.var(dim=-2, keepdims=True) + eps).sqrt()
+    feat_mean = feat.mean(dim=-2, keepdims=True)
+    return feat_mean, feat_std
+
+def adain(feat: T) -> T:
+    feat_mean, feat_std = calc_mean_std(feat)
+    feat_style_mean = expand_first(feat_mean)
+    feat_style_std = expand_first(feat_std)
+    feat = (feat - feat_mean) / feat_std
+    feat = feat * feat_style_std + feat_style_mean
+    return feat
+
+class SharedAttentionProcessor:
+    def __init__(self, args: StyleAlignedArgs, scale: float):
+        self.args = args
+        self.scale = scale
+
+    def __call__(self, q, k, v, extra_options):
+        if self.args.adain_queries:
+            q = adain(q)
+        if self.args.adain_keys:
+            k = adain(k)
+        if self.args.adain_values:
+            v = adain(v)
+        if self.args.share_attention:
+            k = concat_first(k, -2, scale=self.scale)
+            v = concat_first(v, -2)
+
+        return q, k, v
+
+
+def get_norm_layers(
+    layer: nn.Module,
+    norm_layers_: "dict[str, list[Union[nn.GroupNorm, nn.LayerNorm]]]",
+    share_layer_norm: bool,
+    share_group_norm: bool,
+):
+    if isinstance(layer, nn.LayerNorm) and share_layer_norm:
+        norm_layers_["layer"].append(layer)
+    if isinstance(layer, nn.GroupNorm) and share_group_norm:
+        norm_layers_["group"].append(layer)
+    else:
+        for child_layer in layer.children():
+            get_norm_layers(
+                child_layer, norm_layers_, share_layer_norm, share_group_norm
+            )
+
+
+def register_norm_forward(
+    norm_layer: Union[nn.GroupNorm, nn.LayerNorm],
+) -> Union[nn.GroupNorm, nn.LayerNorm]:
+    if not hasattr(norm_layer, "orig_forward"):
+        setattr(norm_layer, "orig_forward", norm_layer.forward)
+    orig_forward = norm_layer.orig_forward
+
+    def forward_(hidden_states: T) -> T:
+        n = hidden_states.shape[-2]
+        hidden_states = concat_first(hidden_states, dim=-2)
+        hidden_states = orig_forward(hidden_states)  # type: ignore
+        return hidden_states[..., :n, :]
+
+    norm_layer.forward = forward_  # type: ignore
+    return norm_layer
+
+
+def register_shared_norm(
+    model: ModelPatcher,
+    share_group_norm: bool = True,
+    share_layer_norm: bool = True,
+):
+    norm_layers = {"group": [], "layer": []}
+    get_norm_layers(model.model, norm_layers, share_layer_norm, share_group_norm)
+    print(
+        f"Patching {len(norm_layers['group'])} group norms, {len(norm_layers['layer'])} layer norms."
+    )
+    return [register_norm_forward(layer) for layer in norm_layers["group"]] + [
+        register_norm_forward(layer) for layer in norm_layers["layer"]
+    ]
+
+
+SHARE_NORM_OPTIONS = ["both", "group", "layer", "disabled"]
+SHARE_ATTN_OPTIONS = ["q+k", "q+k+v", "disabled"]
+
+class StyleAlignedSampleReferenceLatents:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    { 
+                    "reference_image": ("IMAGE",),
+                    "positive": ("CONDITIONING",),
+                    "negative": ("CONDITIONING", ),
+                    "model": ("MODEL",),
+                    "vae": ("VAE", ),
+                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
+
+                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS.reverse(), ), 
+                    "denoise": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0, "step": 0.01}),
+                    
+                     }
+                }
+
+    RETURN_TYPES = ("STEP_LATENTS","LATENT")
+    RETURN_NAMES = ("ref_latents", "noised_output")
+
+    FUNCTION = "run"
+
+    # CATEGORY = "style_aligned"
+    CATEGORY = "♾️Mixlab/Style"
+
+    def run(self, reference_image, positive, negative, model, vae, seed, steps, cfg,scheduler,denoise):
+
+        # TODO noise_mask?
+        def vae_encode_crop_pixels(pixels):
+            x = (pixels.shape[1] // 8) * 8
+            y = (pixels.shape[2] // 8) * 8
+            if pixels.shape[1] != x or pixels.shape[2] != y:
+                x_offset = (pixels.shape[1] % 8) // 2
+                y_offset = (pixels.shape[2] % 8) // 2
+                pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
+            return pixels
+
+        pixels=vae_encode_crop_pixels(reference_image)
+        t = vae.encode(pixels[:,:,:,:3])
+        latent_image =  {"samples":t}
+
+        noise_seed=seed
+
+        sampler_name="ddim"
+        
+        sampler = comfy.samplers.sampler_object(sampler_name)
+
+        total_steps = steps
+        if denoise < 1.0:
+            total_steps = int(steps/denoise)
+
+        comfy.model_management.load_models_gpu([model])
+        sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, scheduler, total_steps).cpu()
+        sigmas = sigmas[-(steps + 1):]
+
+        sigmas = sigmas.flip(0)
+        if sigmas[0] == 0:
+            sigmas[0] = 0.0001
+
+        latent = latent_image
+        latent_image = latent["samples"]
+        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+
+    
+        noise_mask = None
+        if "noise_mask" in latent:
+            noise_mask = latent["noise_mask"]
+
+        ref_latents = []
+        def callback(step: int, x0: T, x: T, steps: int):
+            ref_latents.insert(0, x[0])
+        
+        disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+        samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
+
+        out = latent.copy()
+        out["samples"] = samples
+        out_noised = out
+
+        ref_latents = torch.stack(ref_latents)
+
+        return (ref_latents, out_noised)
+
+class StyleAlignedReferenceSampler:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+
+                "ref_latents": ("STEP_LATENTS",),
+                "reference_image_text": ("STRING", {"multiline": True}),
+                "model": ("MODEL",),
+                "clip": ("CLIP", ),
+
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+
+                "share_norm": (SHARE_NORM_OPTIONS,),
+                "share_attn": (SHARE_ATTN_OPTIONS,),
+                "scale": ("FLOAT", {"default": 1, "min": 0, "max": 2.0, "step": 0.01}),
+                "batch_size": ("INT", {"default": 2, "min": 1, "max": 8, "step": 1}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), 
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+            },
+        }
+
+    RETURN_TYPES = ("LATENT", "LATENT")
+    RETURN_NAMES = ("output", "denoised_output")
+    FUNCTION = "patch"
+    # CATEGORY = "style_aligned"
+    CATEGORY = "♾️Mixlab/Style"
+    def patch(
+        self,
+        ref_latents,
+        reference_image_text,
+        model,
+        clip,
+        positive,
+        negative,
+        share_norm,
+        share_attn,
+        scale,
+        batch_size,
+        seed,steps,cfg,scheduler,denoise
+
+    ) -> "tuple[dict, dict]":
+        
+        m = model.clone()
+
+        # ref_latents = vae.encode(reference_image[:,:,:,:3])
+
+        tokens = clip.tokenize(reference_image_text)
+        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
+        ref_positive=[[cond, {"pooled_output": pooled}]]
+
+        noise_seed=seed
+
+
+        total_steps = steps
+        if denoise < 1.0:
+            total_steps = int(steps/denoise)
+
+        # comfy.model_management.load_models_gpu([model])
+        sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, scheduler, total_steps).cpu()
+        sigmas = sigmas[-(steps + 1):]
+
+        sampler_name="ddim"
+        
+        sampler = comfy.samplers.sampler_object(sampler_name)
+
+        args = StyleAlignedArgs(share_attn)
+
+        # Concat batch with style latent
+        style_latent_tensor = ref_latents[0].unsqueeze(0)
+        height, width = style_latent_tensor.shape[-2:]
+        latent_t = torch.zeros(
+            [batch_size, 4, height, width], device=ref_latents.device
+        )
+        latent = {"samples": latent_t}
+        noise = comfy.sample.prepare_noise(latent_t, noise_seed)
+
+        latent_t = torch.cat((style_latent_tensor, latent_t), dim=0)
+        ref_noise = torch.zeros_like(noise[0]).unsqueeze(0)
+        noise = torch.cat((ref_noise, noise), dim=0)
+
+        x0_output = {}
+        preview_callback = latent_preview.prepare_callback(m, sigmas.shape[-1] - 1, x0_output)
+
+        # Replace first latent with the corresponding reference latent after each step
+        def callback(step: int, x0: T, x: T, steps: int):
+            preview_callback(step, x0, x, steps)
+            if (step + 1 < steps):
+                # 当ref_latents的step不够时
+                if step+1>len(ref_latents)-1:
+                    step=len(ref_latents)-2
+                
+                x[0] = ref_latents[step+1]
+                x0[0] = ref_latents[step+1]
+
+        # Register shared norms
+        share_group_norm = share_norm in ["group", "both"]
+        share_layer_norm = share_norm in ["layer", "both"]
+        register_shared_norm(m, share_group_norm, share_layer_norm)
+
+        # Patch cross attn
+        m.set_model_attn1_patch(SharedAttentionProcessor(args, scale))
+
+        # Add reference conditioning to batch 
+        batched_condition = []
+        for i,condition in enumerate(positive):
+            additional = condition[1].copy()
+            batch_with_reference = torch.cat([ref_positive[i][0], condition[0].repeat([batch_size] + [1] * len(condition[0].shape[1:]))], dim=0)
+            if 'pooled_output' in additional and 'pooled_output' in ref_positive[i][1]:
+                # combine pooled output
+                pooled_output = torch.cat([ref_positive[i][1]['pooled_output'], additional['pooled_output'].repeat([batch_size] 
+                    + [1] * len(additional['pooled_output'].shape[1:]))], dim=0)
+                additional['pooled_output'] = pooled_output
+            if 'control' in additional:
+                if 'control' in ref_positive[i][1]:
+                    # combine control conditioning
+                    control_hint = torch.cat([ref_positive[i][1]['control'].cond_hint_original, additional['control'].cond_hint_original.repeat([batch_size] 
+                        + [1] * len(additional['control'].cond_hint_original.shape[1:]))], dim=0)
+                    cloned_controlnet = additional['control'].copy()
+                    cloned_controlnet.set_cond_hint(control_hint, strength=additional['control'].strength, timestep_percent_range=additional['control'].timestep_percent_range)
+                    additional['control'] = cloned_controlnet
+                else:
+                    # add zeros for first in batch
+                    control_hint = torch.cat([torch.zeros_like(additional['control'].cond_hint_original), additional['control'].cond_hint_original.repeat([batch_size] 
+                        + [1] * len(additional['control'].cond_hint_original.shape[1:]))], dim=0)
+                    cloned_controlnet = additional['control'].copy()
+                    cloned_controlnet.set_cond_hint(control_hint, strength=additional['control'].strength, timestep_percent_range=additional['control'].timestep_percent_range)
+                    additional['control'] = cloned_controlnet
+            batched_condition.append([batch_with_reference, additional])
+
+        disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+        samples = comfy.sample.sample_custom(
+            m,
+            noise,
+            cfg,
+            sampler,
+            sigmas,
+            batched_condition,
+            negative,
+            latent_t,
+            callback=callback,
+            disable_pbar=disable_pbar,
+            seed=noise_seed,
+        )
+
+        # remove reference image
+        samples = samples[1:]
+
+        out = latent.copy()
+        out["samples"] = samples
+        if "x0" in x0_output:
+            out_denoised = latent.copy()
+            x0 = x0_output["x0"][1:]
+            out_denoised["samples"] = m.model.process_latent_out(x0.cpu())
+        else:
+            out_denoised = out
+        return (out, out_denoised)
+
+
+class StyleAlignedBatchAlign:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "share_norm": (SHARE_NORM_OPTIONS,),
+                "share_attn": (SHARE_ATTN_OPTIONS,),
+                "scale": ("FLOAT", {"default": 1, "min": 0, "max": 1.0, "step": 0.1}),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+    # CATEGORY = "style_aligned"
+    CATEGORY = "♾️Mixlab/Style"
+    def patch(
+        self,
+        model: ModelPatcher,
+        share_norm: str,
+        share_attn: str,
+        scale: float,
+    ):
+        m = model.clone()
+        share_group_norm = share_norm in ["group", "both"]
+        share_layer_norm = share_norm in ["layer", "both"]
+        register_shared_norm(model, share_group_norm, share_layer_norm)
+        args = StyleAlignedArgs(share_attn)
+        m.set_model_attn1_patch(SharedAttentionProcessor(args, scale))
+        return (m,)
+
+

nodes/TextGenerateNode.py

@@ -12,23 +12,31 @@ import comfy.utils
 # import numpy as np
 import torch
 import random
+from lark import Lark, Transformer, v_args
+
 
 global _available
 _available=True
 
+def get_text_generator_path():
+    try:
+        return folder_paths.get_folder_paths('prompt_generator')[0]
+    except:
+        return os.path.join(folder_paths.models_dir, "prompt_generator")
 
-text_generator_model_path=os.path.join(folder_paths.models_dir, "prompt_generator/text2image-prompt-generator")
+prompt_generator=get_text_generator_path()
+
+text_generator_model_path=os.path.join(prompt_generator, "text2image-prompt-generator")
 if not os.path.exists(text_generator_model_path):
     print(f"## text_generator_model not found: {text_generator_model_path}, pls download from https://huggingface.co/succinctly/text2image-prompt-generator/tree/main")
     text_generator_model_path='succinctly/text2image-prompt-generator'
 
-zh_en_model_path=os.path.join(folder_paths.models_dir, "prompt_generator/opus-mt-zh-en")
+zh_en_model_path=os.path.join(prompt_generator, "opus-mt-zh-en")
 if not os.path.exists(zh_en_model_path):
     print(f"## zh_en_model not found: {zh_en_model_path}, pls download from https://huggingface.co/Helsinki-NLP/opus-mt-zh-en/tree/main")
     zh_en_model_path='Helsinki-NLP/opus-mt-zh-en'
 
 
-
 def is_installed(package):
     try:
         spec = importlib.util.find_spec(package)
@@ -62,7 +70,14 @@ except:
 
 
 
-def translate(zh_en_tokenizer,zh_en_model,text):
+def translate(text):
+    global text_pipe,zh_en_model,zh_en_tokenizer
+    
+    if zh_en_model==None:
+        zh_en_model = AutoModelForSeq2SeqLM.from_pretrained(zh_en_model_path).eval()
+        zh_en_tokenizer = AutoTokenizer.from_pretrained(zh_en_model_path,padding=True, truncation=True)
+    
+    zh_en_model.to("cuda" if torch.cuda.is_available() else "cpu")
     with torch.no_grad():
         encoded = zh_en_tokenizer([text], return_tensors="pt")
         encoded.to(zh_en_model.device)
@@ -102,18 +117,24 @@ def text_generate(text_pipe,input,seed=None):
     
 import re
 
-def correct_prompt_syntax(prompt):
+def correct_prompt_syntax(prompt=""):
 
-    print("input prompt",prompt)
+    # print("input prompt",prompt)
     corrected_elements = []
     # 处理成统一的英文标点
     prompt = prompt.replace('（', '(').replace('）', ')').replace('，', ',').replace(';', ',').replace('。', '.').replace('：',':')
     # 删除多余的空格
     prompt = re.sub(r'\s+', ' ', prompt).strip()
+    prompt = prompt.replace("< ","<").replace(" >",">").replace("( ","(").replace(" )",")").replace("[ ","[").replace(' ]',']')
 
     # 分词
     prompt_elements = prompt.split(',')
 
+    def balance_brackets(element, open_bracket, close_bracket):
+        open_brackets_count = element.count(open_bracket)
+        close_brackets_count = element.count(close_bracket)
+        return element + close_bracket * (open_brackets_count - close_brackets_count)
+
     for element in prompt_elements:
         element = element.strip()
 
@@ -133,21 +154,118 @@ def correct_prompt_syntax(prompt):
         corrected_elements.append(corrected_element)
 
     # 重组修正后的prompt
-    corrected_prompt = ', '.join(corrected_elements)
-    print("output prompt",corrected_prompt)
-    return corrected_prompt
+    return  ','.join(corrected_elements)
 
-def balance_brackets(element, open_bracket, close_bracket):
-    open_brackets_count = element.count(open_bracket)
-    close_brackets_count = element.count(close_bracket)
-    return element + close_bracket * (open_brackets_count - close_brackets_count)
 
 # # 示例使用
 # test_prompt = "((middle-century castles)), [forsaken: 0.8], (mystery dragons: 1.3, mist forests, sunsets, quiet; (((dummy)), [fisting city: 0.5] background, radiant, soft and flavoured,] promising mountains, ((starry: 1.6), [[crowds], [middle-century castle: urban landscapes of the future: 0.5], [yellow: bright sun: 0.7], overlooking"
 # corrected_prompt = correct_prompt_syntax(test_prompt)
 # print(corrected_prompt)
 
+def detect_language(input_str):
+    # 统计中文和英文字符的数量
+    count_cn = count_en = 0
+    for char in input_str:
+        if '\u4e00' <= char <= '\u9fff':
+            count_cn += 1
+        elif char.isalpha():
+            count_en += 1
+
+    # 根据统计的字符数量判断主要语言
+    if count_cn > count_en:
+        return "cn"
+    elif count_en > count_cn:
+        return "en"
+    else:
+        return "unknow"
+
+
+
+
+#定义Prompt文法
+grammar = """
+start: sentence
+sentence: phrase ("," phrase)*
+phrase: emphasis | weight | word | lora | embedding | schedule
+emphasis: "(" sentence ")" -> emphasis
+        | "[" sentence "]" -> weak_emphasis
+weight: "(" word ":" NUMBER ")"
+schedule: "[" word ":" word ":" NUMBER "]"
+lora: "<" WORD ":" WORD (":" NUMBER)? (":" NUMBER)? ">"
+embedding: "embedding" ":" WORD (":" NUMBER)? (":" NUMBER)?
+word: WORD
+
+NUMBER: /\s*-?\d+(\.\d+)?\s*/
+WORD: /[^,:\(\)\[\]<>]+/
+"""
+
+
+
+@v_args(inline=True)  # Decorator to flatten the tree directly into the function arguments
+class ChinesePromptTranslate(Transformer):
+ 
+    def sentence(self, *args):
+        return ", ".join(args)
+
+    def phrase(self, *args):
+        return "".join(args)
+
+    def emphasis(self, *args):
+        # Reconstruct the emphasis with translated content
+        return "(" + "".join(args) + ")"
+
+    def weak_emphasis(self, *args):
+        print('weak_emphasis:',args)
+        return "[" + "".join(args) + "]"
+    
+    def embedding(self,*args):
+        print('prompt embedding',args[0])
+        if len(args) == 1:
+            # print('prompt embedding',str(args[0]))
+            # 只传递了一个参数，意味着只有embedding名称没有数字
+            embedding_name = str(args[0])
+            return f"embedding:{embedding_name}"
+        elif len(args) > 1:
+            embedding_name,*numbers = args
+           
+            if len(numbers)==2:
+                return f"embedding:{embedding_name}:{numbers[0]}:{numbers[1]}"
+            elif len(numbers)==1:
+                return f"embedding:{embedding_name}:{numbers[0]}"
+            else:
+                return f"embedding:{embedding_name}"
+
+    def lora(self,*args):
+        print('lora prompt',*args)
+        if len(args) == 1:
+            return f"<lora:{loar_name}>"
+        elif len(args) > 1: 
+            # print('lora', args)
+            _,loar_name,*numbers = args
+            loar_name = str(loar_name).strip()
+            if len(numbers)==2:
+                return f"<lora:{loar_name}:{numbers[0]}:{numbers[1]}>"
+            elif len(numbers)==1:
+                return f"<lora:{loar_name}:{numbers[0]}>"
+            else:
+                return f"<lora:{loar_name}>"
+    
+    def weight(self, word,number):
+        translated_word = translate(str(word)).rstrip('.')
+        return f"({translated_word}:{str(number).strip()})"
+    
+    def schedule(self,*args):
+        print('prompt schedule',args)
+        data = [str(arg).strip() for arg in args] 
         
+        return f"[{':'.join(data)}]"
+
+    def word(self, word):
+        # Translate each word using the dictionary
+        if detect_language(str(word)) == "cn":
+            return  translate(str(word)).rstrip('.')
+        else:
+            return str(word).rstrip('.')
         
 class ChinesePrompt:
 
@@ -185,16 +303,16 @@ class ChinesePrompt:
     zh_en_tokenizer=None 
 
     def run(self,text,seed,generation):
-        global text_pipe,zh_en_model,zh_en_tokenizer
+        
 
         seed=seed[0]
         generation=generation[0]
 
         # 进度条
         pbar = comfy.utils.ProgressBar(len(text)+1)
-        
         texts = [correct_prompt_syntax(t) for t in text]
-        print('correct_prompt_syntax::',texts)
+
+        global text_pipe,zh_en_model,zh_en_tokenizer
         if zh_en_model==None:
             zh_en_model = AutoModelForSeq2SeqLM.from_pretrained(zh_en_model_path).eval()
             zh_en_tokenizer = AutoTokenizer.from_pretrained(zh_en_model_path,padding=True, truncation=True)
@@ -210,9 +328,16 @@ class ChinesePrompt:
 
         # print('zh_en_model device',zh_en_model.device,text_pipe.model.device,torch.cuda.current_device() )
         en_texts=[]
+
         for t in texts:
-            en_text=translate(zh_en_tokenizer,zh_en_model,t)
-            en_texts.append(en_text)
+            if t:
+                # translated_text =  translated_word = translate(zh_en_tokenizer,zh_en_model,str(t))
+                parser = Lark(grammar, start="start", parser="lalr", transformer=ChinesePromptTranslate())
+                # print('t',t)
+                result = parser.parse(t).children
+                # print('en_result',result)        
+                # en_text=translate(zh_en_tokenizer,zh_en_model,text_without_syntax)
+                en_texts.append(result[0])
 
         zh_en_model.to('cpu')
         print("test en_text",en_texts)
@@ -232,8 +357,11 @@ class ChinesePrompt:
             pbar.update(1)
 
         text_pipe.model.to('cpu')
-        prompt_result = [correct_prompt_syntax(p) for p in prompt_result]
-        
+
+        print('prompt_result',prompt_result,)
+        # prompt_result = [','.join(correct_prompt_syntax(p)) for p in prompt_result]
+        if len(prompt_result)==0:
+            prompt_result=[""]
         return {
             "ui":{
                     "prompt": prompt_result
@@ -241,6 +369,8 @@ class ChinesePrompt:
             "result": (prompt_result,)}
 
 
+
+
 class PromptGenerate:
 
     global _available

nodes/TripoSR.py

@@ -0,0 +1,180 @@
+import sys
+from os import path
+sys.path.insert(0, path.dirname(__file__))
+from PIL import Image
+import numpy as np
+import torch
+
+from folder_paths import get_folder_paths, get_full_path, get_save_image_path, get_output_directory,models_dir
+from comfy.model_management import get_torch_device
+from .tsr.system import TSR
+
+import comfy.utils
+
+
+def get_triposr_model_path():
+    try:
+        return path.join(get_folder_paths('triposr')[0],'model.ckpt')
+    except:
+        return path.join(path.join(models_dir, "triposr"),'model.ckpt')
+    
+triposr_model_path=get_triposr_model_path()
+
+
+# Tensor to PIL
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+# Convert PIL to Tensor
+def pil2tensor(image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+
+def fill_background(image):
+    im = np.array(image).astype(np.float32) / 255.0
+    im = im[:, :, :3] * im[:, :, 3:4] + (1 - im[:, :, 3:4]) * 0.5
+    im = Image.fromarray((im * 255.0).astype(np.uint8))
+    return im
+
+
+class LoadTripoSRModel:
+    def __init__(self):
+        self.initialized_model = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                # "model": (get_filename_list("checkpoints"),),
+                "chunk_size": ("INT", {"default": 8192, "min": 0, "max": 10000})
+            }
+        }
+
+    RETURN_TYPES = ("TRIPOSR_MODEL",)
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/3D/TripoSR"
+
+    def run(self, chunk_size):
+        device = get_torch_device()
+
+        if not torch.cuda.is_available():
+            device = "cpu"
+
+        if not self.initialized_model:
+            # triposr_model_path
+            print("#Loading TripoSR model",triposr_model_path)
+            self.initialized_model = TSR.from_pretrained_custom(
+                weight_path=triposr_model_path,
+                config_path=path.join(path.dirname(__file__), "tsr/config.yaml")
+            )
+            self.initialized_model.renderer.set_chunk_size(chunk_size)
+            self.initialized_model.to(device)
+
+        return (self.initialized_model,)
+
+
+class TripoSRSampler:
+    def __init__(self):
+        self.initialized_model = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("TRIPOSR_MODEL",),
+                "image": ("IMAGE",),
+                "resolution": ("INT", {"default": 256, "min": 128, "max": 12288}),
+                "threshold": ("FLOAT", {"default": 25.0, "min": 0.0, "step": 0.01}),
+                "device":(["auto","cpu"],), 
+            },
+            "optional": {
+                "mask": ("MASK",)
+            }
+        }
+
+    RETURN_TYPES = ("MESH",)
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/3D/TripoSR"
+
+    def run(self, model, image, resolution, threshold,device='auto', mask=None):
+
+        reference_image=image
+        reference_mask=mask
+
+        device = get_torch_device()
+
+        if not torch.cuda.is_available():
+            device = "cpu"
+
+        if device=='cpu':
+            device = "cpu"
+
+        print('#TripoSRSampler device',device)
+
+        to_images=[]
+
+        for i in range(len(reference_image)):
+
+            image = reference_image[i]
+            
+            if reference_mask is not None:
+                mask = reference_mask[i].unsqueeze(2)
+                image = torch.cat((image, mask), dim=2).detach().cpu().numpy()
+                image = Image.fromarray(np.clip(255. * image, 0, 255).astype(np.uint8))
+                image = fill_background(image)
+            else:
+                image = tensor2pil(image)
+
+            image = image.convert('RGB')
+
+            to_images.append(image)
+
+        # 进度条
+        pbar = comfy.utils.ProgressBar(len(to_images))
+        def callback(c):
+            pbar.update(1)    
+
+        scene_codes = model(to_images, device)
+        meshes = model.extract_mesh(scene_codes, resolution=resolution, threshold=threshold,callback=callback)
+
+        del model
+        return (meshes,)
+
+
+class SaveTripoSRMesh:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mesh": ("MESH",),
+                # "format":(["glb","obj"],), 
+                "filename_prefix":("STRING", {"multiline": False,"default": "TripoSR_"})
+            }
+        }
+
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/3D/TripoSR"
+
+    def run(self, mesh,filename_prefix):
+        format='glb'
+        saved = list()
+        full_output_folder, filename, counter, subfolder, filename_prefix = get_save_image_path(filename_prefix,
+                                                                                                get_output_directory())
+
+        for (index, single_mesh) in enumerate(mesh):
+            filename_with_batch_num = filename.replace("%batch_num%", str(index))
+            file = f"{filename_with_batch_num}_{counter:05}_.{format}"
+            single_mesh.apply_transform(np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]]))
+            single_mesh.export(path.join(full_output_folder, file))
+            saved.append({
+                "filename": file,
+                "type": "output",
+                "subfolder": subfolder
+            })
+
+        return {"ui": {"mesh": saved}}
+
+
+

nodes/Utils.py

@@ -6,8 +6,20 @@ import numpy as np
 import folder_paths
 import matplotlib.font_manager as fm
 import torch
+import importlib.util
 
+def create_incrementing_list(min_value, max_value, step, count):
+    l1 = [int(min_value + i * step) for i in range(count) if min_value + i * step <= max_value]
+    l2 = [float(min_value + i * step) for i in range(count) if min_value + i * step <= max_value]
+    return (l1,l2)
 
+def split_list(lst, chunk_size, transition_size):
+    result = []
+    for i in range(0, len(lst), chunk_size):
+        start = i - transition_size
+        end = i + chunk_size + transition_size
+        result.append(lst[max(start, 0):end])
+    return result
 
 def recursive_search(directory, excluded_dir_names=None):
     if not os.path.isdir(directory):
@@ -146,7 +158,6 @@ class ColorInput:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { 
-             
                     "color":("TCOLOR",), 
                              },
                 }
@@ -156,7 +167,7 @@ class ColorInput:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Color"
 
     INPUT_IS_LIST = False
     OUTPUT_IS_LIST = (False,False,False,False,False,)
@@ -185,7 +196,7 @@ class FontInput:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Input"
 
     INPUT_IS_LIST = False
     OUTPUT_IS_LIST = (False,)
@@ -218,7 +229,7 @@ class TextToNumber:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Text"
 
     INPUT_IS_LIST = False
     OUTPUT_IS_LIST = (False,)
@@ -273,10 +284,10 @@ class FloatSlider:
                 }
     
     RETURN_TYPES = ("FLOAT",) 
-
+    RETURN_NAMES = ('FLOAT',)
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Input"
 
     INPUT_IS_LIST = False
     OUTPUT_IS_LIST = (False,)
@@ -286,9 +297,7 @@ class FloatSlider:
             number = min_value
         elif number > max_value:
             number = max_value
-        scaled_number = (number - min_value) / (max_value - min_value)
-        return (scaled_number,)
-    
+        return (number,)    
   
 class IntNumber:
     @classmethod
@@ -329,7 +338,7 @@ class IntNumber:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Input"
 
     INPUT_IS_LIST = False
     OUTPUT_IS_LIST = (False,)
@@ -347,7 +356,7 @@ class MultiplicationNode:
         return {"required": {
                     "numberA":(any_type,),
                     "multiply_by":("FLOAT", {
-                        "default": 0, 
+                        "default": 1, 
                         "min": -2, #Minimum value
                         "max": 0xffffffffffffffff,
                         "step": 0.01, #Slider's step
@@ -389,7 +398,7 @@ class TextInput:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Input"
 
     INPUT_IS_LIST = False
     OUTPUT_IS_LIST = (False,)
@@ -398,6 +407,61 @@ class TextInput:
        
         return (text,)
     
+
+class IncrementingListNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "min_value": ("FLOAT", {
+                        "default": 0, 
+                        "min": -2000, #Minimum value
+                        "max": 0xffffffffffffffff,
+                        "step": 0.01, #Slider's step
+                        "display": "number" # Cosmetic only: display as "number" or "slider"
+                    }),
+                    "max_value": ("FLOAT", {
+                        "default": 10, 
+                        "min": -2000, #Minimum value
+                        "max": 0xffffffffffffffff,
+                        "step": 0.01, #Slider's step
+                        "display": "number" # Cosmetic only: display as "number" or "slider"
+                    }),
+                    "step": ("FLOAT", {
+                        "default": 0, 
+                        "min": -2000, #Minimum value
+                        "max": 0xffffffffffffffff,
+                        "step": 0.01, #Slider's step
+                        "display": "number" # Cosmetic only: display as "number" or "slider"
+                    }),
+                    "count": ("INT", {
+                        "default": 1, 
+                        "min": 1, #Minimum value
+                        "max": 0xffffffffffffffff,
+                        "step":1, #Slider's step
+                        "display": "number" # Cosmetic only: display as "number" or "slider"
+                    })
+                             },
+                             "optional":{
+                    "seed":("INT", {"default": -1, "min": -1, "max": 1000000}), 
+                    
+                },
+                }
+    
+    RETURN_TYPES = ("INT","FLOAT",) 
+    RETURN_NAMES = ('int_list','float_list',)
+
+    FUNCTION = "run"
+
+    CATEGORY = "♾️Mixlab/Video"
+
+    INPUT_IS_LIST = False
+    OUTPUT_IS_LIST = (True,True,)
+
+    def run(self,min_value,max_value,step,count,seed):
+        print('create_incrementing_list',seed)
+        l1,l2=create_incrementing_list(min_value,max_value,step,count)
+        return (l1,l2,)
+    
 # 接收一个值，然后根据字符串或数值长度计算延迟时间，用户可以自定义延迟"字/s"，延迟之后将转化
 
 import comfy.samplers
@@ -560,9 +624,11 @@ class SwitchByIndex:
     @classmethod
     def INPUT_TYPES(cls):
         return {
-            "required": {
-                "A":(any_type,),
-                "B":(any_type,),
+             "optional":{
+                  "A":(any_type,),
+                  "B":(any_type,),
+            },
+            "required": { 
                 "index":("INT", {
                         "default": -1, 
                         "min": -1, 
@@ -574,17 +640,17 @@ class SwitchByIndex:
             }
         }
 
-    RETURN_TYPES = (any_type,)
-    RETURN_NAMES = ("C",)
+    RETURN_TYPES = (any_type,"INT",)
+    RETURN_NAMES = ("list", "count",)
 
     FUNCTION = "run"
 
     CATEGORY = "♾️Mixlab/Utils"
 
     INPUT_IS_LIST = True
-    OUTPUT_IS_LIST = (True,)
+    OUTPUT_IS_LIST = (True, False,)
 
-    def run(self, A,B,index,flat):
+    def run(self, A=[],B=[],index=-1,flat='on'):
 
         flat=flat[0]
 
@@ -603,10 +669,43 @@ class SwitchByIndex:
             try:
                 C=[C[index]]
             except Exception as e:
-                C=[]
-        
-        return (C,)
+                C=[C[-1]] #最后一个
 
+        return (C, len(C),)
+
+class ListSplit:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+             "optional":{
+                  "A":(any_type,), 
+            },
+            "required": { 
+                "chunk_size": ("INT", {"default": 10, "min": 1, "step": 1}),
+                "transition_size": ("INT", {"default": 0, "min": 0, "step": 1}),
+                "index": ("INT", {"default": -1, "min": -1, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = (any_type,)
+    RETURN_NAMES = ("B",)
+
+    FUNCTION = "run"
+
+    CATEGORY = "♾️Mixlab/Utils"
+
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True,)
+
+    def run(self, A=[],chunk_size=[10],transition_size=[0],index=[-1]):
+        # print(len(A))
+        B=split_list(A,chunk_size[0],transition_size[0])
+
+        if index[0]>-1:
+            B=B[index[0]]
+
+        return (B,)
+    
 
 
 class LimitNumber:
@@ -637,7 +736,7 @@ class LimitNumber:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Input"
 
     INPUT_IS_LIST = False
     OUTPUT_IS_LIST = (False,)
@@ -704,19 +803,40 @@ class TESTNODE_:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/__TEST"
+    CATEGORY = "♾️Mixlab/Test"
 
     OUTPUT_NODE = True
     INPUT_IS_LIST = True
     OUTPUT_IS_LIST = (True,)
 
     def run(self,ANY):
-        print(ANY)
+        print(type(ANY))
+        try:
+            print(ANY[0].shape)
+            img= tensor2pil(ANY[0])
+            print(img.size)
+        except:
+            print('')
+        
         # data=ANY
         list_stats = ListStatistics()
 
         # 调用count_types方法进行统计
         result = list_stats.count_types(ANY)
+ 
+
+        # 假设我们有一个模块文件名为 my_module.py，它位于 'importables' 目录下
+        module_path = os.path.join(os.path.dirname(__file__),'test.py')
+
+        # 使用 spec_from_file_location 获取模块的元数据（名称、定义等）
+        spec = importlib.util.spec_from_file_location('test', module_path)
+
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+
+        functions = getattr(module, 'run')  # 获取函数
+        
+        functions(ANY)
 
 
         return {"ui": {"data": result,"type":[str(type(ANY[0]))]}, "result": (ANY,)}
@@ -735,7 +855,7 @@ class TESTNODE_TOKEN:
 
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/__TEST"
+    CATEGORY = "♾️Mixlab/Test"
 
     OUTPUT_NODE = True
     INPUT_IS_LIST = False
@@ -772,7 +892,7 @@ class CreateSeedNode:
     OUTPUT_NODE = True
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Experiment"
 
     def run(self, seed):
         return (seed,)
@@ -799,7 +919,7 @@ class CreateCkptNames:
     # OUTPUT_NODE = True
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Experiment"
 
     def run(self, ckpt_names):
         ckpt_names=ckpt_names.split('\n')
@@ -828,7 +948,7 @@ class CreateLoraNames:
     # OUTPUT_NODE = True
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Experiment"
 
     def run(self, lora_names):
         lora_names=lora_names.split('\n')
@@ -859,7 +979,7 @@ class CreateSampler_names:
     # OUTPUT_NODE = True
     FUNCTION = "run"
 
-    CATEGORY = "♾️Mixlab/Utils"
+    CATEGORY = "♾️Mixlab/Experiment"
 
     def run(self, sampler_names):
         sampler_names=sampler_names.split('\n')

nodes/Video.py

@@ -0,0 +1,693 @@
+import os
+import hashlib
+import json
+import subprocess
+import shutil
+import re
+import time,math
+import numpy as np
+from typing import List
+import torch
+from PIL import Image, ImageOps
+from PIL.PngImagePlugin import PngInfo
+import cv2,random,string
+from pathlib import Path
+
+import folder_paths
+from comfy.k_diffusion.utils import FolderOfImages
+from comfy.utils import common_upscale
+
+
+
+
+def generate_folder_name(directory,video_path):
+    # Get the directory and filename from the video path
+    _, filename = os.path.split(video_path)
+    # Generate a random string of lowercase letters and digits
+    random_string = ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))
+    # Create the folder name by combining the random string and the filename
+    folder_name = random_string + '_' + filename
+    # Create the full folder path by joining the directory and the folder name
+    folder_path = os.path.join(directory, folder_name)
+    return folder_path
+
+def create_folder(directory,video_path):
+    folder_path = generate_folder_name(directory,video_path)
+    os.makedirs(folder_path)
+    return folder_path
+
+
+def split_video(video_path, video_segment_frames, transition_frames, output_dir):
+    # 读取视频文件
+    video_capture = cv2.VideoCapture(video_path)
+    
+    # 获取视频的总帧数和帧率
+    total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = video_capture.get(cv2.CAP_PROP_FPS)
+    
+    # 计算每个视频片段的总帧数，包括过渡帧
+    segment_total_frames = video_segment_frames + transition_frames
+    
+    # 计算可以分割的片段数量，向上取整
+    num_segments = (total_frames + transition_frames - 1) // segment_total_frames
+    
+    vs=[]
+    # 计算每个片段的起始帧和结束帧
+    start_frame = 0
+    for i in range(num_segments):
+        # 计算当前片段的结束帧，注意最后一个片段可能没有过渡帧
+        end_frame = min(start_frame + segment_total_frames, total_frames)
+        
+        # 打印当前片段的起始帧和结束帧
+        print(f"Segment {i+1}: Start Frame {start_frame}, End Frame {end_frame}")
+        
+        # 保存当前片段为一个视频文件
+        segment_video_path = f"{output_dir}/segment_{i+1}.avi"
+
+        fourcc = cv2.VideoWriter_fourcc(*'XVID')
+        segment_video = cv2.VideoWriter(segment_video_path, fourcc, fps, (int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)),
+                                                              int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))))
+        
+        for frame_num in range(start_frame, end_frame):
+            ret, frame = video_capture.read()
+            if ret:
+                segment_video.write(frame)
+            else:
+                break  # 如果读取失败，则退出循环
+        
+        # 更新起始帧为下一个片段的起始位置
+        start_frame = end_frame + transition_frames
+        vs.append(segment_video_path)
+    
+    # 释放视频捕获对象
+    video_capture.release()
+    # print(vs)
+    return (vs,total_frames,fps)
+
+
+folder_paths.folder_names_and_paths["video_formats"] = (
+    [
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), ".", "video_formats"),
+    ],
+    [".json"]
+)
+
+ffmpeg_path = shutil.which("ffmpeg")
+if ffmpeg_path is None:
+    print("ffmpeg could not be found. Using ffmpeg from imageio-ffmpeg.")
+    from imageio_ffmpeg import get_ffmpeg_exe
+    try:
+        ffmpeg_path = get_ffmpeg_exe()
+    except:
+        print("ffmpeg could not be found. Outputs that require it have been disabled")
+
+# Tensor to PIL
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+# Convert PIL to Tensor
+def pil2tensor(image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+def count_files(directory):
+    count = 0
+    for root, dirs, files in os.walk(directory):
+        count += len(files)
+    return count
+
+def create_temp_file(image):
+    output_dir = folder_paths.get_temp_directory()
+
+    c=count_files(output_dir)
+
+    (
+            full_output_folder,
+            filename,
+            counter,
+            subfolder,
+            _,
+        ) = folder_paths.get_save_image_path('temp_', output_dir)
+
+    
+    image=tensor2pil(image)
+ 
+    image_file = f"{filename}_{c}_{counter:05}.png"
+     
+    image_path=os.path.join(full_output_folder, image_file)
+
+    image.save(image_path,compress_level=4)
+
+    return [{
+                "filename": image_file,
+                "subfolder": subfolder,
+                "type": "temp"
+                }]
+
+
+def split_list(lst, chunk_size, transition_size):
+    result = []
+    for i in range(0, len(lst), chunk_size):
+        start = i - transition_size
+        end = i + chunk_size + transition_size
+        result.append(lst[max(start, 0):end])
+    return result
+
+# images = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+# chunk_size = 3
+# transition_size = 1
+
+# result = split_list(images, chunk_size, transition_size)
+# print(result)
+
+class ImageListReplace:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "images": ("IMAGE",),
+            "start_index":("INT", {"default": 0, "min": 0, "step": 1}),
+            "end_index":("INT", {"default": 0, "min": 0, "step": 1}),
+            "invert": ("BOOLEAN", {"default": False}),
+        },
+         "optional":{
+                    "image_replace": ("IMAGE",),
+                    "images_replace": ("IMAGE",),
+                    
+                    }
+        }
+
+    RETURN_TYPES = ("IMAGE","IMAGE",)
+    RETURN_NAMES = ("images","select_images",)
+    FUNCTION = "run"
+    CATEGORY = "♾️Mixlab/Video"
+
+    OUTPUT_NODE = True
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True,True,)
+
+    def run(self, images,start_index=[0],end_index=[0],invert=[False],image_replace=None,images_replace=None):
+        start_index=start_index[0]
+        end_index=end_index[0]
+        invert=invert[0]
+
+        image_rs=[]
+
+        if image_replace!=None:
+            for i in range(end_index-start_index+1):
+                image_rs.append(image_replace[0])
+
+        if images_replace!=None:
+            image_rs=images_replace
+
+        # 如果image replace 为空
+        if image_replace==None and images_replace==None:
+            # print('如果image replace 为空',images[0])
+            # [[tensor(
+            # tensor([[[[0.
+            first_image=tensor2pil(images[0][0])
+            width, height = first_image.size
+            image_replace=Image.new("RGB", (width, height), (0, 0, 0))
+            image_replace=pil2tensor(image_replace)
+            for i in range(end_index-start_index+1):
+                image_rs.append(image_replace)
+
+        
+        new_images=[]
+        select_images=[]
+        k=0
+        for i in range(len(images)):
+            if i>=start_index and i<=end_index:
+                if invert:
+                    new_images.append(images[i])
+                else:
+                    new_images.append(image_rs[k])
+                    select_images.append(images[i])
+                    k+=1
+            else:
+                if invert:
+                    new_images.append(image_rs[k])
+                    select_images.append(images[i])
+                    k+=1
+                else:
+                    new_images.append(images[i])
+
+        imss=[]
+        # print(len(images))
+        for i in range(len(images)):
+            t=images[i][0]
+            t=tensor2pil(t)
+            t = t.convert("RGB")
+            original_width, original_height = t.size
+            scale = 300 / original_width
+            new_height = int(original_height * scale)
+            t = t.resize((300, new_height))
+
+            ims=create_temp_file(pil2tensor(t))
+            imss.append(ims[0])
+
+        # image_replace=create_temp_file(image_replace)
+
+        return {"ui":{"_images": imss},"result": (new_images,select_images,)}
+
+# The code is based on ComfyUI-VideoHelperSuite modification.
+class LoadVideoAndSegment:
+    @classmethod
+    def INPUT_TYPES(s):
+        video_extensions = ['webm', 'mp4', 'mkv', 'gif']
+        input_dir = folder_paths.get_input_directory()
+        files = []
+        for f in os.listdir(input_dir):
+            if os.path.isfile(os.path.join(input_dir, f)):
+                file_parts = f.split('.')
+                if len(file_parts) > 1 and (file_parts[-1] in video_extensions):
+                    files.append(f)
+        return {"required": {
+                    "video": (sorted(files), {"video_upload": True}),
+                     "video_segment_frames": ("INT", {"default": 10, "min": 1, "step": 1}),
+                     "transition_frames": ("INT", {"default": 0, "min": 0, "step": 1}), 
+                     },}
+
+    CATEGORY = "♾️Mixlab/Video"
+
+    RETURN_TYPES = ("SCENE_VIDEO","INT", "INT","INT",)
+    RETURN_NAMES = ("scenes_video","scenes_count","frame_count","fps",)
+    FUNCTION = "load_video"
+    OUTPUT_NODE = True
+    OUTPUT_IS_LIST = (True,False,False,False,)
+
+
+    def is_gif(self, filename):
+        file_parts = filename.split('.')
+        return len(file_parts) > 1 and file_parts[-1] == "gif"
+
+    def load_video_cv_fallback(self, video, frame_load_cap, skip_first_frames):
+        try:
+            video_cap = cv2.VideoCapture(folder_paths.get_annotated_filepath(video))
+            if not video_cap.isOpened():
+                raise ValueError(f"{video} could not be loaded with cv fallback.")
+            # set video_cap to look at start_index frame
+            images = []
+            total_frame_count = 0
+            frames_added = 0
+            base_frame_time = 1/video_cap.get(cv2.CAP_PROP_FPS)
+            
+            target_frame_time = base_frame_time
+
+            time_offset=0.0
+            while video_cap.isOpened():
+                if time_offset < target_frame_time:
+                    is_returned, frame = video_cap.read()
+                    # if didn't return frame, video has ended
+                    if not is_returned:
+                        break
+                    time_offset += base_frame_time
+                if time_offset < target_frame_time:
+                    continue
+                time_offset -= target_frame_time
+                # if not at start_index, skip doing anything with frame
+                total_frame_count += 1
+                if total_frame_count <= skip_first_frames:
+                    continue
+                # TODO: do whatever operations need to happen, like force_size, etc
+
+                # opencv loads images in BGR format (yuck), so need to convert to RGB for ComfyUI use
+                # follow up: can videos ever have an alpha channel?
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                # convert frame to comfyui's expected format (taken from comfy's load image code)
+                image = Image.fromarray(frame)
+                image = ImageOps.exif_transpose(image)
+                image = np.array(image, dtype=np.float32) / 255.0
+                image = torch.from_numpy(image)[None,]
+                images.append(image)
+                frames_added += 1
+                # if cap exists and we've reached it, stop processing frames
+                if frame_load_cap > 0 and frames_added >= frame_load_cap:
+                    break
+        finally:
+            video_cap.release()
+        images = torch.cat(images, dim=0)
+       
+        return (images, frames_added)
+
+    def load_video(self, video,video_segment_frames,transition_frames ):
+        
+        video_path = folder_paths.get_annotated_filepath(video)
+        
+        # check if video is a gif - will need to use cv fallback to read frames
+        # use cv fallback if ffmpeg not installed or gif
+        # if ffmpeg_path is None:
+        #     return self.load_video_cv_fallback(video, frame_load_cap, skip_first_frames)
+        # otherwise, continue with ffmpeg
+        
+        # args_dummy = [ffmpeg_path, "-i", video_path, "-f", "null", "-"]
+        # try:
+        #     with subprocess.Popen(args_dummy, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) as proc:
+        #         for line in proc.stderr.readlines():
+        #             match = re.search(", ([1-9]|\\d{2,})x(\\d+)",line.decode('utf-8'))
+        #             if match is not None:
+        #                 size = [int(match.group(1)), int(match.group(2))]
+        #                 break
+        # except Exception as e:
+        #     print(f"Retrying with opencv due to ffmpeg error: {e}")
+        #     return self.load_video_cv_fallback(video, frame_load_cap, skip_first_frames)
+        # args_all_frames = [ffmpeg_path, "-i", video_path, "-v", "error",
+        #                      "-pix_fmt", "rgb24"]
+
+        # vfilters = []
+       
+        # if skip_first_frames > 0:
+        #     vfilters.append(f"select=gt(n\\,{skip_first_frames-1})")
+        # if frame_load_cap > 0:
+        #     vfilters.append(f"select=gt({frame_load_cap}\\,n)")
+        # #manually calculate aspect ratio to ensure reads remain aligned
+        
+        # if len(vfilters) > 0:
+        #     args_all_frames += ["-vf", ",".join(vfilters)]
+
+        # args_all_frames += ["-f", "rawvideo", "-"]
+        # images = []
+        # try:
+        #     with subprocess.Popen(args_all_frames, stdout=subprocess.PIPE) as proc:
+        #         #Manually buffer enough bytes for an image
+        #         bpi = size[0]*size[1]*3
+        #         current_bytes = bytearray(bpi)
+        #         current_offset=0
+        #         while True:
+        #             bytes_read = proc.stdout.read(bpi - current_offset)
+        #             if bytes_read is None:#sleep to wait for more data
+        #                 time.sleep(.2)
+        #                 continue
+        #             if len(bytes_read) == 0:#EOF
+        #                 break
+        #             current_bytes[current_offset:len(bytes_read)] = bytes_read
+        #             current_offset+=len(bytes_read)
+        #             if current_offset == bpi:
+        #                 images.append(np.array(current_bytes, dtype=np.float32).reshape(size[1], size[0], 3) / 255.0)
+        #                 current_offset = 0
+        # except Exception as e:
+        #     print(f"Retrying with opencv due to ffmpeg error: {e}")
+        #     return self.load_video_cv_fallback(video, frame_load_cap, skip_first_frames)
+
+        # imgs=split_list(images,video_segment_frames,transition_frames)
+
+        # temp path 
+        tp=folder_paths.get_temp_directory()
+        basename = os.path.basename(video_path)  # 获取文件名
+        name_without_extension = os.path.splitext(basename)[0]  # 去掉文件后缀
+
+        folder_path = create_folder(tp,name_without_extension)
+        
+ 
+        # 导出的数据
+        scenes_video,total_frames,fps=split_video(video_path,video_segment_frames,
+                                 transition_frames,folder_path)
+      
+
+        # imgs=[torch.from_numpy(np.stack(im)) for im in imgs]
+
+        # images = torch.from_numpy(np.stack(images))
+
+        return (scenes_video,len(scenes_video), total_frames,fps,)
+
+    @classmethod
+    def IS_CHANGED(s, video, **kwargs):
+        image_path = folder_paths.get_annotated_filepath(video)
+        m = hashlib.sha256()
+        with open(image_path, 'rb') as f:
+            m.update(f.read())
+        return m.digest().hex()
+
+    @classmethod
+    def VALIDATE_INPUTS(s, video, **kwargs):
+        if not folder_paths.exists_annotated_filepath(video):
+            return "Invalid image file: {}".format(video)
+
+        return True
+    
+
+# The code is based on ComfyUI-VideoHelperSuite modification.
+class VideoCombine_Adv:
+    @classmethod
+    def INPUT_TYPES(s):
+        #Hide ffmpeg formats if ffmpeg isn't available
+        if ffmpeg_path is not None:
+            ffmpeg_formats = ["video/"+x[:-5] for x in folder_paths.get_filename_list("video_formats")]
+        else:
+            ffmpeg_formats = []
+        # ffmpeg_formats =["video/"+x for x in  ['webm', 'mp4', 'mkv']]
+        return {
+            "required": {
+                "image_batch": ("IMAGE",),
+                "frame_rate": (
+                    "INT",
+                    {"default": 8, "min": 1, "step": 1},
+                ),
+                "loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
+                "filename_prefix": ("STRING", {"default": "Comfyui"}),
+                "format": (["image/gif", "image/webp"] + ffmpeg_formats,),
+                "pingpong": ("BOOLEAN", {"default": False}),
+                "save_image": ("BOOLEAN", {"default": True}),
+                "metadata": ("BOOLEAN", {"default": False}),
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+                "extra_pnginfo": "EXTRA_PNGINFO",
+            },
+        }
+
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+    CATEGORY = "♾️Mixlab/Video"
+    FUNCTION = "run"
+
+    def save_with_tempfile(self, args, metadata, file_path, frames, env):
+        #Ensure temp directory exists
+        os.makedirs(folder_paths.get_temp_directory(), exist_ok=True)
+
+        metadata_path = os.path.join(folder_paths.get_temp_directory(), "metadata.txt")
+        #metadata from file should  escape = ; # \ and newline
+        #From my testing, though, only backslashes need escapes and = in particular causes problems
+        #It is likely better to prioritize future compatibility with containers that don't support
+        #or shouldn't use the comment tag for embedding metadata
+        metadata = metadata.replace("\\","\\\\")
+        metadata = metadata.replace(";","\\;")
+        metadata = metadata.replace("#","\\#")
+        #metadata = metadata.replace("=","\\=")
+        metadata = metadata.replace("\n","\\\n")
+        with open(metadata_path, "w") as f:
+            f.write(";FFMETADATA1\n")
+            f.write(metadata)
+        args = args[:1] + ["-i", metadata_path] + args[1:] + [file_path]
+        with subprocess.Popen(args, stdin=subprocess.PIPE, env=env) as proc:
+            for frame in frames:
+                proc.stdin.write(frame.tobytes())
+
+    def run(
+        self,
+        image_batch,
+        frame_rate: int,
+        loop_count: int,
+        filename_prefix="AnimateDiff",
+        format="image/gif",
+        pingpong=False,
+        save_image=True,
+        metadata=False,
+        prompt=None,
+        extra_pnginfo=None,
+    ):
+        images=image_batch
+
+        frames: List[Image.Image] = []
+        for image in images:
+            img = 255.0 * image.cpu().numpy()
+            img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
+            # resize 保证
+            # 检查图像的高度是否是2的倍数，如果不是，则调整高度
+            if img.height % 2 != 0:
+                img = img.resize((img.width, img.height + 1))
+
+            # 检查图像的宽度是否是2的倍数，如果不是，则调整宽度
+            if img.width % 2 != 0:
+                img = img.resize((img.width + 1, img.height))
+                
+            frames.append(img)
+
+        # get output information
+        output_dir = (
+            folder_paths.get_output_directory()
+            if save_image
+            else folder_paths.get_temp_directory()
+        )
+        (
+            full_output_folder,
+            filename,
+            counter,
+            subfolder,
+            _,
+        ) = folder_paths.get_save_image_path(filename_prefix, output_dir)
+
+        metadata = PngInfo()
+        video_metadata = {}
+        if prompt is not None:
+            metadata.add_text("prompt", json.dumps(prompt))
+            video_metadata["prompt"] = prompt
+        if extra_pnginfo is not None:
+            for x in extra_pnginfo:
+                metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+                video_metadata[x] = extra_pnginfo[x]
+
+        # 取消保存metadata
+        if metadata==False:
+            metadata = PngInfo()
+
+        # save first frame as png to keep metadata
+        file = f"{filename}_{counter:05}_.png"
+        file_path = os.path.join(full_output_folder, file)
+        frames[0].save(
+            file_path,
+            pnginfo=metadata,
+            compress_level=4,
+        )
+        if pingpong:
+            frames = frames + frames[-2:0:-1]
+
+        format_type, format_ext = format.split("/")
+        file = f"{filename}_{counter:05}_.{format_ext}"
+        file_path = os.path.join(full_output_folder, file)
+        if format_type == "image":
+            # Use pillow directly to save an animated image
+            frames[0].save(
+                file_path,
+                format=format_ext.upper(),
+                save_all=True,
+                append_images=frames[1:],
+                duration=round(1000 / frame_rate),
+                loop=loop_count,
+                compress_level=4,
+            )
+        else:
+            # Use ffmpeg to save a video
+            if ffmpeg_path is None:
+                #Should never be reachable
+                raise ProcessLookupError("Could not find ffmpeg")
+
+            video_format_path = folder_paths.get_full_path("video_formats", format_ext + ".json")
+            with open(video_format_path, 'r') as stream:
+                video_format = json.load(stream)
+            file = f"{filename}_{counter:05}_.{video_format['extension']}"
+            file_path = os.path.join(full_output_folder, file)
+            dimensions = f"{frames[0].width}x{frames[0].height}"
+            metadata_args = ["-metadata", "comment=" + json.dumps(video_metadata)]
+            args = [ffmpeg_path, "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
+                    "-s", dimensions, "-r", str(frame_rate), "-i", "-"] \
+                    + video_format['main_pass']
+            # On linux, max arg length is Pagesize * 32 -> 131072
+            # On windows, this around 32767 but seems to vary wildly by > 500
+            # in a manor not solely related to other arguments
+            if os.name == 'posix':
+                max_arg_length = 4096*32
+            else:
+                max_arg_length = 32767 - len(" ".join(args + [metadata_args[0]] + [file_path])) - 1
+            #test max limit
+            #metadata_args[1] = metadata_args[1] + "a"*(max_arg_length - len(metadata_args[1])-1)
+
+            env=os.environ.copy()
+            if  "environment" in video_format:
+                env.update(video_format["environment"])
+            if len(metadata_args[1]) >= max_arg_length:
+                print(f"Using fallback file for extremely long metadata: {len(metadata_args[1])}/{max_arg_length}")
+                self.save_with_tempfile(args, metadata_args[1], file_path, frames, env)
+            else:
+                try:
+                    with subprocess.Popen(args + metadata_args + [file_path],
+                                          stdin=subprocess.PIPE, env=env) as proc:
+                        for frame in frames:
+                            proc.stdin.write(frame.tobytes())
+                except FileNotFoundError as e:
+                    if "winerror" in dir(e) and e.winerror == 206:
+                        print("Metadata was too long. Retrying with fallback file")
+                        self.save_with_tempfile(args, metadata_args[1], file_path, frames, env)
+                    else:
+                        raise
+                except OSError as e:
+                    if "errno" in dir(e) and e.errno == 7:
+                        print("Metadata was too long. Retrying with fallback file")
+                        self.save_with_tempfile(args, metadata_args[1], file_path, frames, env)
+                    else:
+                        raise
+
+        previews = [
+            {
+                "filename": file,
+                "subfolder": subfolder,
+                "type": "output" if save_image else "temp",
+                "format": format,
+            }
+        ]
+        return {"ui": {"gifs": previews}}
+
+
+class VAEEncodeForInpaint_Frames:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { 
+            "vae": ("VAE", ), 
+            "images": ("IMAGE", ), 
+            "masks": ("MASK", ), 
+            "grow_mask_by": ("INT", {"default": 6, "min": 0, "max": 64, "step": 1}),
+            }}
+    
+    FUNCTION = "encode"
+
+
+    RETURN_TYPES = ("LATENT",)
+    RETURN_NAMES = ("LATENT",)
+ 
+    CATEGORY = "♾️Mixlab/Video"
+
+    OUTPUT_NODE = True
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True,)
+
+
+    def encode(self, vae, images, masks, grow_mask_by=[6]):
+        vae=vae[0]
+        grow_mask_by=grow_mask_by[0]
+
+        result=[]
+
+        for i in range(len(images)):
+            pixels=images[i]
+            mask=masks[i]
+
+
+            x = (pixels.shape[1] // 8) * 8
+            y = (pixels.shape[2] // 8) * 8
+            mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
+
+            pixels = pixels.clone()
+            if pixels.shape[1] != x or pixels.shape[2] != y:
+                x_offset = (pixels.shape[1] % 8) // 2
+                y_offset = (pixels.shape[2] % 8) // 2
+                pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
+                mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]
+
+            #grow mask by a few pixels to keep things seamless in latent space
+            if grow_mask_by == 0:
+                mask_erosion = mask
+            else:
+                kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
+                padding = math.ceil((grow_mask_by - 1) / 2)
+
+                mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)
+
+            m = (1.0 - mask.round()).squeeze(1)
+            for i in range(3):
+                pixels[:,:,:,i] -= 0.5
+                pixels[:,:,:,i] *= m
+                pixels[:,:,:,i] += 0.5
+            t = vae.encode(pixels)
+
+            result.append({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())})
+
+
+        return (result, )

nodes/VisualStylePrompting/attention_functions.py

@@ -0,0 +1,45 @@
+from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked
+from .style_functions import adain, concat_first
+
+class VisualStyleProcessor(object):
+    def __init__(self, 
+        module_self, 
+        keys_scale: float = 1.0,
+        enabled: bool = True, 
+        adain_queries: bool = True,
+        adain_keys: bool = True,
+        adain_values: bool = False 
+    ):
+        self.module_self = module_self
+        self.keys_scale = keys_scale
+        self.enabled = enabled
+        self.adain_queries = adain_queries
+        self.adain_keys = adain_keys
+        self.adain_values = adain_values
+
+    def visual_style_forward(self, x, context, value, mask=None):
+        q = self.module_self.to_q(x)
+        context = default(context, x)
+        k = self.module_self.to_k(context)
+        if value is not None:
+            v = self.module_self.to_v(value)
+            del value
+        else:
+            v = self.module_self.to_v(context)
+
+        if self.enabled:
+            if self.adain_queries:
+                q = adain(q)
+            if self.adain_keys:
+                k = adain(k)
+            if self.adain_values:
+                v = adain(v)
+            
+            k = concat_first(k, -2, self.keys_scale)
+            v = concat_first(v, -2)
+
+        if mask is None:
+            out = optimized_attention(q, k, v, self.module_self.heads)
+        else:
+            out = optimized_attention_masked(q, k, v, self.module_self.heads, mask)
+        return self.module_self.to_out(out)

nodes/VisualStylePrompting/style_functions.py

@@ -0,0 +1,60 @@
+import torch
+
+from einops import rearrange
+from dataclasses import dataclass
+
+T = torch.Tensor
+
+@dataclass(frozen=True)
+class StyleAlignedArgs:
+    share_group_norm: bool = True
+    share_layer_norm: bool = True,
+    share_attention: bool = True
+    adain_queries: bool = True
+    adain_keys: bool = True
+    adain_values: bool = False
+    full_attention_share: bool = False
+    keys_scale: float = 1.
+    only_self_level: float = 0.
+
+def expand_first(feat: T, scale=1., ) -> T:
+    b = feat.shape[0]
+    feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1)
+    if scale == 1:
+        feat_style = feat_style.expand(2, b // 2, *feat.shape[1:])
+    else:
+        feat_style = feat_style.repeat(1, b // 2, 1, 1, 1)
+        feat_style = torch.cat([feat_style[:, :1], scale * feat_style[:, 1:]], dim=1)
+    return feat_style.reshape(*feat.shape)
+
+
+def concat_first(feat: T, dim=2, scale=1.) -> T:
+    feat_style = expand_first(feat, scale=scale)
+    return torch.cat((feat, feat_style), dim=dim)
+
+
+def calc_mean_std(feat, eps: float = 1e-5) -> tuple[T, T]:
+    feat_std = (feat.var(dim=-2, keepdims=True) + eps).sqrt()
+    feat_mean = feat.mean(dim=-2, keepdims=True)
+    return feat_mean, feat_std
+
+
+def adain(feat: T) -> T:
+    feat_mean, feat_std = calc_mean_std(feat)
+    feat_style_mean = expand_first(feat_mean)
+    feat_style_std = expand_first(feat_std)
+    feat = (feat - feat_mean) / feat_std
+    feat = feat * feat_style_std + feat_style_mean
+    return feat
+
+def swapping_attention(key, value, chunk_size=2):
+    chunk_length = key.size()[0] // chunk_size  # [text-condition, null-condition]
+    reference_image_index = [0] * chunk_length  # [0 0 0 0 0]
+    key = rearrange(key, "(b f) d c -> b f d c", f=chunk_length)
+    key = key[:, reference_image_index]  # ref to all
+    key = rearrange(key, "b f d c -> (b f) d c")
+    value = rearrange(value, "(b f) d c -> b f d c", f=chunk_length)
+    value = value[:, reference_image_index]  # ref to all
+    value = rearrange(value, "b f d c -> (b f) d c")
+
+    return key, value

nodes/VoiceStreamAI/asr/asr_factory.py

@@ -0,0 +1,12 @@
+from VoiceStreamAI.asr.whisper_asr import WhisperASR
+from VoiceStreamAI.asr.faster_whisper_asr import FasterWhisperASR
+
+class ASRFactory:
+    @staticmethod
+    def create_asr_pipeline(type, **kwargs):
+        if type == "whisper":
+            return WhisperASR(**kwargs)
+        if type == "faster_whisper":
+            return FasterWhisperASR(**kwargs)
+        else:
+            raise ValueError(f"Unknown ASR pipeline type: {type}")

nodes/VoiceStreamAI/asr/asr_interface.py

@@ -0,0 +1,9 @@
+class ASRInterface:
+    async def transcribe(self, client):
+        """
+        Transcribe the given audio data.
+
+        :param client: The client object with all the member variables including the buffer
+        :return: The transcription structure, see for example the faster_whisper_asr.py file.
+        """
+        raise NotImplementedError("This method should be implemented by subclasses.")

nodes/VoiceStreamAI/asr/faster_whisper_asr.py

@@ -0,0 +1,142 @@
+import os
+from faster_whisper import WhisperModel
+
+from VoiceStreamAI.asr.asr_interface import ASRInterface
+from VoiceStreamAI.audio_utils import save_audio_to_file
+
+import folder_paths
+
+language_codes = {
+    "afrikaans": "af",
+    "amharic": "am",
+    "arabic": "ar",
+    "assamese": "as",
+    "azerbaijani": "az",
+    "bashkir": "ba",
+    "belarusian": "be",
+    "bulgarian": "bg",
+    "bengali": "bn",
+    "tibetan": "bo",
+    "breton": "br",
+    "bosnian": "bs",
+    "catalan": "ca",
+    "czech": "cs",
+    "welsh": "cy",
+    "danish": "da",
+    "german": "de",
+    "greek": "el",
+    "english": "en",
+    "spanish": "es",
+    "estonian": "et",
+    "basque": "eu",
+    "persian": "fa",
+    "finnish": "fi",
+    "faroese": "fo",
+    "french": "fr",
+    "galician": "gl",
+    "gujarati": "gu",
+    "hausa": "ha",
+    "hawaiian": "haw",
+    "hebrew": "he",
+    "hindi": "hi",
+    "croatian": "hr",
+    "haitian": "ht",
+    "hungarian": "hu",
+    "armenian": "hy",
+    "indonesian": "id",
+    "icelandic": "is",
+    "italian": "it",
+    "japanese": "ja",
+    "javanese": "jw",
+    "georgian": "ka",
+    "kazakh": "kk",
+    "khmer": "km",
+    "kannada": "kn",
+    "korean": "ko",
+    "latin": "la",
+    "luxembourgish": "lb",
+    "lingala": "ln",
+    "lao": "lo",
+    "lithuanian": "lt",
+    "latvian": "lv",
+    "malagasy": "mg",
+    "maori": "mi",
+    "macedonian": "mk",
+    "malayalam": "ml",
+    "mongolian": "mn",
+    "marathi": "mr",
+    "malay": "ms",
+    "maltese": "mt",
+    "burmese": "my",
+    "nepali": "ne",
+    "dutch": "nl",
+    "norwegian nynorsk": "nn",
+    "norwegian": "no",
+    "occitan": "oc",
+    "punjabi": "pa",
+    "polish": "pl",
+    "pashto": "ps",
+    "portuguese": "pt",
+    "romanian": "ro",
+    "russian": "ru",
+    "sanskrit": "sa",
+    "sindhi": "sd",
+    "sinhalese": "si",
+    "slovak": "sk",
+    "slovenian": "sl",
+    "shona": "sn",
+    "somali": "so",
+    "albanian": "sq",
+    "serbian": "sr",
+    "sundanese": "su",
+    "swedish": "sv",
+    "swahili": "sw",
+    "tamil": "ta",
+    "telugu": "te",
+    "tajik": "tg",
+    "thai": "th",
+    "turkmen": "tk",
+    "tagalog": "tl",
+    "turkish": "tr",
+    "tatar": "tt",
+    "ukrainian": "uk",
+    "urdu": "ur",
+    "uzbek": "uz",
+    "vietnamese": "vi",
+    "yiddish": "yi",
+    "yoruba": "yo",
+    "chinese": "zh",
+    "cantonese": "yue",
+}
+
+
+class FasterWhisperASR(ASRInterface):
+    def __init__(self, **kwargs):
+        model_size = kwargs.get('model_size', "large-v3")
+        device = kwargs.get('device', "cuda")
+        model_root = os.path.join(folder_paths.models_dir, "whisper") 
+        # Run on GPU with FP16
+        self.asr_pipeline = WhisperModel(model_size, device=device, compute_type="float16",download_root=model_root)
+
+    async def transcribe(self, client):
+        file_path = await save_audio_to_file(client.scratch_buffer, client.get_file_name())
+
+        language = None if client.config['language'] is None else language_codes.get(client.config['language'].lower())
+        segments, info = self.asr_pipeline.transcribe(file_path, word_timestamps=True, language=language)
+
+        segments = list(segments)  # The transcription will actually run here.
+        os.remove(file_path)
+
+        flattened_words = [word for segment in segments for word in segment.words]
+
+        to_return = {
+        "language": info.language,
+        "language_probability": info.language_probability,
+        "text": ' '.join([s.text.strip() for s in segments]),
+        "words":
+            [
+                {"word": w.word, "start": w.start, "end": w.end, "probability":w.probability} for w in flattened_words
+            ]
+        }
+        return to_return
+

nodes/VoiceStreamAI/asr/whisper_asr.py

@@ -0,0 +1,27 @@
+from transformers import pipeline
+from VoiceStreamAI.asr.asr_interface import ASRInterface
+from VoiceStreamAI.audio_utils import save_audio_to_file
+import os
+
+class WhisperASR(ASRInterface):
+    def __init__(self, **kwargs):
+        model_name = kwargs.get('model_name', "openai/whisper-large-v3")
+        self.asr_pipeline = pipeline("automatic-speech-recognition", model=model_name)
+
+    async def transcribe(self, client):
+        file_path = await save_audio_to_file(client.scratch_buffer, client.get_file_name())
+        
+        if client.config['language'] is not None:
+            to_return = self.asr_pipeline(file_path, generate_kwargs={"language": client.config['language']})['text']
+        else:
+            to_return = self.asr_pipeline(file_path)['text']
+
+        os.remove(file_path)
+
+        to_return = {
+            "language": "UNSUPPORTED_BY_HUGGINGFACE_WHISPER",
+            "language_probability": None,
+            "text": to_return.strip(),
+            "words": "UNSUPPORTED_BY_HUGGINGFACE_WHISPER"
+        }
+        return to_return

nodes/VoiceStreamAI/audio_utils.py

@@ -0,0 +1,26 @@
+import wave
+import os
+
+async def save_audio_to_file(audio_data, file_name, audio_dir="audio_files", audio_format="wav"):
+    """
+    Saves the audio data to a file.
+
+    :param client_id: Unique identifier for the client.
+    :param audio_data: The audio data to save.
+    :param file_counters: Dictionary to keep track of file counts for each client.
+    :param audio_dir: Directory where audio files will be saved.
+    :param audio_format: Format of the audio file.
+    :return: Path to the saved audio file.
+    """
+
+    os.makedirs(audio_dir, exist_ok=True)
+    
+    file_path = os.path.join(audio_dir, file_name)
+
+    with wave.open(file_path, 'wb') as wav_file:
+        wav_file.setnchannels(1)  # Assuming mono audio
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(16000)
+        wav_file.writeframes(audio_data)
+
+    return file_path

nodes/VoiceStreamAI/buffering_strategy/buffering_strategies.py

@@ -0,0 +1,142 @@
+import os
+import asyncio
+import json
+import time
+
+from VoiceStreamAI.buffering_strategy.buffering_strategy_interface import BufferingStrategyInterface
+from openai import OpenAI
+
+class SilenceAtEndOfChunk(BufferingStrategyInterface):
+    """
+    A buffering strategy that processes audio at the end of each chunk with silence detection.
+
+    This class is responsible for handling audio chunks, detecting silence at the end of each chunk,
+    and initiating the transcription process for the chunk.
+
+    Attributes:
+        client (Client): The client instance associated with this buffering strategy.
+        chunk_length_seconds (float): Length of each audio chunk in seconds.
+        chunk_offset_seconds (float): Offset time in seconds to be considered for processing audio chunks.
+    """
+
+    def __init__(self, client, **kwargs):
+        """
+        Initialize the SilenceAtEndOfChunk buffering strategy.
+
+        Args:
+            client (Client): The client instance associated with this buffering strategy.
+            **kwargs: Additional keyword arguments, including 'chunk_length_seconds' and 'chunk_offset_seconds'.
+        """
+        self.client = client
+
+        self.chunk_length_seconds = os.environ.get('BUFFERING_CHUNK_LENGTH_SECONDS')
+        if not self.chunk_length_seconds:
+            self.chunk_length_seconds = kwargs.get('chunk_length_seconds')
+        self.chunk_length_seconds = float(self.chunk_length_seconds)
+
+        self.chunk_offset_seconds = os.environ.get('BUFFERING_CHUNK_OFFSET_SECONDS')
+        if not self.chunk_offset_seconds:
+            self.chunk_offset_seconds = kwargs.get('chunk_offset_seconds')
+        self.chunk_offset_seconds = float(self.chunk_offset_seconds)
+
+        self.error_if_not_realtime = os.environ.get('ERROR_IF_NOT_REALTIME')
+        if not self.error_if_not_realtime:
+            self.error_if_not_realtime = kwargs.get('error_if_not_realtime', False)
+        
+        self.processing_flag = False
+
+        self.messages=[]
+
+    def process_audio(self, websocket, vad_pipeline, asr_pipeline,llm_port):
+        """
+        Process audio chunks by checking their length and scheduling asynchronous processing.
+
+        This method checks if the length of the audio buffer exceeds the chunk length and, if so,
+        it schedules asynchronous processing of the audio.
+
+        Args:
+            websocket (Websocket): The WebSocket connection for sending transcriptions.
+            vad_pipeline: The voice activity detection pipeline.
+            asr_pipeline: The automatic speech recognition pipeline.
+        """
+        chunk_length_in_bytes = self.chunk_length_seconds * self.client.sampling_rate * self.client.samples_width
+        if len(self.client.buffer) > chunk_length_in_bytes:
+            if self.processing_flag:
+                exit("Error in realtime processing: tried processing a new chunk while the previous one was still being processed")
+
+            self.client.scratch_buffer += self.client.buffer
+            self.client.buffer.clear()
+            self.processing_flag = True
+            # Schedule the processing in a separate task
+            asyncio.create_task(self.process_audio_async(websocket, vad_pipeline, asr_pipeline,llm_port))
+    
+    async def process_audio_async(self, websocket, vad_pipeline, asr_pipeline,llm_port):
+        """
+        Asynchronously process audio for activity detection and transcription.
+
+        This method performs heavy processing, including voice activity detection and transcription of
+        the audio data. It sends the transcription results through the WebSocket connection.
+
+        Args:
+            websocket (Websocket): The WebSocket connection for sending transcriptions.
+            vad_pipeline: The voice activity detection pipeline.
+            asr_pipeline: The automatic speech recognition pipeline.
+        """   
+        start = time.time()
+        vad_results = await vad_pipeline.detect_activity(self.client)
+
+        if len(vad_results) == 0:
+            self.client.scratch_buffer.clear()
+            self.client.buffer.clear()
+            self.processing_flag = False
+            return
+
+        last_segment_should_end_before = ((len(self.client.scratch_buffer) / (self.client.sampling_rate * self.client.samples_width)) - self.chunk_offset_seconds)
+        if vad_results[-1]['end'] < last_segment_should_end_before:
+            transcription = await asr_pipeline.transcribe(self.client)
+            if transcription['text'] != '':
+                end = time.time()
+                transcription['processing_time'] = end - start
+
+                transcription['status']="chat_start"
+
+                json_transcription = json.dumps(transcription)
+
+                await websocket.send(json_transcription)
+
+                # Point to the local server
+                client = OpenAI(base_url=f"http://localhost:{llm_port}/v1", api_key="lm-studio")
+
+               
+                messages=[
+                    {"role": "system", "content": "You are a friendly and engaging AI designed to interact with users in a conversational manner. Your personality is that of a sophisticated and polite young professional who is both a designer and a programmer. You are well-mannered, articulate, and possess a good sense of humor. Your goal is to provide helpful and insightful responses while maintaining a pleasant and enjoyable conversation. Be sure to use your knowledge in design and programming to enrich the dialogue and offer relevant advice or information when appropriate. Always be respectful and considerate of the user's feelings and perspectives. Additionally, you are fluent in both English and Chinese, and can seamlessly switch between the two languages to best assist users."}, 
+                ]+self.messages[-10:0]+[{"role": "user", "content":transcription['text']}]
+                # print('#messages',messages)
+                
+                completion = client.chat.completions.create(
+                    model="model-identifier",
+                    messages=messages,
+                    temperature=0.7,
+                    )
+
+                transcription['asistant'] = completion.choices[0].message.content
+
+                transcription['status']="chat_end"
+
+                json_transcription = json.dumps(transcription) 
+
+                self.messages.append({
+                    "role": "user", 
+                    "content":transcription['text']})
+                
+                self.messages.append({
+                    "role": "asistant", 
+                    "content": transcription['asistant']
+                                      })
+                # print('#messages',completion.choices[0].message.content)
+
+                await websocket.send(json_transcription)
+            self.client.scratch_buffer.clear()
+            self.client.increment_file_counter()
+        
+        self.processing_flag = False

nodes/VoiceStreamAI/buffering_strategy/buffering_strategy_factory.py

@@ -0,0 +1,41 @@
+from VoiceStreamAI.buffering_strategy.buffering_strategies import SilenceAtEndOfChunk
+
+class BufferingStrategyFactory:
+    """
+    A factory class for creating instances of different buffering strategies.
+
+    This factory provides a centralized way to instantiate various buffering strategies
+    based on the type specified. It abstracts the creation logic, making it easier to
+    manage and extend with new buffering strategy types.
+
+    Methods:
+        create_buffering_strategy: Creates and returns an instance of a specified buffering strategy.
+    """
+
+    @staticmethod
+    def create_buffering_strategy(type, client, **kwargs):
+        """
+        Creates an instance of a buffering strategy based on the specified type.
+
+        This method acts as a factory for creating buffering strategy objects. It returns
+        an instance of the strategy corresponding to the given type. If the type is not
+        recognized, it raises a ValueError.
+
+        Args:
+            type (str): The type of buffering strategy to create. Currently supports 'silence_at_end_of_chunk'.
+            client (Client): The client instance to be associated with the buffering strategy.
+            **kwargs: Additional keyword arguments specific to the buffering strategy being created.
+
+        Returns:
+            An instance of the specified buffering strategy.
+
+        Raises:
+            ValueError: If the specified type is not recognized or supported.
+
+        Example:
+            strategy = BufferingStrategyFactory.create_buffering_strategy("silence_at_end_of_chunk", client)
+        """
+        if type == "silence_at_end_of_chunk":
+            return SilenceAtEndOfChunk(client, **kwargs)
+        else:
+            raise ValueError(f"Unknown buffering strategy type: {type}")

nodes/VoiceStreamAI/buffering_strategy/buffering_strategy_interface.py

@@ -0,0 +1,31 @@
+class BufferingStrategyInterface:
+    """
+    An interface class for buffering strategies in audio processing systems.
+
+    This class defines the structure for buffering strategies used in handling
+    and processing audio data. It serves as a template for creating custom buffering
+    strategies that fit specific requirements of an audio processing pipeline.
+
+    Subclasses should implement the methods defined in this interface to ensure
+    consistency and compatibility with the system's audio processing framework.
+
+    Methods:
+        process_audio: Process audio data. This method should be implemented by subclasses.
+    """
+
+    def process_audio(self, websocket, vad_pipeline, asr_pipeline):
+        """
+        Process audio data using the given WebSocket connection, VAD pipeline, and ASR pipeline.
+
+        This method is intended to be overridden in subclasses to provide specific logic
+        for handling and processing audio data in different buffering strategies.
+
+        Args:
+            websocket (Websocket): The WebSocket connection for communication with clients.
+            vad_pipeline: The Voice Activity Detection (VAD) pipeline used for detecting speech in the audio.
+            asr_pipeline: The Automatic Speech Recognition (ASR) pipeline used for transcribing speech in the audio.
+
+        Raises:
+            NotImplementedError: If the method is not implemented in the subclass.
+        """
+        raise NotImplementedError("This method should be implemented by subclasses.")

nodes/VoiceStreamAI/client.py

@@ -0,0 +1,54 @@
+from VoiceStreamAI.buffering_strategy.buffering_strategy_factory import BufferingStrategyFactory
+
+class Client:
+    """
+    Represents a client connected to the VoiceStreamAI server.
+
+    This class maintains the state for each connected client, including their
+    unique identifier, audio buffer, configuration, and a counter for processed audio files.
+
+    Attributes:
+        client_id (str): A unique identifier for the client.
+        buffer (bytearray): A buffer to store incoming audio data.
+        config (dict): Configuration settings for the client, like chunk length and offset.
+        file_counter (int): Counter for the number of audio files processed.
+        total_samples (int): Total number of audio samples received from this client.
+        sampling_rate (int): The sampling rate of the audio data in Hz.
+        samples_width (int): The width of each audio sample in bits.
+    """
+    def __init__(self, client_id, sampling_rate, samples_width):
+        self.client_id = client_id
+        self.buffer = bytearray()
+        self.scratch_buffer = bytearray()
+        self.config = {"language": None,
+                       "processing_strategy": "silence_at_end_of_chunk", 
+                       "processing_args": {
+                           "chunk_length_seconds": 5, 
+                           "chunk_offset_seconds": 0.1
+                           }
+                       }
+        self.file_counter = 0
+        self.total_samples = 0
+        self.sampling_rate = sampling_rate
+        self.samples_width = samples_width
+        self.buffering_strategy = BufferingStrategyFactory.create_buffering_strategy(self.config['processing_strategy'], self, **self.config['processing_args'])
+
+    def update_config(self, config_data):
+        self.config.update(config_data)
+        self.buffering_strategy = BufferingStrategyFactory.create_buffering_strategy(self.config['processing_strategy'], self, **self.config['processing_args'])
+
+    def append_audio_data(self, audio_data):
+        self.buffer.extend(audio_data)
+        self.total_samples += len(audio_data) / self.samples_width
+
+    def clear_buffer(self):
+        self.buffer.clear()
+
+    def increment_file_counter(self):
+        self.file_counter += 1
+
+    def get_file_name(self):
+        return f"{self.client_id}_{self.file_counter}.wav"
+    
+    def process_audio(self, websocket, vad_pipeline, asr_pipeline,llm_port):
+        self.buffering_strategy.process_audio(websocket, vad_pipeline, asr_pipeline,llm_port)

nodes/VoiceStreamAI/main.py

@@ -0,0 +1,54 @@
+import argparse
+import asyncio
+import json
+
+import os
+import sys
+from pathlib import Path
+
+
+# 获取当前文件的绝对路径
+current_file_path = os.path.abspath(__file__)
+
+# 获取当前文件的目录
+current_directory = os.path.dirname(current_file_path)
+sys.path.append(str(Path(current_directory).parent))
+# print("sys.path", current_directory)
+
+
+from VoiceStreamAI.server import Server
+from VoiceStreamAI.asr.asr_factory import ASRFactory
+from VoiceStreamAI.vad.vad_factory import VADFactory
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="VoiceStreamAI Server: Real-time audio transcription using self-hosted Whisper and WebSocket")
+    parser.add_argument("--vad-type", type=str, default="pyannote", help="Type of VAD pipeline to use (e.g., 'pyannote')")
+    parser.add_argument("--vad-args", type=str, default='{"auth_token": "huggingface_token"}', help="JSON string of additional arguments for VAD pipeline")
+    parser.add_argument("--asr-type", type=str, default="faster_whisper", help="Type of ASR pipeline to use (e.g., 'whisper')")
+    parser.add_argument("--asr-args", type=str, default='{"model_size": "large-v3"}', help="JSON string of additional arguments for ASR pipeline")
+    parser.add_argument("--host", type=str, default="127.0.0.1", help="Host for the WebSocket server")
+    parser.add_argument("--port", type=int, default=8765, help="Port for the WebSocket server")
+    parser.add_argument("--certfile", type=str, default=None, help="The path to the SSL certificate (cert file) if using secure websockets")
+    parser.add_argument("--keyfile", type=str, default=None, help="The path to the SSL key file if using secure websockets")
+    return parser.parse_args()
+
+def main():
+    args = parse_args()
+
+    try:
+        vad_args = json.loads(args.vad_args)
+        asr_args = json.loads(args.asr_args)
+    except json.JSONDecodeError as e:
+        print(f"Error parsing JSON arguments: {e}")
+        return
+
+    vad_pipeline = VADFactory.create_vad_pipeline(args.vad_type, **vad_args)
+    asr_pipeline = ASRFactory.create_asr_pipeline(args.asr_type, **asr_args)
+
+    server = Server(vad_pipeline, asr_pipeline, host=args.host, port=args.port, sampling_rate=16000, samples_width=2, certfile=args.certfile, keyfile=args.keyfile)
+
+    asyncio.get_event_loop().run_until_complete(server.start())
+    asyncio.get_event_loop().run_forever()
+
+if __name__ == "__main__":
+    main()

nodes/VoiceStreamAI/requirements.txt

@@ -0,0 +1,7 @@
+websockets
+speechbrain
+pyannote-audio
+asyncio
+sentence-transformers
+transformers
+faster-whisper

nodes/VoiceStreamAI/server.py

@@ -0,0 +1,88 @@
+import websockets
+import uuid
+import json
+import asyncio
+import ssl
+
+from VoiceStreamAI.audio_utils import save_audio_to_file
+from VoiceStreamAI.client import Client
+
+class Server:
+    """
+    Represents the WebSocket server for handling real-time audio transcription.
+
+    This class manages WebSocket connections, processes incoming audio data,
+    and interacts with VAD and ASR pipelines for voice activity detection and
+    speech recognition.
+
+    Attributes:
+        vad_pipeline: An instance of a voice activity detection pipeline.
+        asr_pipeline: An instance of an automatic speech recognition pipeline.
+        host (str): Host address of the server.
+        port (int): Port on which the server listens.
+        sampling_rate (int): The sampling rate of audio data in Hz.
+        samples_width (int): The width of each audio sample in bits.
+        connected_clients (dict): A dictionary mapping client IDs to Client objects.
+    """
+    def __init__(self, vad_pipeline, asr_pipeline, host='localhost', port=8765, sampling_rate=16000, samples_width=2, certfile = None, keyfile = None,llm_port=9000):
+        self.vad_pipeline = vad_pipeline
+        self.asr_pipeline = asr_pipeline
+        self.host = host
+        self.port = port
+        self.sampling_rate = sampling_rate
+        self.samples_width = samples_width
+        self.certfile = certfile
+        self.keyfile = keyfile
+        self.connected_clients = {}
+
+        self.llm_port=llm_port
+
+    async def handle_audio(self, client, websocket):
+        while True:
+            message = await websocket.recv()
+
+            if isinstance(message, bytes):
+                client.append_audio_data(message)
+            elif isinstance(message, str):
+                config = json.loads(message)
+                if config.get('type') == 'config':
+                    client.update_config(config['data'])
+                    continue
+            else:
+                print(f"Unexpected message type from {client.client_id}")
+
+            # this is synchronous, any async operation is in BufferingStrategy
+            client.process_audio(websocket, self.vad_pipeline, self.asr_pipeline,self.llm_port)
+
+
+    async def handle_websocket(self, websocket, path):
+        client_id = str(uuid.uuid4())
+        client = Client(client_id, self.sampling_rate, self.samples_width)
+        self.connected_clients[client_id] = client
+
+        print(f"Client {client_id} connected")
+
+        try:
+            await self.handle_audio(client, websocket)
+        except websockets.ConnectionClosed as e:
+            print(f"Connection with {client_id} closed: {e}")
+        finally:
+            del self.connected_clients[client_id]
+
+    def start(self):
+        if self.certfile:
+            # Create an SSL context to enforce encrypted connections
+            ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+            
+            # Load your server's certificate and private key
+            # Replace 'your_cert_path.pem' and 'your_key_path.pem' with the actual paths to your files
+            ssl_context.load_cert_chain(certfile=self.certfile, keyfile=self.keyfile)
+
+            print(f"WebSocket server ready to accept secure connections on {self.host}:{self.port}")
+            
+            # Pass the SSL context to the serve function along with the host and port
+            # Ensure the secure flag is set to True if using a secure WebSocket protocol (wss://)
+            return websockets.serve(self.handle_websocket, self.host, self.port, ssl=ssl_context)
+        else:
+            print(f"WebSocket server ready to accept secure connections on {self.host}:{self.port}")
+            return websockets.serve(self.handle_websocket, self.host, self.port)

nodes/VoiceStreamAI/vad/pyannote_vad.py

@@ -0,0 +1,50 @@
+from os import remove
+import os
+
+from pyannote.core import Segment
+from pyannote.audio import Model
+from pyannote.audio.pipelines import VoiceActivityDetection
+
+from VoiceStreamAI.vad.vad_interface import VADInterface
+from VoiceStreamAI.audio_utils import save_audio_to_file
+
+
+class PyannoteVAD(VADInterface):
+    """
+    Pyannote-based implementation of the VADInterface.
+    """
+
+    def __init__(self, **kwargs):
+        """
+        Initializes Pyannote's VAD pipeline.
+
+        Args:
+            model_name (str): The model name for Pyannote.
+            auth_token (str, optional): Authentication token for Hugging Face.
+        """
+        
+        model_name = kwargs.get('model_name', "pyannote/segmentation")
+
+        auth_token = os.environ.get('PYANNOTE_AUTH_TOKEN')
+        if not auth_token:
+            auth_token = kwargs.get('auth_token')
+        
+        if auth_token is None:
+            raise ValueError("Missing required env var in PYANNOTE_AUTH_TOKEN or argument in --vad-args: 'auth_token'")
+        
+        pyannote_args = kwargs.get('pyannote_args', {"onset": 0.5, "offset": 0.5, "min_duration_on": 0.3, "min_duration_off": 0.3})
+        self.model = Model.from_pretrained(model_name, use_auth_token=auth_token)
+        self.vad_pipeline = VoiceActivityDetection(segmentation=self.model)
+        self.vad_pipeline.instantiate(pyannote_args)
+
+    async def detect_activity(self, client):
+        audio_file_path = await save_audio_to_file(client.scratch_buffer, client.get_file_name())
+        vad_results = self.vad_pipeline(audio_file_path)
+        remove(audio_file_path)
+        vad_segments = []
+        if len(vad_results) > 0:
+            vad_segments = [
+                {"start": segment.start, "end": segment.end, "confidence": 1.0}
+                for segment in vad_results.itersegments()
+            ]
+        return vad_segments

nodes/VoiceStreamAI/vad/vad_factory.py

@@ -0,0 +1,23 @@
+from VoiceStreamAI.vad.pyannote_vad import PyannoteVAD
+
+class VADFactory:
+    """
+    Factory for creating instances of VAD systems.
+    """
+
+    @staticmethod
+    def create_vad_pipeline(type, **kwargs):
+        """
+        Creates a VAD pipeline based on the specified type.
+
+        Args:
+            type (str): The type of VAD pipeline to create (e.g., 'pyannote').
+            kwargs: Additional arguments for the VAD pipeline creation.
+
+        Returns:
+            VADInterface: An instance of a class that implements VADInterface.
+        """
+        if type == "pyannote":
+            return PyannoteVAD(**kwargs)
+        else:
+            raise ValueError(f"Unknown VAD pipeline type: {type}")

nodes/VoiceStreamAI/vad/vad_interface.py

@@ -0,0 +1,16 @@
+class VADInterface:
+    """
+    Interface for voice activity detection (VAD) systems.
+    """
+
+    async def detect_activity(self, client):
+        """
+        Detects voice activity in the given audio data.
+
+        Args:
+            client (src.Client): The client to detect on
+
+        Returns:
+            List: VAD result, a list of objects containing "start", "end", "confidence"
+        """
+        raise NotImplementedError("This method should be implemented by subclasses.")

nodes/test.py

@@ -0,0 +1,8 @@
+import folder_paths
+
+# 外挂一个文件，用来编写新的节点
+def run(v):
+
+    output_dir = folder_paths.get_temp_directory()
+
+    print('1323',v,output_dir)

nodes/tsr/config.yaml

@@ -0,0 +1,38 @@
+cond_image_size: 512
+
+image_tokenizer_cls: tsr.models.tokenizers.image.DINOSingleImageTokenizer
+image_tokenizer:
+  pretrained_model_name_or_path: "facebook/dino-vitb16"
+
+tokenizer_cls: tsr.models.tokenizers.triplane.Triplane1DTokenizer
+tokenizer:
+  plane_size: 32
+  num_channels: 1024
+
+backbone_cls: tsr.models.transformer.transformer_1d.Transformer1D
+backbone:
+  in_channels: ${tokenizer.num_channels}
+  num_attention_heads: 16
+  attention_head_dim: 64
+  num_layers: 16
+  cross_attention_dim: 768
+
+post_processor_cls: tsr.models.network_utils.TriplaneUpsampleNetwork
+post_processor:
+  in_channels: 1024
+  out_channels: 40
+
+decoder_cls: tsr.models.network_utils.NeRFMLP
+decoder:
+  in_channels: 120 # 3 * 40
+  n_neurons: 64
+  n_hidden_layers: 9
+  activation: silu
+
+renderer_cls: tsr.models.nerf_renderer.TriplaneNeRFRenderer
+renderer:
+  radius: 0.87 # slightly larger than 0.5 * sqrt(3)
+  feature_reduction: concat
+  density_activation: exp
+  density_bias: -1.0
+  num_samples_per_ray: 128

nodes/tsr/models/isosurface.py

@@ -0,0 +1,51 @@
+from typing import Callable, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from skimage import measure
+
+
+class IsosurfaceHelper(nn.Module):
+    points_range: Tuple[float, float] = (0, 1)
+
+    @property
+    def grid_vertices(self) -> torch.FloatTensor:
+        raise NotImplementedError
+
+
+class MarchingCubeHelper(IsosurfaceHelper):
+    def __init__(self, resolution: int) -> None:
+        super().__init__()
+        self.resolution = resolution
+        #self.mc_func: Callable = marching_cubes
+        self._grid_vertices: Optional[torch.FloatTensor] = None
+
+    @property
+    def grid_vertices(self) -> torch.FloatTensor:
+        if self._grid_vertices is None:
+            # keep the vertices on CPU so that we can support very large resolution
+            x, y, z = (
+                torch.linspace(*self.points_range, self.resolution),
+                torch.linspace(*self.points_range, self.resolution),
+                torch.linspace(*self.points_range, self.resolution),
+            )
+            x, y, z = torch.meshgrid(x, y, z, indexing="ij")
+            verts = torch.cat(
+                [x.reshape(-1, 1), y.reshape(-1, 1), z.reshape(-1, 1)], dim=-1
+            ).reshape(-1, 3)
+            self._grid_vertices = verts
+        return self._grid_vertices
+
+    def forward(
+        self,
+        level: torch.FloatTensor,
+    ) -> Tuple[torch.FloatTensor, torch.LongTensor]:
+        level = -level.view(self.resolution, self.resolution, self.resolution)
+        v_pos, t_pos_idx, _, __ = measure.marching_cubes((level.detach().cpu() if level.is_cuda else level.detach()).numpy(), 0.0) #self.mc_func(level.detach(), 0.0)
+        v_pos = torch.from_numpy(v_pos.copy()).type(torch.FloatTensor).to(level.device)
+        t_pos_idx = torch.from_numpy(t_pos_idx.copy()).type(torch.LongTensor).to(level.device)
+        v_pos = v_pos[..., [0, 1, 2]]
+        t_pos_idx = t_pos_idx[..., [1, 0, 2]]
+        v_pos = v_pos / (self.resolution - 1.0)
+        return v_pos, t_pos_idx

nodes/tsr/models/nerf_renderer.py

@@ -0,0 +1,180 @@
+from dataclasses import dataclass
+from typing import Dict
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, reduce
+
+from ..utils import (
+    BaseModule,
+    chunk_batch,
+    get_activation,
+    rays_intersect_bbox,
+    scale_tensor,
+)
+
+
+class TriplaneNeRFRenderer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        radius: float
+
+        feature_reduction: str = "concat"
+        density_activation: str = "trunc_exp"
+        density_bias: float = -1.0
+        color_activation: str = "sigmoid"
+        num_samples_per_ray: int = 128
+        randomized: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        assert self.cfg.feature_reduction in ["concat", "mean"]
+        self.chunk_size = 0
+
+    def set_chunk_size(self, chunk_size: int):
+        assert (
+            chunk_size >= 0
+        ), "chunk_size must be a non-negative integer (0 for no chunking)."
+        self.chunk_size = chunk_size
+
+    def query_triplane(
+        self,
+        decoder: torch.nn.Module,
+        positions: torch.Tensor,
+        triplane: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        input_shape = positions.shape[:-1]
+        positions = positions.view(-1, 3)
+
+        # positions in (-radius, radius)
+        # normalized to (-1, 1) for grid sample
+        positions = scale_tensor(
+            positions, (-self.cfg.radius, self.cfg.radius), (-1, 1)
+        )
+
+        def _query_chunk(x):
+            indices2D: torch.Tensor = torch.stack(
+                (x[..., [0, 1]], x[..., [0, 2]], x[..., [1, 2]]),
+                dim=-3,
+            )
+            out: torch.Tensor = F.grid_sample(
+                rearrange(triplane, "Np Cp Hp Wp -> Np Cp Hp Wp", Np=3),
+                rearrange(indices2D, "Np N Nd -> Np () N Nd", Np=3),
+                align_corners=False,
+                mode="bilinear",
+            )
+            if self.cfg.feature_reduction == "concat":
+                out = rearrange(out, "Np Cp () N -> N (Np Cp)", Np=3)
+            elif self.cfg.feature_reduction == "mean":
+                out = reduce(out, "Np Cp () N -> N Cp", Np=3, reduction="mean")
+            else:
+                raise NotImplementedError
+
+            net_out: Dict[str, torch.Tensor] = decoder(out)
+            return net_out
+
+        if self.chunk_size > 0:
+            net_out = chunk_batch(_query_chunk, self.chunk_size, positions)
+        else:
+            net_out = _query_chunk(positions)
+
+        net_out["density_act"] = get_activation(self.cfg.density_activation)(
+            net_out["density"] + self.cfg.density_bias
+        )
+        net_out["color"] = get_activation(self.cfg.color_activation)(
+            net_out["features"]
+        )
+
+        net_out = {k: v.view(*input_shape, -1) for k, v in net_out.items()}
+
+        return net_out
+
+    def _forward(
+        self,
+        decoder: torch.nn.Module,
+        triplane: torch.Tensor,
+        rays_o: torch.Tensor,
+        rays_d: torch.Tensor,
+        **kwargs,
+    ):
+        rays_shape = rays_o.shape[:-1]
+        rays_o = rays_o.view(-1, 3)
+        rays_d = rays_d.view(-1, 3)
+        n_rays = rays_o.shape[0]
+
+        t_near, t_far, rays_valid = rays_intersect_bbox(rays_o, rays_d, self.cfg.radius)
+        t_near, t_far = t_near[rays_valid], t_far[rays_valid]
+
+        t_vals = torch.linspace(
+            0, 1, self.cfg.num_samples_per_ray + 1, device=triplane.device
+        )
+        t_mid = (t_vals[:-1] + t_vals[1:]) / 2.0
+        z_vals = t_near * (1 - t_mid[None]) + t_far * t_mid[None]  # (N_rays, N_samples)
+
+        xyz = (
+            rays_o[:, None, :] + z_vals[..., None] * rays_d[..., None, :]
+        )  # (N_rays, N_sample, 3)
+
+        mlp_out = self.query_triplane(
+            decoder=decoder,
+            positions=xyz,
+            triplane=triplane,
+        )
+
+        eps = 1e-10
+        # deltas = z_vals[:, 1:] - z_vals[:, :-1] # (N_rays, N_samples)
+        deltas = t_vals[1:] - t_vals[:-1]  # (N_rays, N_samples)
+        alpha = 1 - torch.exp(
+            -deltas * mlp_out["density_act"][..., 0]
+        )  # (N_rays, N_samples)
+        accum_prod = torch.cat(
+            [
+                torch.ones_like(alpha[:, :1]),
+                torch.cumprod(1 - alpha[:, :-1] + eps, dim=-1),
+            ],
+            dim=-1,
+        )
+        weights = alpha * accum_prod  # (N_rays, N_samples)
+        comp_rgb_ = (weights[..., None] * mlp_out["color"]).sum(dim=-2)  # (N_rays, 3)
+        opacity_ = weights.sum(dim=-1)  # (N_rays)
+
+        comp_rgb = torch.zeros(
+            n_rays, 3, dtype=comp_rgb_.dtype, device=comp_rgb_.device
+        )
+        opacity = torch.zeros(n_rays, dtype=opacity_.dtype, device=opacity_.device)
+        comp_rgb[rays_valid] = comp_rgb_
+        opacity[rays_valid] = opacity_
+
+        comp_rgb += 1 - opacity[..., None]
+        comp_rgb = comp_rgb.view(*rays_shape, 3)
+
+        return comp_rgb
+
+    def forward(
+        self,
+        decoder: torch.nn.Module,
+        triplane: torch.Tensor,
+        rays_o: torch.Tensor,
+        rays_d: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        if triplane.ndim == 4:
+            comp_rgb = self._forward(decoder, triplane, rays_o, rays_d)
+        else:
+            comp_rgb = torch.stack(
+                [
+                    self._forward(decoder, triplane[i], rays_o[i], rays_d[i])
+                    for i in range(triplane.shape[0])
+                ],
+                dim=0,
+            )
+
+        return comp_rgb
+
+    def train(self, mode=True):
+        self.randomized = mode and self.cfg.randomized
+        return super().train(mode=mode)
+
+    def eval(self):
+        self.randomized = False
+        return super().eval()

nodes/tsr/models/network_utils.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from ..utils import BaseModule
+
+
+class TriplaneUpsampleNetwork(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int
+        out_channels: int
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.upsample = nn.ConvTranspose2d(
+            self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2
+        )
+
+    def forward(self, triplanes: torch.Tensor) -> torch.Tensor:
+        triplanes_up = rearrange(
+            self.upsample(
+                rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
+            ),
+            "(B Np) Co Hp Wp -> B Np Co Hp Wp",
+            Np=3,
+        )
+        return triplanes_up
+
+
+class NeRFMLP(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int
+        n_neurons: int
+        n_hidden_layers: int
+        activation: str = "relu"
+        bias: bool = True
+        weight_init: Optional[str] = "kaiming_uniform"
+        bias_init: Optional[str] = None
+
+    cfg: Config
+
+    def configure(self) -> None:
+        layers = [
+            self.make_linear(
+                self.cfg.in_channels,
+                self.cfg.n_neurons,
+                bias=self.cfg.bias,
+                weight_init=self.cfg.weight_init,
+                bias_init=self.cfg.bias_init,
+            ),
+            self.make_activation(self.cfg.activation),
+        ]
+        for i in range(self.cfg.n_hidden_layers - 1):
+            layers += [
+                self.make_linear(
+                    self.cfg.n_neurons,
+                    self.cfg.n_neurons,
+                    bias=self.cfg.bias,
+                    weight_init=self.cfg.weight_init,
+                    bias_init=self.cfg.bias_init,
+                ),
+                self.make_activation(self.cfg.activation),
+            ]
+        layers += [
+            self.make_linear(
+                self.cfg.n_neurons,
+                4,  # density 1 + features 3
+                bias=self.cfg.bias,
+                weight_init=self.cfg.weight_init,
+                bias_init=self.cfg.bias_init,
+            )
+        ]
+        self.layers = nn.Sequential(*layers)
+
+    def make_linear(
+        self,
+        dim_in,
+        dim_out,
+        bias=True,
+        weight_init=None,
+        bias_init=None,
+    ):
+        layer = nn.Linear(dim_in, dim_out, bias=bias)
+
+        if weight_init is None:
+            pass
+        elif weight_init == "kaiming_uniform":
+            torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
+        else:
+            raise NotImplementedError
+
+        if bias:
+            if bias_init is None:
+                pass
+            elif bias_init == "zero":
+                torch.nn.init.zeros_(layer.bias)
+            else:
+                raise NotImplementedError
+
+        return layer
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def forward(self, x):
+        inp_shape = x.shape[:-1]
+        x = x.reshape(-1, x.shape[-1])
+
+        features = self.layers(x)
+        features = features.reshape(*inp_shape, -1)
+        out = {"density": features[..., 0:1], "features": features[..., 1:4]}
+
+        return out

nodes/tsr/models/tokenizers/image.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from transformers.models.vit.modeling_vit import ViTModel
+
+from ...utils import BaseModule
+import os
+import folder_paths
+model_path=os.path.join(folder_paths.models_dir,'triposr')
+
+
+class DINOSingleImageTokenizer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pretrained_model_name_or_path: str = "facebook/dino-vitb16"
+        enable_gradient_checkpointing: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        print('#Loading ViTModel:',os.path.join(model_path,self.cfg.pretrained_model_name_or_path))
+        self.model: ViTModel = ViTModel(
+            ViTModel.config_class.from_pretrained(
+                hf_hub_download(
+                    repo_id=self.cfg.pretrained_model_name_or_path,
+                    filename="config.json",
+                    local_dir=model_path,
+                    endpoint='https://hf-mirror.com'
+                )
+            )
+        )
+
+        if self.cfg.enable_gradient_checkpointing:
+            self.model.encoder.gradient_checkpointing = True
+
+        self.register_buffer(
+            "image_mean",
+            torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+        self.register_buffer(
+            "image_std",
+            torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+
+    def forward(self, images: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        packed = False
+        if images.ndim == 4:
+            packed = True
+            images = images.unsqueeze(1)
+
+        batch_size, n_input_views = images.shape[:2]
+        images = (images - self.image_mean) / self.image_std
+        out = self.model(
+            rearrange(images, "B N C H W -> (B N) C H W"), interpolate_pos_encoding=True
+        )
+        local_features, global_features = out.last_hidden_state, out.pooler_output
+        local_features = local_features.permute(0, 2, 1)
+        local_features = rearrange(
+            local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size
+        )
+        if packed:
+            local_features = local_features.squeeze(1)
+
+        return local_features
+
+    def detokenize(self, *args, **kwargs):
+        raise NotImplementedError

nodes/tsr/models/tokenizers/triplane.py

@@ -0,0 +1,45 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+from ...utils import BaseModule
+
+
+class Triplane1DTokenizer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        plane_size: int
+        num_channels: int
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.embeddings = nn.Parameter(
+            torch.randn(
+                (3, self.cfg.num_channels, self.cfg.plane_size, self.cfg.plane_size),
+                dtype=torch.float32,
+            )
+            * 1
+            / math.sqrt(self.cfg.num_channels)
+        )
+
+    def forward(self, batch_size: int) -> torch.Tensor:
+        return rearrange(
+            repeat(self.embeddings, "Np Ct Hp Wp -> B Np Ct Hp Wp", B=batch_size),
+            "B Np Ct Hp Wp -> B Ct (Np Hp Wp)",
+        )
+
+    def detokenize(self, tokens: torch.Tensor) -> torch.Tensor:
+        batch_size, Ct, Nt = tokens.shape
+        assert Nt == self.cfg.plane_size**2 * 3
+        assert Ct == self.cfg.num_channels
+        return rearrange(
+            tokens,
+            "B Ct (Np Hp Wp) -> B Np Ct Hp Wp",
+            Np=3,
+            Hp=self.cfg.plane_size,
+            Wp=self.cfg.plane_size,
+        )

nodes/tsr/models/transformer/attention.py

@@ -0,0 +1,653 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# --------
+#
+# Modified 2024 by the Tripo AI and Stability AI Team.
+#
+# Copyright (c) 2024 Tripo AI & Stability AI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Attention(nn.Module):
+    r"""
+    A cross attention layer.
+
+    Parameters:
+        query_dim (`int`):
+            The number of channels in the query.
+        cross_attention_dim (`int`, *optional*):
+            The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.
+        heads (`int`,  *optional*, defaults to 8):
+            The number of heads to use for multi-head attention.
+        dim_head (`int`,  *optional*, defaults to 64):
+            The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0):
+            The dropout probability to use.
+        bias (`bool`, *optional*, defaults to False):
+            Set to `True` for the query, key, and value linear layers to contain a bias parameter.
+        upcast_attention (`bool`, *optional*, defaults to False):
+            Set to `True` to upcast the attention computation to `float32`.
+        upcast_softmax (`bool`, *optional*, defaults to False):
+            Set to `True` to upcast the softmax computation to `float32`.
+        cross_attention_norm (`str`, *optional*, defaults to `None`):
+            The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
+        cross_attention_norm_num_groups (`int`, *optional*, defaults to 32):
+            The number of groups to use for the group norm in the cross attention.
+        added_kv_proj_dim (`int`, *optional*, defaults to `None`):
+            The number of channels to use for the added key and value projections. If `None`, no projection is used.
+        norm_num_groups (`int`, *optional*, defaults to `None`):
+            The number of groups to use for the group norm in the attention.
+        spatial_norm_dim (`int`, *optional*, defaults to `None`):
+            The number of channels to use for the spatial normalization.
+        out_bias (`bool`, *optional*, defaults to `True`):
+            Set to `True` to use a bias in the output linear layer.
+        scale_qk (`bool`, *optional*, defaults to `True`):
+            Set to `True` to scale the query and key by `1 / sqrt(dim_head)`.
+        only_cross_attention (`bool`, *optional*, defaults to `False`):
+            Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if
+            `added_kv_proj_dim` is not `None`.
+        eps (`float`, *optional*, defaults to 1e-5):
+            An additional value added to the denominator in group normalization that is used for numerical stability.
+        rescale_output_factor (`float`, *optional*, defaults to 1.0):
+            A factor to rescale the output by dividing it with this value.
+        residual_connection (`bool`, *optional*, defaults to `False`):
+            Set to `True` to add the residual connection to the output.
+        _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`):
+            Set to `True` if the attention block is loaded from a deprecated state dict.
+        processor (`AttnProcessor`, *optional*, defaults to `None`):
+            The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and
+            `AttnProcessor` otherwise.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: Optional[int] = None,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        upcast_attention: bool = False,
+        upcast_softmax: bool = False,
+        cross_attention_norm: Optional[str] = None,
+        cross_attention_norm_num_groups: int = 32,
+        added_kv_proj_dim: Optional[int] = None,
+        norm_num_groups: Optional[int] = None,
+        out_bias: bool = True,
+        scale_qk: bool = True,
+        only_cross_attention: bool = False,
+        eps: float = 1e-5,
+        rescale_output_factor: float = 1.0,
+        residual_connection: bool = False,
+        _from_deprecated_attn_block: bool = False,
+        processor: Optional["AttnProcessor"] = None,
+        out_dim: int = None,
+    ):
+        super().__init__()
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.cross_attention_dim = (
+            cross_attention_dim if cross_attention_dim is not None else query_dim
+        )
+        self.upcast_attention = upcast_attention
+        self.upcast_softmax = upcast_softmax
+        self.rescale_output_factor = rescale_output_factor
+        self.residual_connection = residual_connection
+        self.dropout = dropout
+        self.fused_projections = False
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        # we make use of this private variable to know whether this class is loaded
+        # with an deprecated state dict so that we can convert it on the fly
+        self._from_deprecated_attn_block = _from_deprecated_attn_block
+
+        self.scale_qk = scale_qk
+        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        # for slice_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set slice_size with `set_attention_slice`
+        self.sliceable_head_dim = heads
+
+        self.added_kv_proj_dim = added_kv_proj_dim
+        self.only_cross_attention = only_cross_attention
+
+        if self.added_kv_proj_dim is None and self.only_cross_attention:
+            raise ValueError(
+                "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
+            )
+
+        if norm_num_groups is not None:
+            self.group_norm = nn.GroupNorm(
+                num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True
+            )
+        else:
+            self.group_norm = None
+
+        self.spatial_norm = None
+
+        if cross_attention_norm is None:
+            self.norm_cross = None
+        elif cross_attention_norm == "layer_norm":
+            self.norm_cross = nn.LayerNorm(self.cross_attention_dim)
+        elif cross_attention_norm == "group_norm":
+            if self.added_kv_proj_dim is not None:
+                # The given `encoder_hidden_states` are initially of shape
+                # (batch_size, seq_len, added_kv_proj_dim) before being projected
+                # to (batch_size, seq_len, cross_attention_dim). The norm is applied
+                # before the projection, so we need to use `added_kv_proj_dim` as
+                # the number of channels for the group norm.
+                norm_cross_num_channels = added_kv_proj_dim
+            else:
+                norm_cross_num_channels = self.cross_attention_dim
+
+            self.norm_cross = nn.GroupNorm(
+                num_channels=norm_cross_num_channels,
+                num_groups=cross_attention_norm_num_groups,
+                eps=1e-5,
+                affine=True,
+            )
+        else:
+            raise ValueError(
+                f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
+            )
+
+        linear_cls = nn.Linear
+
+        self.linear_cls = linear_cls
+        self.to_q = linear_cls(query_dim, self.inner_dim, bias=bias)
+
+        if not self.only_cross_attention:
+            # only relevant for the `AddedKVProcessor` classes
+            self.to_k = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
+            self.to_v = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
+        else:
+            self.to_k = None
+            self.to_v = None
+
+        if self.added_kv_proj_dim is not None:
+            self.add_k_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
+            self.add_v_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(linear_cls(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(nn.Dropout(dropout))
+
+        # set attention processor
+        # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+        # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+        # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+        if processor is None:
+            processor = (
+                AttnProcessor2_0()
+                if hasattr(F, "scaled_dot_product_attention") and self.scale_qk
+                else AttnProcessor()
+            )
+        self.set_processor(processor)
+
+    def set_processor(self, processor: "AttnProcessor") -> None:
+        self.processor = processor
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        **cross_attention_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        The forward method of the `Attention` class.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                The hidden states of the query.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                The hidden states of the encoder.
+            attention_mask (`torch.Tensor`, *optional*):
+                The attention mask to use. If `None`, no mask is applied.
+            **cross_attention_kwargs:
+                Additional keyword arguments to pass along to the cross attention.
+
+        Returns:
+            `torch.Tensor`: The output of the attention layer.
+        """
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**cross_attention_kwargs` is empty
+        return self.processor(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs,
+        )
+
+    def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:
+        r"""
+        Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads`
+        is the number of heads initialized while constructing the `Attention` class.
+
+        Args:
+            tensor (`torch.Tensor`): The tensor to reshape.
+
+        Returns:
+            `torch.Tensor`: The reshaped tensor.
+        """
+        head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(
+            batch_size // head_size, seq_len, dim * head_size
+        )
+        return tensor
+
+    def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:
+        r"""
+        Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is
+        the number of heads initialized while constructing the `Attention` class.
+
+        Args:
+            tensor (`torch.Tensor`): The tensor to reshape.
+            out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is
+                reshaped to `[batch_size * heads, seq_len, dim // heads]`.
+
+        Returns:
+            `torch.Tensor`: The reshaped tensor.
+        """
+        head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3)
+
+        if out_dim == 3:
+            tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)
+
+        return tensor
+
+    def get_attention_scores(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attention_mask: torch.Tensor = None,
+    ) -> torch.Tensor:
+        r"""
+        Compute the attention scores.
+
+        Args:
+            query (`torch.Tensor`): The query tensor.
+            key (`torch.Tensor`): The key tensor.
+            attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
+
+        Returns:
+            `torch.Tensor`: The attention probabilities/scores.
+        """
+        dtype = query.dtype
+        if self.upcast_attention:
+            query = query.float()
+            key = key.float()
+
+        if attention_mask is None:
+            baddbmm_input = torch.empty(
+                query.shape[0],
+                query.shape[1],
+                key.shape[1],
+                dtype=query.dtype,
+                device=query.device,
+            )
+            beta = 0
+        else:
+            baddbmm_input = attention_mask
+            beta = 1
+
+        attention_scores = torch.baddbmm(
+            baddbmm_input,
+            query,
+            key.transpose(-1, -2),
+            beta=beta,
+            alpha=self.scale,
+        )
+        del baddbmm_input
+
+        if self.upcast_softmax:
+            attention_scores = attention_scores.float()
+
+        attention_probs = attention_scores.softmax(dim=-1)
+        del attention_scores
+
+        attention_probs = attention_probs.to(dtype)
+
+        return attention_probs
+
+    def prepare_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        target_length: int,
+        batch_size: int,
+        out_dim: int = 3,
+    ) -> torch.Tensor:
+        r"""
+        Prepare the attention mask for the attention computation.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                The attention mask to prepare.
+            target_length (`int`):
+                The target length of the attention mask. This is the length of the attention mask after padding.
+            batch_size (`int`):
+                The batch size, which is used to repeat the attention mask.
+            out_dim (`int`, *optional*, defaults to `3`):
+                The output dimension of the attention mask. Can be either `3` or `4`.
+
+        Returns:
+            `torch.Tensor`: The prepared attention mask.
+        """
+        head_size = self.heads
+        if attention_mask is None:
+            return attention_mask
+
+        current_length: int = attention_mask.shape[-1]
+        if current_length != target_length:
+            if attention_mask.device.type == "mps":
+                # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+                # Instead, we can manually construct the padding tensor.
+                padding_shape = (
+                    attention_mask.shape[0],
+                    attention_mask.shape[1],
+                    target_length,
+                )
+                padding = torch.zeros(
+                    padding_shape,
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device,
+                )
+                attention_mask = torch.cat([attention_mask, padding], dim=2)
+            else:
+                # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+                #       we want to instead pad by (0, remaining_length), where remaining_length is:
+                #       remaining_length: int = target_length - current_length
+                # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+                attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
+
+        if out_dim == 3:
+            if attention_mask.shape[0] < batch_size * head_size:
+                attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
+        elif out_dim == 4:
+            attention_mask = attention_mask.unsqueeze(1)
+            attention_mask = attention_mask.repeat_interleave(head_size, dim=1)
+
+        return attention_mask
+
+    def norm_encoder_hidden_states(
+        self, encoder_hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        r"""
+        Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the
+        `Attention` class.
+
+        Args:
+            encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder.
+
+        Returns:
+            `torch.Tensor`: The normalized encoder hidden states.
+        """
+        assert (
+            self.norm_cross is not None
+        ), "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
+
+        if isinstance(self.norm_cross, nn.LayerNorm):
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+        elif isinstance(self.norm_cross, nn.GroupNorm):
+            # Group norm norms along the channels dimension and expects
+            # input to be in the shape of (N, C, *). In this case, we want
+            # to norm along the hidden dimension, so we need to move
+            # (batch_size, sequence_length, hidden_size) ->
+            # (batch_size, hidden_size, sequence_length)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+        else:
+            assert False
+
+        return encoder_hidden_states
+
+    @torch.no_grad()
+    def fuse_projections(self, fuse=True):
+        is_cross_attention = self.cross_attention_dim != self.query_dim
+        device = self.to_q.weight.data.device
+        dtype = self.to_q.weight.data.dtype
+
+        if not is_cross_attention:
+            # fetch weight matrices.
+            concatenated_weights = torch.cat(
+                [self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]
+            )
+            in_features = concatenated_weights.shape[1]
+            out_features = concatenated_weights.shape[0]
+
+            # create a new single projection layer and copy over the weights.
+            self.to_qkv = self.linear_cls(
+                in_features, out_features, bias=False, device=device, dtype=dtype
+            )
+            self.to_qkv.weight.copy_(concatenated_weights)
+
+        else:
+            concatenated_weights = torch.cat(
+                [self.to_k.weight.data, self.to_v.weight.data]
+            )
+            in_features = concatenated_weights.shape[1]
+            out_features = concatenated_weights.shape[0]
+
+            self.to_kv = self.linear_cls(
+                in_features, out_features, bias=False, device=device, dtype=dtype
+            )
+            self.to_kv.weight.copy_(concatenated_weights)
+
+        self.fused_projections = fuse
+
+
+class AttnProcessor:
+    r"""
+    Default processor for performing attention-related computations.
+    """
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask, sequence_length, batch_size
+        )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+class AttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states

nodes/tsr/models/transformer/basic_transformer_block.py

@@ -0,0 +1,334 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# --------
+#
+# Modified 2024 by the Tripo AI and Stability AI Team.
+#
+# Copyright (c) 2024 Tripo AI & Stability AI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .attention import Attention
+
+
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        upcast_attention (`bool`, *optional*):
+            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
+        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
+        final_dropout (`bool` *optional*, defaults to False):
+            Whether to apply a final dropout after the last feed-forward layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+
+        assert norm_type == "layer_norm"
+
+        # Define 3 blocks. Each block has its own normalization layer.
+        # 1. Self-Attn
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+        )
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None or double_self_attention:
+            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+            # the second cross attention block.
+            self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=(
+                    cross_attention_dim if not double_self_attention else None
+                ),
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+            )  # is self-attn if encoder_hidden_states is none
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        # 3. Feed-forward
+        self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+        )
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 0. Self-Attention
+        norm_hidden_states = self.norm1(hidden_states)
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=(
+                encoder_hidden_states if self.only_cross_attention else None
+            ),
+            attention_mask=attention_mask,
+        )
+
+        hidden_states = attn_output + hidden_states
+
+        # 3. Cross-Attention
+        if self.attn2 is not None:
+            norm_hidden_states = self.norm2(hidden_states)
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+            )
+            hidden_states = attn_output + hidden_states
+
+        # 4. Feed-forward
+        norm_hidden_states = self.norm3(hidden_states)
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
+                raise ValueError(
+                    f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+                )
+
+            num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
+            ff_output = torch.cat(
+                [
+                    self.ff(hid_slice)
+                    for hid_slice in norm_hidden_states.chunk(
+                        num_chunks, dim=self._chunk_dim
+                    )
+                ],
+                dim=self._chunk_dim,
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        hidden_states = ff_output + hidden_states
+
+        return hidden_states
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+        activation_fn: str = "geglu",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        linear_cls = nn.Linear
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh")
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim)
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # project dropout
+        self.net.append(nn.Dropout(dropout))
+        # project out
+        self.net.append(linear_cls(inner_dim, dim_out))
+        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
+        if final_dropout:
+            self.net.append(nn.Dropout(dropout))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+class GELU(nn.Module):
+    r"""
+    GELU activation function with tanh approximation support with `approximate="tanh"`.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+        self.approximate = approximate
+
+    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+        if gate.device.type != "mps":
+            return F.gelu(gate, approximate=self.approximate)
+        # mps: gelu is not implemented for float16
+        return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(
+            dtype=gate.dtype
+        )
+
+    def forward(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        hidden_states = self.gelu(hidden_states)
+        return hidden_states
+
+
+class GEGLU(nn.Module):
+    r"""
+    A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int):
+        super().__init__()
+        linear_cls = nn.Linear
+
+        self.proj = linear_cls(dim_in, dim_out * 2)
+
+    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+        if gate.device.type != "mps":
+            return F.gelu(gate)
+        # mps: gelu is not implemented for float16
+        return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
+
+    def forward(self, hidden_states, scale: float = 1.0):
+        args = ()
+        hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1)
+        return hidden_states * self.gelu(gate)
+
+
+class ApproximateGELU(nn.Module):
+    r"""
+    The approximate form of Gaussian Error Linear Unit (GELU). For more details, see section 2:
+    https://arxiv.org/abs/1606.08415.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        return x * torch.sigmoid(1.702 * x)

nodes/tsr/models/transformer/transformer_1d.py

@@ -0,0 +1,219 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# --------
+#
+# Modified 2024 by the Tripo AI and Stability AI Team.
+#
+# Copyright (c) 2024 Tripo AI & Stability AI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ...utils import BaseModule
+from .basic_transformer_block import BasicTransformerBlock
+
+
+class Transformer1D(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        num_attention_heads: int = 16
+        attention_head_dim: int = 88
+        in_channels: Optional[int] = None
+        out_channels: Optional[int] = None
+        num_layers: int = 1
+        dropout: float = 0.0
+        norm_num_groups: int = 32
+        cross_attention_dim: Optional[int] = None
+        attention_bias: bool = False
+        activation_fn: str = "geglu"
+        only_cross_attention: bool = False
+        double_self_attention: bool = False
+        upcast_attention: bool = False
+        norm_type: str = "layer_norm"
+        norm_elementwise_affine: bool = True
+        gradient_checkpointing: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.num_attention_heads = self.cfg.num_attention_heads
+        self.attention_head_dim = self.cfg.attention_head_dim
+        inner_dim = self.num_attention_heads * self.attention_head_dim
+
+        linear_cls = nn.Linear
+
+        # 2. Define input layers
+        self.in_channels = self.cfg.in_channels
+
+        self.norm = torch.nn.GroupNorm(
+            num_groups=self.cfg.norm_num_groups,
+            num_channels=self.cfg.in_channels,
+            eps=1e-6,
+            affine=True,
+        )
+        self.proj_in = linear_cls(self.cfg.in_channels, inner_dim)
+
+        # 3. Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    self.num_attention_heads,
+                    self.attention_head_dim,
+                    dropout=self.cfg.dropout,
+                    cross_attention_dim=self.cfg.cross_attention_dim,
+                    activation_fn=self.cfg.activation_fn,
+                    attention_bias=self.cfg.attention_bias,
+                    only_cross_attention=self.cfg.only_cross_attention,
+                    double_self_attention=self.cfg.double_self_attention,
+                    upcast_attention=self.cfg.upcast_attention,
+                    norm_type=self.cfg.norm_type,
+                    norm_elementwise_affine=self.cfg.norm_elementwise_affine,
+                )
+                for d in range(self.cfg.num_layers)
+            ]
+        )
+
+        # 4. Define output layers
+        self.out_channels = (
+            self.cfg.in_channels
+            if self.cfg.out_channels is None
+            else self.cfg.out_channels
+        )
+
+        self.proj_out = linear_cls(inner_dim, self.cfg.in_channels)
+
+        self.gradient_checkpointing = self.cfg.gradient_checkpointing
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        The [`Transformer1DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
+                Input `hidden_states`.
+            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
+                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
+                self-attention.
+            attention_mask ( `torch.Tensor`, *optional*):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            encoder_attention_mask ( `torch.Tensor`, *optional*):
+                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
+
+                    * Mask `(batch, sequence_length)` True = keep, False = discard.
+                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
+
+                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
+                above. This bias will be added to the cross-attention scores.
+
+        Returns:
+            torch.FloatTensor
+        """
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+        #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+        #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None and attention_mask.ndim == 2:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+            encoder_attention_mask = (
+                1 - encoder_attention_mask.to(hidden_states.dtype)
+            ) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 1. Input
+        batch, _, seq_len = hidden_states.shape
+        residual = hidden_states
+
+        hidden_states = self.norm(hidden_states)
+        inner_dim = hidden_states.shape[1]
+        hidden_states = hidden_states.permute(0, 2, 1).reshape(
+            batch, seq_len, inner_dim
+        )
+        hidden_states = self.proj_in(hidden_states)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            if self.training and self.gradient_checkpointing:
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    block,
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+
+        # 3. Output
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = (
+            hidden_states.reshape(batch, seq_len, inner_dim)
+            .permute(0, 2, 1)
+            .contiguous()
+        )
+
+        output = hidden_states + residual
+
+        return output

nodes/tsr/system.py

@@ -0,0 +1,218 @@
+import math
+import os
+from dataclasses import dataclass, field
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+import trimesh
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from omegaconf import OmegaConf
+from PIL import Image
+
+from .models.isosurface import MarchingCubeHelper
+from .utils import (
+    BaseModule,
+    ImagePreprocessor,
+    find_class,
+    get_spherical_cameras,
+    scale_tensor,
+)
+
+
+class TSR(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        cond_image_size: int
+
+        image_tokenizer_cls: str
+        image_tokenizer: dict
+
+        tokenizer_cls: str
+        tokenizer: dict
+
+        backbone_cls: str
+        backbone: dict
+
+        post_processor_cls: str
+        post_processor: dict
+
+        decoder_cls: str
+        decoder: dict
+
+        renderer_cls: str
+        renderer: dict
+
+    cfg: Config
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: str, config_name: str, weight_name: str
+    ):
+        if os.path.isdir(pretrained_model_name_or_path):
+            config_path = os.path.join(pretrained_model_name_or_path, config_name)
+            weight_path = os.path.join(pretrained_model_name_or_path, weight_name)
+        else:
+            config_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path, filename=config_name
+            )
+            weight_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path, filename=weight_name
+            )
+
+        cfg = OmegaConf.load(config_path)
+        OmegaConf.resolve(cfg)
+        model = cls(cfg)
+        ckpt = torch.load(weight_path, map_location="cpu")
+        model.load_state_dict(ckpt)
+        return model
+
+    @classmethod
+    def from_pretrained_custom(
+            cls, weight_path: str, config_path: str
+    ):
+        cfg = OmegaConf.load(config_path)
+        OmegaConf.resolve(cfg)
+        model = cls(cfg)
+        ckpt = torch.load(weight_path, map_location="cpu")
+        model.load_state_dict(ckpt)
+        return model
+
+    def configure(self):
+        self.image_tokenizer = find_class(self.cfg.image_tokenizer_cls)(
+            self.cfg.image_tokenizer
+        )
+        self.tokenizer = find_class(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
+        self.backbone = find_class(self.cfg.backbone_cls)(self.cfg.backbone)
+        self.post_processor = find_class(self.cfg.post_processor_cls)(
+            self.cfg.post_processor
+        )
+        self.decoder = find_class(self.cfg.decoder_cls)(self.cfg.decoder)
+        self.renderer = find_class(self.cfg.renderer_cls)(self.cfg.renderer)
+        self.image_processor = ImagePreprocessor()
+        self.isosurface_helper = None
+
+    def forward(
+        self,
+        image: Union[
+            PIL.Image.Image,
+            np.ndarray,
+            torch.FloatTensor,
+            List[PIL.Image.Image],
+            List[np.ndarray],
+            List[torch.FloatTensor],
+        ],
+        device: str,
+    ) -> torch.FloatTensor:
+        rgb_cond = self.image_processor(image, self.cfg.cond_image_size)[:, None].to(
+            device
+        )
+        batch_size = rgb_cond.shape[0]
+
+        input_image_tokens: torch.Tensor = self.image_tokenizer(
+            rearrange(rgb_cond, "B Nv H W C -> B Nv C H W", Nv=1),
+        )
+
+        input_image_tokens = rearrange(
+            input_image_tokens, "B Nv C Nt -> B (Nv Nt) C", Nv=1
+        )
+
+        tokens: torch.Tensor = self.tokenizer(batch_size)
+
+        tokens = self.backbone(
+            tokens,
+            encoder_hidden_states=input_image_tokens,
+        )
+
+        scene_codes = self.post_processor(self.tokenizer.detokenize(tokens))
+        return scene_codes
+
+    def render(
+        self,
+        scene_codes,
+        n_views: int,
+        elevation_deg: float = 0.0,
+        camera_distance: float = 1.9,
+        fovy_deg: float = 40.0,
+        height: int = 256,
+        width: int = 256,
+        return_type: str = "pil",
+    ):
+        rays_o, rays_d = get_spherical_cameras(
+            n_views, elevation_deg, camera_distance, fovy_deg, height, width
+        )
+        rays_o, rays_d = rays_o.to(scene_codes.device), rays_d.to(scene_codes.device)
+
+        def process_output(image: torch.FloatTensor):
+            if return_type == "pt":
+                return image
+            elif return_type == "np":
+                return image.detach().cpu().numpy()
+            elif return_type == "pil":
+                return Image.fromarray(
+                    (image.detach().cpu().numpy() * 255.0).astype(np.uint8)
+                )
+            else:
+                raise NotImplementedError
+
+        images = []
+        for scene_code in scene_codes:
+            images_ = []
+            for i in range(n_views):
+                with torch.no_grad():
+                    image = self.renderer(
+                        self.decoder, scene_code, rays_o[i], rays_d[i]
+                    )
+                images_.append(process_output(image))
+            images.append(images_)
+
+        return images
+
+    def set_marching_cubes_resolution(self, resolution: int):
+        if (
+            self.isosurface_helper is not None
+            and self.isosurface_helper.resolution == resolution
+        ):
+            return
+        self.isosurface_helper = MarchingCubeHelper(resolution)
+
+    def extract_mesh(self, scene_codes, resolution: int = 256, threshold: float = 25.0,callback=None):
+        self.set_marching_cubes_resolution(resolution)
+        meshes = []
+        for scene_code in scene_codes:
+            with torch.no_grad():
+                density = self.renderer.query_triplane(
+                    self.decoder,
+                    scale_tensor(
+                        self.isosurface_helper.grid_vertices.to(scene_codes.device),
+                        self.isosurface_helper.points_range,
+                        (-self.renderer.cfg.radius, self.renderer.cfg.radius),
+                    ),
+                    scene_code,
+                )["density_act"]
+            v_pos, t_pos_idx = self.isosurface_helper(-(density - threshold))
+            v_pos = scale_tensor(
+                v_pos,
+                self.isosurface_helper.points_range,
+                (-self.renderer.cfg.radius, self.renderer.cfg.radius),
+            )
+            with torch.no_grad():
+                color = self.renderer.query_triplane(
+                    self.decoder,
+                    v_pos,
+                    scene_code,
+                )["color"]
+            mesh = trimesh.Trimesh(
+                vertices=v_pos.cpu().numpy(),
+                faces=t_pos_idx.cpu().numpy(),
+                vertex_colors=color.cpu().numpy(),
+            )
+            meshes.append(mesh)
+
+            if callback:
+                callback(len(meshes))
+
+        return meshes

nodes/tsr/utils.py

@@ -0,0 +1,475 @@
+import importlib
+import math
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import imageio
+import numpy as np
+import PIL.Image
+#import rembg
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import trimesh
+from omegaconf import DictConfig, OmegaConf
+#from PIL import Image
+
+
+def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
+    scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg)
+    return scfg
+
+
+def find_class(cls_string):
+    module_string = ".".join(cls_string.split(".")[:-1])
+    cls_name = cls_string.split(".")[-1]
+    module = importlib.import_module(module_string, package=None)
+    cls = getattr(module, cls_name)
+    return cls
+
+
+def get_intrinsic_from_fov(fov, H, W, bs=-1):
+    focal_length = 0.5 * H / np.tan(0.5 * fov)
+    intrinsic = np.identity(3, dtype=np.float32)
+    intrinsic[0, 0] = focal_length
+    intrinsic[1, 1] = focal_length
+    intrinsic[0, 2] = W / 2.0
+    intrinsic[1, 2] = H / 2.0
+
+    if bs > 0:
+        intrinsic = intrinsic[None].repeat(bs, axis=0)
+
+    return torch.from_numpy(intrinsic)
+
+
+class BaseModule(nn.Module):
+    @dataclass
+    class Config:
+        pass
+
+    cfg: Config  # add this to every subclass of BaseModule to enable static type checking
+
+    def __init__(
+        self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
+    ) -> None:
+        super().__init__()
+        self.cfg = parse_structured(self.Config, cfg)
+        self.configure(*args, **kwargs)
+
+    def configure(self, *args, **kwargs) -> None:
+        raise NotImplementedError
+
+
+class ImagePreprocessor:
+    def convert_and_resize(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+        size: int,
+    ):
+        if isinstance(image, PIL.Image.Image):
+            image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
+        elif isinstance(image, np.ndarray):
+            if image.dtype == np.uint8:
+                image = torch.from_numpy(image.astype(np.float32) / 255.0)
+            else:
+                image = torch.from_numpy(image)
+        elif isinstance(image, torch.Tensor):
+            pass
+
+        batched = image.ndim == 4
+
+        if not batched:
+            image = image[None, ...]
+        image = F.interpolate(
+            image.permute(0, 3, 1, 2),
+            (size, size),
+            mode="bilinear",
+            align_corners=False,
+            antialias=True,
+        ).permute(0, 2, 3, 1)
+        if not batched:
+            image = image[0]
+        return image
+
+    def __call__(
+        self,
+        image: Union[
+            PIL.Image.Image,
+            np.ndarray,
+            torch.FloatTensor,
+            List[PIL.Image.Image],
+            List[np.ndarray],
+            List[torch.FloatTensor],
+        ],
+        size: int,
+    ) -> Any:
+        if isinstance(image, (np.ndarray, torch.FloatTensor)) and image.ndim == 4:
+            image = self.convert_and_resize(image, size)
+        else:
+            if not isinstance(image, list):
+                image = [image]
+            image = [self.convert_and_resize(im, size) for im in image]
+            image = torch.stack(image, dim=0)
+        return image
+
+
+def rays_intersect_bbox(
+    rays_o: torch.Tensor,
+    rays_d: torch.Tensor,
+    radius: float,
+    near: float = 0.0,
+    valid_thresh: float = 0.01,
+):
+    input_shape = rays_o.shape[:-1]
+    rays_o, rays_d = rays_o.view(-1, 3), rays_d.view(-1, 3)
+    rays_d_valid = torch.where(
+        rays_d.abs() < 1e-6, torch.full_like(rays_d, 1e-6), rays_d
+    )
+    if type(radius) in [int, float]:
+        radius = torch.FloatTensor(
+            [[-radius, radius], [-radius, radius], [-radius, radius]]
+        ).to(rays_o.device)
+    radius = (
+        1.0 - 1.0e-3
+    ) * radius  # tighten the radius to make sure the intersection point lies in the bounding box
+    interx0 = (radius[..., 1] - rays_o) / rays_d_valid
+    interx1 = (radius[..., 0] - rays_o) / rays_d_valid
+    t_near = torch.minimum(interx0, interx1).amax(dim=-1).clamp_min(near)
+    t_far = torch.maximum(interx0, interx1).amin(dim=-1)
+
+    # check wheter a ray intersects the bbox or not
+    rays_valid = t_far - t_near > valid_thresh
+
+    t_near[torch.where(~rays_valid)] = 0.0
+    t_far[torch.where(~rays_valid)] = 0.0
+
+    t_near = t_near.view(*input_shape, 1)
+    t_far = t_far.view(*input_shape, 1)
+    rays_valid = rays_valid.view(*input_shape)
+
+    return t_near, t_far, rays_valid
+
+
+def chunk_batch(func: Callable, chunk_size: int, *args, **kwargs) -> Any:
+    if chunk_size <= 0:
+        return func(*args, **kwargs)
+    B = None
+    for arg in list(args) + list(kwargs.values()):
+        if isinstance(arg, torch.Tensor):
+            B = arg.shape[0]
+            break
+    assert (
+        B is not None
+    ), "No tensor found in args or kwargs, cannot determine batch size."
+    out = defaultdict(list)
+    out_type = None
+    # max(1, B) to support B == 0
+    for i in range(0, max(1, B), chunk_size):
+        out_chunk = func(
+            *[
+                arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
+                for arg in args
+            ],
+            **{
+                k: arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
+                for k, arg in kwargs.items()
+            },
+        )
+        if out_chunk is None:
+            continue
+        out_type = type(out_chunk)
+        if isinstance(out_chunk, torch.Tensor):
+            out_chunk = {0: out_chunk}
+        elif isinstance(out_chunk, tuple) or isinstance(out_chunk, list):
+            chunk_length = len(out_chunk)
+            out_chunk = {i: chunk for i, chunk in enumerate(out_chunk)}
+        elif isinstance(out_chunk, dict):
+            pass
+        else:
+            print(
+                f"Return value of func must be in type [torch.Tensor, list, tuple, dict], get {type(out_chunk)}."
+            )
+            exit(1)
+        for k, v in out_chunk.items():
+            v = v if torch.is_grad_enabled() else v.detach()
+            out[k].append(v)
+
+    if out_type is None:
+        return None
+
+    out_merged: Dict[Any, Optional[torch.Tensor]] = {}
+    for k, v in out.items():
+        if all([vv is None for vv in v]):
+            # allow None in return value
+            out_merged[k] = None
+        elif all([isinstance(vv, torch.Tensor) for vv in v]):
+            out_merged[k] = torch.cat(v, dim=0)
+        else:
+            raise TypeError(
+                f"Unsupported types in return value of func: {[type(vv) for vv in v if not isinstance(vv, torch.Tensor)]}"
+            )
+
+    if out_type is torch.Tensor:
+        return out_merged[0]
+    elif out_type in [tuple, list]:
+        return out_type([out_merged[i] for i in range(chunk_length)])
+    elif out_type is dict:
+        return out_merged
+
+
+ValidScale = Union[Tuple[float, float], torch.FloatTensor]
+
+
+def scale_tensor(dat: torch.FloatTensor, inp_scale: ValidScale, tgt_scale: ValidScale):
+    if inp_scale is None:
+        inp_scale = (0, 1)
+    if tgt_scale is None:
+        tgt_scale = (0, 1)
+    if isinstance(tgt_scale, torch.FloatTensor):
+        assert dat.shape[-1] == tgt_scale.shape[-1]
+    dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0])
+    dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0]
+    return dat
+
+
+def get_activation(name) -> Callable:
+    if name is None:
+        return lambda x: x
+    name = name.lower()
+    if name == "none":
+        return lambda x: x
+    elif name == "exp":
+        return lambda x: torch.exp(x)
+    elif name == "sigmoid":
+        return lambda x: torch.sigmoid(x)
+    elif name == "tanh":
+        return lambda x: torch.tanh(x)
+    elif name == "softplus":
+        return lambda x: F.softplus(x)
+    else:
+        try:
+            return getattr(F, name)
+        except AttributeError:
+            raise ValueError(f"Unknown activation function: {name}")
+
+
+def get_ray_directions(
+    H: int,
+    W: int,
+    focal: Union[float, Tuple[float, float]],
+    principal: Optional[Tuple[float, float]] = None,
+    use_pixel_centers: bool = True,
+    normalize: bool = True,
+) -> torch.FloatTensor:
+    """
+    Get ray directions for all pixels in camera coordinate.
+    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
+               ray-tracing-generating-camera-rays/standard-coordinate-systems
+
+    Inputs:
+        H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
+    Outputs:
+        directions: (H, W, 3), the direction of the rays in camera coordinate
+    """
+    pixel_center = 0.5 if use_pixel_centers else 0
+
+    if isinstance(focal, float):
+        fx, fy = focal, focal
+        cx, cy = W / 2, H / 2
+    else:
+        fx, fy = focal
+        assert principal is not None
+        cx, cy = principal
+
+    i, j = torch.meshgrid(
+        torch.arange(W, dtype=torch.float32) + pixel_center,
+        torch.arange(H, dtype=torch.float32) + pixel_center,
+        indexing="xy",
+    )
+
+    directions = torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)
+
+    if normalize:
+        directions = F.normalize(directions, dim=-1)
+
+    return directions
+
+
+def get_rays(
+    directions,
+    c2w,
+    keepdim=False,
+    normalize=False,
+) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+    # Rotate ray directions from camera coordinate to the world coordinate
+    assert directions.shape[-1] == 3
+
+    if directions.ndim == 2:  # (N_rays, 3)
+        if c2w.ndim == 2:  # (4, 4)
+            c2w = c2w[None, :, :]
+        assert c2w.ndim == 3  # (N_rays, 4, 4) or (1, 4, 4)
+        rays_d = (directions[:, None, :] * c2w[:, :3, :3]).sum(-1)  # (N_rays, 3)
+        rays_o = c2w[:, :3, 3].expand(rays_d.shape)
+    elif directions.ndim == 3:  # (H, W, 3)
+        assert c2w.ndim in [2, 3]
+        if c2w.ndim == 2:  # (4, 4)
+            rays_d = (directions[:, :, None, :] * c2w[None, None, :3, :3]).sum(
+                -1
+            )  # (H, W, 3)
+            rays_o = c2w[None, None, :3, 3].expand(rays_d.shape)
+        elif c2w.ndim == 3:  # (B, 4, 4)
+            rays_d = (directions[None, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
+                -1
+            )  # (B, H, W, 3)
+            rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
+    elif directions.ndim == 4:  # (B, H, W, 3)
+        assert c2w.ndim == 3  # (B, 4, 4)
+        rays_d = (directions[:, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
+            -1
+        )  # (B, H, W, 3)
+        rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
+
+    if normalize:
+        rays_d = F.normalize(rays_d, dim=-1)
+    if not keepdim:
+        rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3)
+
+    return rays_o, rays_d
+
+
+def get_spherical_cameras(
+    n_views: int,
+    elevation_deg: float,
+    camera_distance: float,
+    fovy_deg: float,
+    height: int,
+    width: int,
+):
+    azimuth_deg = torch.linspace(0, 360.0, n_views + 1)[:n_views]
+    elevation_deg = torch.full_like(azimuth_deg, elevation_deg)
+    camera_distances = torch.full_like(elevation_deg, camera_distance)
+
+    elevation = elevation_deg * math.pi / 180
+    azimuth = azimuth_deg * math.pi / 180
+
+    # convert spherical coordinates to cartesian coordinates
+    # right hand coordinate system, x back, y right, z up
+    # elevation in (-90, 90), azimuth from +x to +y in (-180, 180)
+    camera_positions = torch.stack(
+        [
+            camera_distances * torch.cos(elevation) * torch.cos(azimuth),
+            camera_distances * torch.cos(elevation) * torch.sin(azimuth),
+            camera_distances * torch.sin(elevation),
+        ],
+        dim=-1,
+    )
+
+    # default scene center at origin
+    center = torch.zeros_like(camera_positions)
+    # default camera up direction as +z
+    up = torch.as_tensor([0, 0, 1], dtype=torch.float32)[None, :].repeat(n_views, 1)
+
+    fovy = torch.full_like(elevation_deg, fovy_deg) * math.pi / 180
+
+    lookat = F.normalize(center - camera_positions, dim=-1)
+    right = F.normalize(torch.cross(lookat, up), dim=-1)
+    up = F.normalize(torch.cross(right, lookat), dim=-1)
+    c2w3x4 = torch.cat(
+        [torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
+        dim=-1,
+    )
+    c2w = torch.cat([c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1)
+    c2w[:, 3, 3] = 1.0
+
+    # get directions by dividing directions_unit_focal by focal length
+    focal_length = 0.5 * height / torch.tan(0.5 * fovy)
+    directions_unit_focal = get_ray_directions(
+        H=height,
+        W=width,
+        focal=1.0,
+    )
+    directions = directions_unit_focal[None, :, :, :].repeat(n_views, 1, 1, 1)
+    directions[:, :, :, :2] = (
+        directions[:, :, :, :2] / focal_length[:, None, None, None]
+    )
+    # must use normalize=True to normalize directions here
+    rays_o, rays_d = get_rays(directions, c2w, keepdim=True, normalize=True)
+
+    return rays_o, rays_d
+
+
+# def remove_background(
+#     image: PIL.Image.Image,
+#     rembg_session: Any = None,
+#     force: bool = False,
+#     **rembg_kwargs,
+# ) -> PIL.Image.Image:
+#     do_remove = True
+#     if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
+#         do_remove = False
+#     do_remove = do_remove or force
+#     if do_remove:
+#         image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
+#     return image
+
+
+def resize_foreground(
+    image: PIL.Image.Image,
+    ratio: float,
+) -> PIL.Image.Image:
+    image = np.array(image)
+    assert image.shape[-1] == 4
+    alpha = np.where(image[..., 3] > 0)
+    y1, y2, x1, x2 = (
+        alpha[0].min(),
+        alpha[0].max(),
+        alpha[1].min(),
+        alpha[1].max(),
+    )
+    # crop the foreground
+    fg = image[y1:y2, x1:x2]
+    # pad to square
+    size = max(fg.shape[0], fg.shape[1])
+    ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
+    ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
+    new_image = np.pad(
+        fg,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+
+    # compute padding according to the ratio
+    new_size = int(new_image.shape[0] / ratio)
+    # pad to size, double side
+    ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
+    ph1, pw1 = new_size - size - ph0, new_size - size - pw0
+    new_image = np.pad(
+        new_image,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+    new_image = PIL.Image.fromarray(new_image)
+    return new_image
+
+
+def save_video(
+    frames: List[PIL.Image.Image],
+    output_path: str,
+    fps: int = 30,
+):
+    # use imageio to save video
+    frames = [np.array(frame) for frame in frames]
+    writer = imageio.get_writer(output_path, fps=fps)
+    for frame in frames:
+        writer.append_data(frame)
+    writer.close()
+
+
+def to_gradio_3d_orientation(mesh):
+    mesh.apply_transform(trimesh.transformations.rotation_matrix(-np.pi/2, [1, 0, 0]))
+    mesh.apply_scale([1, 1, -1])
+    mesh.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2, [0, 1, 0]))
+    return mesh

nodes/video_formats/av1-webm.json

@@ -0,0 +1,10 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libsvtav1",
+        "-pix_fmt", "yuv420p10le",
+        "-crf", "23"
+    ],
+     "extension": "webm",
+     "environment": {"SVT_LOG": "1"}
+}

nodes/video_formats/h264-mp4.json

@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libx264",
+        "-pix_fmt", "yuv420p",
+        "-crf", "19"
+    ],
+     "extension": "mp4"
+}

nodes/video_formats/h265-mp4.json

@@ -0,0 +1,11 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libx265",
+        "-pix_fmt", "yuv420p10le",
+        "-preset", "medium",
+        "-crf", "22",
+        "-x265-params", "log-level=quiet"
+    ],
+     "extension": "mp4"
+}

nodes/video_formats/webm.json

@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n",
+        "-pix_fmt", "yuv420p",
+        "-crf", "23"
+    ],
+     "extension": "webm"
+}

pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "comfyui-mixlab-nodes"
+description = "3D, ScreenShareNode & FloatingVideoNode, SpeechRecognition & SpeechSynthesis, GPT, LoadImagesFromLocal, Layers, Other Nodes, ..."
+version = "0.28.3"
+license = "MIT"
+dependencies = ["numpy", "pyOpenSSL", "watchdog", "opencv-python-headless", "matplotlib", "openai", "simple-lama-inpainting", "clip-interrogator==0.6.0", "transformers>=4.36.0", "lark-parser", "imageio-ffmpeg", "rembg[gpu]", "omegaconf==2.3.0", "Pillow>=9.5.0", "einops==0.7.0", "trimesh>=4.0.5", "huggingface-hub", "scikit-image"]
+
+[project.urls]
+Repository = "https://github.com/shadowcz007/comfyui-mixlab-nodes"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "shadow"
+DisplayName = "comfyui-mixlab-nodes"
+Icon = ""

requirements.txt

@@ -6,4 +6,13 @@ matplotlib
 openai
 simple-lama-inpainting
 clip-interrogator==0.6.0
-transformers>=4.36.0
+transformers>=4.36.0
+lark-parser
+imageio-ffmpeg
+rembg[gpu]
+omegaconf==2.3.0
+Pillow>=9.5.0
+einops==0.7.0
+trimesh>=4.0.5
+huggingface-hub
+scikit-image

web/javascript/app_mixlab.js

@@ -2,6 +2,46 @@ import { app } from '../../../scripts/app.js'
 import { $el } from '../../../scripts/ui.js'
 import { api } from '../../../scripts/api.js'
 
+//本机安装的插件节点全集
+window._nodesAll = null
+
+//获取当前系统的插件，节点清单
+function getObjectInfo () {
+  return new Promise(async (resolve, reject) => {
+    let url = getUrl()
+
+    try {
+      const response = await fetch(`${url}/object_info`)
+      const data = await response.json()
+      resolve(data)
+    } catch (error) {
+      reject(error)
+    }
+  })
+}
+
+const base64Df =
+  'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAwAAAAMCAYAAABWdVznAAAAAXNSR0IArs4c6QAAALZJREFUKFOFkLERwjAQBPdbgBkInECGaMLUQDsE0AkRVRAYWqAByxldPPOWHwnw4OBGye1p50UDSoA+W2ABLPN7i+C5dyC6R/uiAUXRQCs0bXoNIu4QPQzAxDKxHoALOrZcqtiyR/T6CXw7+3IGHhkYcy6BOR2izwT8LptG8rbMiCRAUb+CQ6WzQVb0SNOi5Z2/nX35DRyb/ENazhpWKoGwrpD6nICp5c2qogc4of+c7QcrhgF4Aa/aoAFHiL+RAAAAAElFTkSuQmCC'
+
+const parseImageToBase64 = url => {
+  return new Promise((res, rej) => {
+    fetch(url)
+      .then(response => response.blob())
+      .then(blob => {
+        const reader = new FileReader()
+        reader.onloadend = () => {
+          const base64data = reader.result
+          res(base64data)
+          // 在这里可以将base64数据用于进一步处理或显示图片
+        }
+        reader.readAsDataURL(blob)
+      })
+      .catch(error => {
+        console.log('发生错误:', error)
+      })
+  })
+}
+
 function get_position_style (ctx, widget_width, y, node_height) {
   const MARGIN = 12 // the margin around the html element
 
@@ -28,20 +68,21 @@ function get_position_style (ctx, widget_width, y, node_height) {
     // height: `${node_height * 0.3 - MARGIN * 2}px`,
     // background: '#EEEEEE',
     display: 'flex',
-    flexDirection: 'row',
+    flexDirection: 'column',
     // alignItems: 'center',
-    justifyContent: 'flex-start'
+    justifyContent: 'flex-start',
+    zIndex: 9999999
   }
 }
 
-async function drawImageToCanvas (imageUrl) {
+async function drawImageToCanvas (imageUrl, sFactor = 320) {
   var canvas = document.createElement('canvas')
   var ctx = canvas.getContext('2d')
   var img = new Image()
 
   await new Promise((resolve, reject) => {
     img.onload = function () {
-      var scaleFactor = 320 / img.width
+      var scaleFactor = sFactor / img.width
       var canvasWidth = img.width * scaleFactor
       var canvasHeight = img.height * scaleFactor
 
@@ -66,7 +107,11 @@ async function drawImageToCanvas (imageUrl) {
   // 可以在这里执行其他操作，比如将Base64数据保存到服务器或显示在页面上
 }
 
-function extractInputAndOutputData (jsonData, inputIds = [], outputIds = []) {
+async function extractInputAndOutputData (
+  jsonData,
+  inputIds = [],
+  outputIds = []
+) {
   // workflow
   // const workflow=jsonData.workflow;
   // const nodes=workflow.nodes;
@@ -75,6 +120,7 @@ function extractInputAndOutputData (jsonData, inputIds = [], outputIds = []) {
   let input = []
   let output = []
   const seed = {}
+  const seedTitle = {}
 
   for (const id in data) {
     if (data.hasOwnProperty(id)) {
@@ -114,16 +160,43 @@ function extractInputAndOutputData (jsonData, inputIds = [], outputIds = []) {
           }
         }
 
+        if (node.type == 'ImagesPrompt_') {
+          //图库
+          // console.log('ImagesPrompt_', data[id])
+          let image_base64 = data[id].inputs.image_base64
+          let img_index = 0
+          let imgsData = JSON.parse(data[id].inputs.upload)
+          for (let index = 0; index < imgsData.length; index++) {
+            const imgd = imgsData[index].imgurl
+            imgsData[index].index = index
+            //TODO缩放大小
+            imgsData[index].imgurl = await parseImageToBase64(imgd)
+            if (image_base64 == imgsData[index].imgurl) {
+              img_index = index
+            }
+          }
+          options.images = imgsData
+          delete data[id].inputs.upload
+          delete data[id].inputs.image_base64
+
+          data[id].inputs.imageIndex = img_index
+        }
+
         if (node.type == 'Color') {
         }
 
-        // loadImage的mask支持
         if (node.type === 'LoadImage') {
+          // loadImage的mask支持
           let output = node.outputs.filter(ot => ot.type == 'MASK')[0]
           if (output.links) {
             // 有输出
             options.hasMask = true
           }
+          // loadImage的默认图，转为base64
+          let imgurl = app.graph.getNodeById(id).imgs[0].src + '&channel=rgb'
+
+          options.defaultImage = await drawImageToCanvas(imgurl, 512)
+          console.log('#loadImage的默认图', options)
         }
 
         input[inputIds.indexOf(id)] = {
@@ -135,17 +208,41 @@ function extractInputAndOutputData (jsonData, inputIds = [], outputIds = []) {
         // input.push()
       }
       if (outputIds.includes(id)) {
+        let options = {}
+        //输出的默认图
+        if (
+          node.type === 'SaveImageAndMetadata_' &&
+          app.graph.getNodeById(id).imgs
+        ) {
+          // SaveImageAndMetadata_的默认图，转为base64
+          let imgurl = app.graph.getNodeById(id).imgs[0].src
+
+          options.defaultImage = await drawImageToCanvas(imgurl, 512)
+          console.log('#SaveImageAndMetadata_的默认图', options)
+        }
+
         // let node = app.graph.getNodeById(id)
         // output.push()
-        output[outputIds.indexOf(id)] = { ...data[id], title: node.title, id }
+        output[outputIds.indexOf(id)] = {
+          ...data[id],
+          title: node.title,
+          id,
+          options
+        }
       }
 
-      if (node.type === 'KSampler' || node.type == 'SamplerCustom') {
+      if (
+        node.type === 'KSampler' ||
+        node.type == 'SamplerCustom' ||
+        node.type === 'ChinesePrompt_Mix' ||
+        node.type === 'Seed_'
+      ) {
         // seed 的类型收集
         try {
           seed[id] = node.widgets.filter(
             w => w.name === 'seed' || w.name == 'noise_seed'
           )[0].linkedWidgets[0].value
+          seedTitle[id] = node.title
         } catch (error) {}
       }
     }
@@ -155,7 +252,7 @@ function extractInputAndOutputData (jsonData, inputIds = [], outputIds = []) {
   input = input.filter(i => i)
   output = output.filter(i => i)
 
-  return { input, output, seed }
+  return { input, output, seed, seedTitle }
 }
 
 function getUrl () {
@@ -207,7 +304,10 @@ function downloadJsonFile (jsonData, fileName = 'mix_app.json') {
 }
 
 async function save (json, download = false, showInfo = true) {
-  console.log('####SAVE', json[0])
+  let nodesAll = window._nodesAll || (await getObjectInfo())
+
+  console.log('####SAVE', nodesAll, json[0])
+
   const name = json[0],
     version = json[5],
     share_prefix = json[6], //用于分享的功能扩展
@@ -228,12 +328,26 @@ async function save (json, download = false, showInfo = true) {
   try {
     let data = await app.graphToPrompt()
 
-    let { input, output, seed } = extractInputAndOutputData(
+    //从output数据里把工作流的节点，插件数据统计出来
+    data.nodesMap = {}
+    for (const id in data.output) {
+      data.nodesMap[data.output[id].class_type] =
+        nodesAll[data.output[id].class_type]
+    }
+
+    let { input, output, seed, seedTitle } = await extractInputAndOutputData(
       data,
       inputIds,
       outputIds
     )
 
+    let authorAvatar =
+        localStorage.getItem('_mixlab_author_avatar') || base64Df,
+      authorName =
+        localStorage.getItem('_mixlab_author_name') ||
+        localStorage.getItem('Comfy.userName'),
+      authorLink = localStorage.getItem('_mixlab_author_link') || ''
+
     data.app = {
       name,
       description,
@@ -241,10 +355,16 @@ async function save (json, download = false, showInfo = true) {
       input,
       output,
       seed, //控制是fixed 还是random
+      seedTitle,
       share_prefix,
       link,
       category,
-      filename: `${name}_${version}.json`
+      filename: `${name}_${version}.json`,
+      author: {
+        avatar: authorAvatar,
+        name: authorName,
+        link: authorLink
+      }
     }
 
     try {
@@ -277,12 +397,13 @@ async function save (json, download = false, showInfo = true) {
 
 function getInputsAndOutputs () {
   const inputs =
-      `LoadImage VHS_LoadVideo CLIPTextEncode PromptSlide TextInput_ Color FloatSlider IntNumber CheckpointLoaderSimple LoraLoader`.split(
+      `LoadImage LoadImagesToBatch ImagesPrompt_ VHS_LoadVideo CLIPTextEncode PromptSlide TextInput_ Color FloatSlider IntNumber CheckpointLoaderSimple LoraLoader`.split(
         ' '
       ),
-    outputs = `PreviewImage SaveImage ShowTextForGPT VHS_VideoCombine`.split(
-      ' '
-    )
+    outputs =
+      `SaveTripoSRMesh,PreviewImage,SaveImage,TransparentImage,ShowTextForGPT,VHS_VideoCombine,VideoCombine_Adv,Image Save,SaveImageAndMetadata_,ClipInterrogator`.split(
+        ','
+      )
 
   let inputsId = [],
     outputsId = []
@@ -305,6 +426,11 @@ function getInputsAndOutputs () {
 
 app.registerExtension({
   name: 'Mixlab.utils.AppInfo',
+  init () {
+    if (!window._nodesAll) {
+      getObjectInfo().then(r => (window._nodesAll = r))
+    }
+  },
   async beforeRegisterNodeDef (nodeType, nodeData, app) {
     if (nodeType.comfyClass == 'AppInfo') {
       const orig_nodeCreated = nodeType.prototype.onNodeCreated
@@ -377,10 +503,164 @@ app.registerExtension({
           }
         })
 
-        document.body.appendChild(widget.div)
-        widget.div.appendChild(btn)
-        widget.div.appendChild(download)
+        // author
+        let author = document.createElement('div')
+        // author.style=`display: flex`
 
+        let authorAvatar = document.createElement('img')
+        authorAvatar.className = `${'comfy-multiline-input'}`
+        authorAvatar.style = `outline: none;
+          border: none;
+          padding: 4px;
+          width: 32px;
+          cursor: pointer;
+          height: 32px;`
+
+        if (localStorage.getItem('_mixlab_author_avatar')) {
+          authorAvatar.src =
+            localStorage.getItem('_mixlab_author_avatar') || base64Df
+        }
+
+        let authorAvatarUpload = document.createElement('input')
+        authorAvatarUpload.type = 'file'
+        authorAvatarUpload.style = `display:none`
+
+        let authorAvatarInput = document.createElement('div')
+        authorAvatarInput.style = `display: flex;justify-content: flex-start;
+        align-items: center;`
+        let authorAvatarInputLabel = document.createElement('p')
+        authorAvatarInputLabel.innerText = 'Author Avatar'
+        authorAvatarInputLabel.className = `${'comfy-multiline-input'}`
+        authorAvatarInputLabel.style = `font-size:12px`
+
+        authorAvatar.addEventListener('click', e => {
+          authorAvatarUpload.click()
+        })
+
+        authorAvatarInputLabel.addEventListener('click', e => {
+          authorAvatarUpload.click()
+        })
+
+        authorAvatarUpload.addEventListener('change', event => {
+          const file = event.target.files[0]
+          const reader = new FileReader()
+
+          reader.onload = async e => {
+            let im = new Image()
+            im.src = e.target.result
+            authorAvatar.src = e.target.result
+            im.onload = () => {
+              let c = document.createElement('canvas')
+              let ctx = c.getContext('2d')
+              c.width = 72
+              c.height = 72
+              ctx.drawImage(
+                im,
+                0,
+                0,
+                im.naturalWidth,
+                im.naturalHeight,
+                0,
+                0,
+                c.width,
+                c.height
+              )
+              window._mixlab_author_avatar = c.toDataURL()
+              localStorage.setItem(
+                '_mixlab_author_avatar',
+                window._mixlab_author_avatar
+              )
+            }
+          }
+
+          // 以文本形式读取文件
+          reader.readAsDataURL(file)
+        })
+
+        author.appendChild(authorAvatarInput)
+        authorAvatarInput.appendChild(authorAvatarInputLabel)
+        authorAvatarInput.appendChild(authorAvatar)
+        authorAvatarInput.appendChild(authorAvatarUpload)
+
+        let authorName = document.createElement('input')
+        authorName.type = 'text'
+        authorName.value =
+          localStorage.getItem('_mixlab_author_name') ||
+          localStorage.getItem('Comfy.userName')
+        authorName.placeholder = 'author name'
+        authorName.className = `${'comfy-multiline-input'}`
+        authorName.style = `
+          outline: none;
+          border: none;
+          padding: 4px;
+          width: 100%;
+          cursor: pointer;
+          height: 32px;`
+
+        let authorNameInput = document.createElement('div')
+        authorNameInput.style = `display: flex;justify-content: flex-start;
+          align-items: center;`
+        let authorNameInputLabel = document.createElement('p')
+        authorNameInputLabel.innerText = 'Author Name'
+        authorNameInputLabel.className = `${'comfy-multiline-input'}`
+        authorNameInputLabel.style = `font-size:12px;width: 110px`
+
+        authorName.addEventListener('change', e => {
+          window._mixlab_author_name = authorName.value.trim()
+          localStorage.setItem(
+            '_mixlab_author_name',
+            window._mixlab_author_name
+          )
+        })
+
+        author.appendChild(authorNameInput)
+        authorNameInput.appendChild(authorNameInputLabel)
+        authorNameInput.appendChild(authorName)
+
+        // 社交链接
+        let authorLink = document.createElement('input')
+        authorLink.type = 'text'
+        authorLink.value = localStorage.getItem('_mixlab_author_link') || ''
+        authorLink.placeholder = 'author link'
+        authorLink.className = `${'comfy-multiline-input'}`
+        authorLink.style = `
+          outline: none;
+          border: none;
+          padding: 4px;
+          width: 100%;
+          cursor: pointer;
+          height: 32px;`
+
+        let authorLinkInput = document.createElement('div')
+        authorLinkInput.style = `display: flex;justify-content: flex-start;
+          align-items: center;`
+        let authorLinkInputLabel = document.createElement('p')
+        authorLinkInputLabel.innerText = 'Author Link'
+        authorLinkInputLabel.className = `${'comfy-multiline-input'}`
+        authorLinkInputLabel.style = `font-size:12px;width: 110px`
+
+        authorLink.addEventListener('change', e => {
+          window._mixlab_author_link = authorLink.value.trim()
+          localStorage.setItem(
+            '_mixlab_author_link',
+            window._mixlab_author_link
+          )
+        })
+
+        author.appendChild(authorLinkInput)
+        authorLinkInput.appendChild(authorLinkInputLabel)
+        authorLinkInput.appendChild(authorLink)
+
+        widget.div.appendChild(author)
+
+        let btns = document.createElement('div')
+
+        widget.div.appendChild(btns)
+
+        btns.appendChild(btn)
+        btns.appendChild(download)
+
+        document.body.appendChild(widget.div)
         this.addCustomWidget(widget)
 
         const onRemoved = this.onRemoved

web/javascript/chat.js

@@ -0,0 +1,106 @@
+async function* completion (url, messages, controller) {
+  let data = {
+    model: 'gpt-3.5-turbo-16k',
+    messages,
+    temperature: 0.05,
+    stream: true
+  }
+  // if (imageNode) {
+  //   data = { ...data, image_data: [imageNode] }
+  // }
+
+  // let controller = new AbortController()
+
+  let response = await fetch(url, {
+    method: 'POST',
+    body: JSON.stringify(data),
+    headers: {
+      Connection: 'keep-alive',
+      'Content-Type': 'application/json',
+      Accept: 'text/event-stream'
+    },
+    signal: controller.signal
+  })
+
+  const reader = response.body.getReader()
+  const decoder = new TextDecoder()
+
+  let content = ''
+  let leftover = '' // Buffer for partially read lines
+
+  try {
+    let cont = true
+    while (cont) {
+      let result = await reader.read()
+      if (result.done) {
+        break
+      }
+
+      // Add any leftover data to the current chunk of data
+      const text = leftover + decoder.decode(result.value)
+
+      // Check if the last character is a line break
+      const endsWithLineBreak = text.endsWith('\n')
+
+      // Split the text into lines
+      let lines = text.split('\n')
+
+      // If the text doesn't end with a line break, then the last line is incomplete
+      // Store it in leftover to be added to the next chunk of data
+      if (!endsWithLineBreak) {
+        leftover = lines.pop()
+      } else {
+        leftover = '' // Reset leftover if we have a line break at the end
+      }
+
+      // Parse all sse events and add them to result
+      const regex = /^(\S+):\s(.*)$/gm
+      for (const line of lines) {
+        const match = regex.exec(line)
+        if (match) {
+          result[match[1]] = match[2]
+          // since we know this is llama.cpp, let's just decode the json in data
+          if (result.data) {
+            result.data = JSON.parse(result.data)
+            // console.log('#result.data',result.data)
+
+            content += result.data.choices[0].delta?.content || ''
+
+            // yield
+            yield result
+
+            // if we got a stop token from server, we will break here
+            if (result.data.choices[0].finish_reason == 'stop') {
+              if (result.data.generation_settings) {
+                // generation_settings = result.data.generation_settings;
+              }
+              cont = false
+              break
+            }
+          }
+        }
+      }
+    }
+  } catch (e) {
+    console.error('llama error: ', e)
+    throw e
+  } finally {
+    controller.abort()
+  }
+
+  return content
+  // return (await response.json()).content
+}
+
+export async function completion_ (url, messages, controller, callback) {
+  let request = await completion(url, messages, controller)
+  for await (const chunk of request) {
+    let content = chunk.data.choices[0].delta.content || ''
+    if (chunk.data.choices[0].role == 'assistant') {
+      //开始
+      content = ''
+    }
+
+    if (callback) callback(content)
+  }
+}

web/javascript/checkVersion_mixlab.js

@@ -3,7 +3,7 @@ import { app } from '../../../scripts/app.js'
 const repoOwner = 'shadowcz007' // 替换为仓库的所有者
 const repoName = 'comfyui-mixlab-nodes' // 替换为仓库的名称
 
-const version = 'v0.17.0'
+const version = 'v0.28.3'
 
 fetch(`https://api.github.com/repos/${repoOwner}/${repoName}/releases/latest`)
   .then(response => response.json())
@@ -17,7 +17,13 @@ fetch(`https://api.github.com/repos/${repoOwner}/${repoName}/releases/latest`)
       return
     if (latestVersion && latestVersion != version) {
       localStorage.setItem('_mixlab_nodes_vesion', latestVersion)
-      app.ui.dialog.show(`<h4 style="font-size: 18px;">${repoName} <br>
+      app.ui.dialog.show(`<a style="color: white;
+      font-size: 18px;
+      font-weight: 800;
+      letter-spacing: 2px;
+    }"
+    href="https://discord.gg/cXs9vZSqeK">Welcome to Mixlab nodes discord</a>
+    <h4 style="font-size: 18px;">${repoName} <br>
       Latest release version: ${latestVersion}</h4>
       <p>Please proceed to the official repository to download the latest version.</p>
       <a style="color: #2196F3;

web/javascript/gpt_mixlab.js

@@ -61,7 +61,7 @@ app.registerExtension({
   async getCustomWidgets (app) {
     return {
       KEY (node, inputName, inputData, app) {
-        console.log('##inputData', inputData)
+        // console.log('##inputData', inputData)
         const widget = {
           type: inputData[0], // the type, CHEESE
           name: inputName, // the name, slice

web/javascript/image_mixlab.js

@@ -1,7 +1,39 @@
 import { app } from '../../../scripts/app.js'
 import { api } from '../../../scripts/api.js'
-import { ComfyWidgets } from '../../../scripts/widgets.js'
+// import { ComfyWidgets } from '../../../scripts/widgets.js'
 import { $el } from '../../../scripts/ui.js'
+import { applyTextReplacements } from '../../../scripts/utils.js'
+
+function loadImageToCanvas (base64Image) {
+  var img = new Image()
+  var canvas = document.createElement('canvas')
+  var ctx = canvas.getContext('2d')
+  return new Promise((res, rej) => {
+    img.onload = function () {
+      // 等比例缩放图片
+      var width = img.width
+      var height = img.height
+      var max_width = 1024
+      if (width > max_width) {
+        height *= max_width / width
+        width = max_width
+      }
+
+      // 设置canvas尺寸
+      canvas.width = width
+      canvas.height = height
+
+      // 在canvas上绘制图片
+      ctx.drawImage(img, 0, 0, width, height)
+
+      // 将canvas转换为base64图片数据
+      var canvasData = canvas.toDataURL()
+      res(canvasData) //  canvas转换后的base64图片数据
+    }
+
+    img.src = base64Image
+  })
+}
 
 async function uploadImage (blob, fileType = '.svg', filename) {
   // const blob = await (await fetch(src)).blob();
@@ -28,6 +60,9 @@ async function uploadImage (blob, fileType = '.svg', filename) {
   return src
 }
 
+const base64Df =
+  'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAwAAAAMCAYAAABWdVznAAAAAXNSR0IArs4c6QAAALZJREFUKFOFkLERwjAQBPdbgBkInECGaMLUQDsE0AkRVRAYWqAByxldPPOWHwnw4OBGye1p50UDSoA+W2ABLPN7i+C5dyC6R/uiAUXRQCs0bXoNIu4QPQzAxDKxHoALOrZcqtiyR/T6CXw7+3IGHhkYcy6BOR2izwT8LptG8rbMiCRAUb+CQ6WzQVb0SNOi5Z2/nX35DRyb/ENazhpWKoGwrpD6nICp5c2qogc4of+c7QcrhgF4Aa/aoAFHiL+RAAAAAElFTkSuQmCC'
+
 function base64ToBlobFromURL (base64URL, contentType) {
   return fetch(base64URL).then(response => response.blob())
 }
@@ -108,7 +143,7 @@ function createImage (url) {
   })
 }
 
-const parseImage = url => {
+const parseImageToBase64 = url => {
   return new Promise((res, rej) => {
     fetch(url)
       .then(response => response.blob())
@@ -406,9 +441,7 @@ app.registerExtension({
 
         this.serialize_widgets = true //需要保存参数
       }
-    };
-
- 
+    }
   },
   async loadedGraphNode (node, app) {
     // Fires every time a node is constructed
@@ -442,3 +475,518 @@ app.registerExtension({
     }
   }
 })
+
+const createSelect = (imgDiv, select, opts, targetWidget, textWidget) => {
+  select.style.display = 'block'
+  let html = ''
+  let isMatch = false
+  for (const opt of opts) {
+    html += `<option value='${opt.keyword}' ${opt.selected ? 'selected' : ''}>${
+      opt.keyword
+    }</option>`
+    if (opt.selected) {
+      isMatch = true
+      imgDiv.src = opt.imgurl
+      // targetWidget.value = opt.keyword
+    }
+  }
+  select.innerHTML = html
+  if (!isMatch) {
+    // targetWidget.value = opts[0].keyword
+    imgDiv.src = opts[0].imgurl
+  }
+
+  // 添加change事件监听器
+  select.addEventListener('change', async function () {
+    // 获取选中的选项的值
+    var selectedOption = select.options[select.selectedIndex].value
+    let t = opts.filter(opt => opt.keyword === selectedOption)[0]
+
+    targetWidget.value = await parseImageToBase64(t.imgurl)
+    imgDiv.src = targetWidget.value
+    textWidget.value = t.keyword
+  })
+  // console.log(select)
+}
+
+app.registerExtension({
+  name: 'Mixlab.prompt.ImagesPrompt_',
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeType.comfyClass == 'ImagesPrompt_') {
+      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+      nodeType.prototype.onNodeCreated = async function () {
+        orig_nodeCreated?.apply(this, arguments)
+
+        const image_prompt = this.widgets.filter(
+          w => w.name == 'image_base64'
+        )[0]
+        const image_text = this.widgets.filter(w => w.name == 'text')[0]
+
+        const node = this
+
+        const widget = {
+          type: 'div',
+          name: 'upload',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, widget_width, y, node.size[1])
+            )
+          }
+        }
+
+        widget.div = $el('div', {})
+
+        // console.log('image_prompt',image_prompt)
+        const img = new Image()
+        img.src = image_prompt?.value || base64Df
+        widget.div.appendChild(img)
+
+        const btn = document.createElement('button')
+        btn.innerText = 'Upload Images JSON'
+
+        btn.style = `cursor: pointer;
+        font-weight: 300;
+        margin: 2px; 
+        color: var(--descrip-text);
+        background-color: var(--comfy-input-bg);
+        border-radius: 8px;
+        border-color: var(--border-color);
+        border-style: solid;height: 30px;min-width: 122px;
+       `
+
+        const select = document.createElement('select')
+        select.style = `display:none;cursor: pointer;
+        font-weight: 300;
+        margin: 2px; 
+        color: var(--descrip-text);
+        background-color: var(--comfy-input-bg);
+        border-radius: 8px;
+        border-color: var(--border-color);
+        border-style: solid;height: 30px;min-width: 100px;
+       `
+        widget.select = select
+
+        // const btn=document.createElement('button');
+        // btn.innerText='Upload'
+        btn.addEventListener('click', () => {
+          let inp = document.createElement('input')
+          inp.type = 'file'
+          inp.accept = '.json'
+          inp.click()
+          inp.addEventListener('change', event => {
+            // 获取选择的文件
+            // [{title,imageUrl}]
+            const file = event.target.files[0]
+            this.title = file.name.split('.')[0]
+
+            // console.log(file.name.split('.')[0])
+            // 创建文件读取器
+            const reader = new FileReader()
+
+            // 定义读取完成事件的回调函数
+            reader.onload = async event => {
+              // 读取完成后的文本内容
+              const json = JSON.parse(event.target.result)
+              console.log(node, json)
+
+              widget.value = JSON.stringify(json)
+
+              let img = widget.div.querySelector('img')
+
+              createSelect(img, select, json, image_prompt, image_text)
+
+              image_prompt.value = await parseImageToBase64(json[0].imgurl)
+              image_text.value = json[0].keyword
+
+              if (img) {
+                img.src = image_prompt.value
+              }
+
+              inp.remove()
+            }
+
+            // 以文本方式读取文件
+            reader.readAsText(file)
+          })
+        })
+
+        widget.div.appendChild(btn)
+        widget.div.appendChild(select)
+        document.body.appendChild(widget.div)
+        this.addCustomWidget(widget)
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          widget.div.remove()
+          return onRemoved?.()
+        }
+
+        if (this.onResize) {
+          this.onResize(this.size)
+        }
+
+        this.serialize_widgets = true //需要保存参数
+      }
+    }
+  },
+  async loadedGraphNode (node, app) {
+    if (node.type === 'ImagesPrompt_') {
+      try {
+        let prompt = node.widgets.filter(w => w.name === 'image_base64')[0]
+        let text = node.widgets.filter(w => w.name === 'text')[0]
+        let uploadWidget = node.widgets.filter(w => w.name == 'upload')[0]
+        // console.log('##prompt',prompt.value)
+        let img = uploadWidget.div.querySelector('img')
+        let json = JSON.parse(uploadWidget.value)
+
+        for (let index = 0; index < json.length; index++) {
+          const j = json[index]
+          let base64 = await parseImageToBase64(j.imgurl)
+          if (base64 === prompt.value) {
+            json[index].selected = true
+          }
+        }
+
+        if (json && json[0]) {
+          uploadWidget.select.style.display = 'block'
+          createSelect(img, uploadWidget.select, json, prompt, text)
+        }
+      } catch (error) {}
+    }
+  }
+})
+
+const createInputImageForBatch = (base64, widget) => {
+  let im = new Image()
+  im.src = base64
+  im.style = `width: 88px;`
+
+  im.addEventListener('click', e => {
+    let newValue = []
+    let items = widget.value?.base64 || []
+    for (const v of items) {
+      if (v != base64) newValue.push(v)
+    }
+    widget.value.base64 = newValue
+    im.remove()
+  })
+
+  return im
+}
+
+app.registerExtension({
+  name: 'Mixlab.Comfy.LoadImagesToBatch',
+  async getCustomWidgets (app) {
+    return {
+      IMAGEBASE64 (node, inputName, inputData, app) {
+        // console.log('##node', node)
+        const widget = {
+          value: {
+            base64: []
+          }, // 不能[x,x,x]
+          type: inputData[0], // the type
+          name: inputName, // the name, slice
+          size: [128, 32], // a default size
+          draw (ctx, node, width, y) {},
+          computeSize (...args) {
+            return [128, 32] // a method to compute the current size of the widget
+          }
+          // serializeValue (nodeId, widgetIndex) {
+          //   return widget.value
+          // },
+        }
+        //  widget.something = something;          // maybe adds stuff to it
+        node.addCustomWidget(widget) // adds it to the node
+        return widget // and returns it.
+      }
+    }
+  },
+
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeType.comfyClass == 'LoadImagesToBatch') {
+
+      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+
+      nodeType.prototype.onNodeCreated = function () {
+        orig_nodeCreated?.apply(this, arguments)
+
+        let imagesWidget = this.widgets.filter(w => w.name == 'images')[0]
+
+        const widget = {
+          type: 'div',
+          name: 'image_base64',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, widget_width, 44, node.size[1])
+            )
+          },
+          serialize: false
+        }
+
+        widget.div = $el('div', {})
+
+        document.body.appendChild(widget.div)
+
+        let imagePreview = document.createElement('div')
+        let imagesDiv = document.createElement('div') //显示图片
+        imagesDiv.className = 'images_preview'
+        imagesDiv.style = `width: calc(100% - 14px);
+        display: flex;
+        flex-wrap: wrap;
+        padding: 7px;    justify-content: space-between;
+        align-items: center;`
+
+        let inputImage = document.createElement('input')
+        inputImage.type = 'file'
+        inputImage.style.display = 'none'
+        inputImage.addEventListener('change', e => {
+          e.preventDefault()
+          const file = e.target.files[0]
+          const reader = new FileReader()
+          reader.onload = async event => {
+            let base64 = event.target.result
+            //压缩图片，控制1024以内
+            base64 = await loadImageToCanvas(base64)
+            // console.log(base64)
+            if (!imagesWidget.value) imagesWidget.value = { base64: [] }
+            imagesWidget.value.base64.push(base64)
+            let im = createInputImageForBatch(base64, imagesWidget)
+            imagesDiv.appendChild(im)
+          }
+          reader.readAsDataURL(file)
+        })
+
+        const btn = document.createElement('button')
+        btn.innerText = 'Upload Image'
+
+        btn.style = `cursor: pointer;
+        font-weight: 300;
+        margin: 2px; 
+        color: var(--descrip-text);
+        background-color: var(--comfy-input-bg);
+        border-radius: 8px;
+        border-color: var(--border-color);
+        border-style: solid;height: 30px;min-width: 122px;
+       `
+
+        btn.addEventListener('click', e => {
+          e.preventDefault()
+          inputImage.click()
+        })
+
+        widget.div.appendChild(imagePreview)
+        imagePreview.appendChild(imagesDiv)
+        imagePreview.appendChild(btn)
+        imagePreview.appendChild(inputImage)
+
+        this.addCustomWidget(widget)
+
+        // document.addEventListener('wheel', handleMouseWheel)
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          inputImage.remove()
+          widget.div.remove()
+          try {
+            // document.removeEventListener('wheel', handleMouseWheel)
+          } catch (error) {
+            console.log(error)
+          }
+
+          return onRemoved?.()
+        }
+
+        this.serialize_widgets = true //需要保存参数
+      }
+    }
+
+    if (nodeData.name === 'SaveImageAndMetadata_') {
+      const onNodeCreated = nodeType.prototype.onNodeCreated
+      // /web/extensions/core/saveImageExtraOutput.js
+      nodeType.prototype.onNodeCreated = function () {
+        const r = onNodeCreated
+          ? onNodeCreated.apply(this, arguments)
+          : undefined
+        const widget = this.widgets.find(w => w.name === 'filename_prefix')
+        widget.serializeValue = () => {
+          return applyTextReplacements(app, widget.value)
+        }
+
+        return r
+      }
+
+      const onExecuted = nodeType.prototype.onExecuted
+      nodeType.prototype.onExecuted = function (message) {
+        onExecuted?.apply(this, arguments)
+        console.log('##onExecuted', this, message)
+        //TODO 是否 保存base64
+        if (message.base64) {
+          if (Array.isArray(message.base64)) {
+          }
+        }
+      }
+    }
+  },
+  async loadedGraphNode (node, app) {
+    if (node.type === 'LoadImagesToBatch') {
+      // await sleep(0)
+      let imagesWidget = node.widgets.filter(w => w.name === 'images')[0]
+      let imagePreview = node.widgets.filter(w => w.name == 'image_base64')[0]
+
+      let pre = imagePreview.div.querySelector('.images_preview')
+      for (const d of imagesWidget.value?.base64 || []) {
+        let im = createInputImageForBatch(d, imagesWidget)
+        pre.appendChild(im)
+      }
+    }
+  }
+})
+
+// 如何引入css
+app.registerExtension({
+  name: 'Mixlab.output.ComparingTwoFrames_',
+  init () {
+    $el('link', {
+      rel: 'stylesheet',
+      href: '/extensions/comfyui-mixlab-nodes/lib/juxtapose.css',
+      parent: document.head
+    })
+
+    $el('style', {
+      textContent: `
+      .juxtapose-name{
+        display: none!important; 
+      }
+			`,
+      parent: document.body
+    })
+  },
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeType.comfyClass == 'ComparingTwoFrames_') {
+      const onNodeCreated = nodeType.prototype.onNodeCreated
+      nodeType.prototype.onNodeCreated = function () {
+        const r = onNodeCreated
+        ? onNodeCreated.apply(this, arguments)
+        : undefined
+
+        this.size = [400, this.size[1]]
+        console.log('##onNodeCreated', this)
+        const widget = {
+          type: 'div',
+          name: 'preview',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, 400, 44, node.size[1])
+            )
+          },
+          serialize: false
+        }
+
+        widget.div = $el('div', {})
+
+        document.body.appendChild(widget.div)
+        this.addCustomWidget(widget)
+        this.serialize_widgets = true //需要保存参数
+
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          widget.div.remove()
+          return onRemoved?.()
+        }
+
+
+       return r
+
+      }
+
+     
+
+      const onExecuted = nodeType.prototype.onExecuted
+      nodeType.prototype.onExecuted = function (message) {
+        onExecuted?.apply(this, arguments)
+        console.log('##onExecuted', this, message)
+
+        this.widgets[0].div.id = 'mix_comparingtowframes_' + this.id
+
+        let after_image = message.after_images[0]
+        let before_image = message.before_images[0]
+
+        after_image = `${window.location.protocol}//${
+          window.location.hostname
+        }:${window.location.port}/view?filename=${encodeURIComponent(
+          after_image.filename
+        )}&type=${after_image.type}&subfolder=${encodeURIComponent(
+          after_image.subfolder
+        )}&t=${+new Date()}`
+
+        before_image = `${window.location.protocol}//${
+          window.location.hostname
+        }:${window.location.port}/view?filename=${encodeURIComponent(
+          before_image.filename
+        )}&type=${before_image.type}&subfolder=${encodeURIComponent(
+          before_image.subfolder
+        )}&t=${+new Date()}`
+
+        this.widgets[0].div.innerHTML = ''
+
+        let slider = new juxtapose.JXSlider(
+          '#mix_comparingtowframes_' + this.id,
+          [
+            {
+              src: before_image,
+              label: 'Before'
+            },
+            {
+              src: after_image,
+              label: 'After'
+            }
+          ],
+          {
+            animate: true,
+            showLabels: true,
+            showCredits: false,
+            startingPosition: '50%',
+            makeResponsive: false
+          }
+        )
+
+        this.widgets_values = [
+          {
+            src: before_image,
+            label: 'Before'
+          },
+          {
+            src: after_image,
+            label: 'After'
+          }
+        ]
+        this.size=[this.size[0],300]
+      }
+    }
+  },
+  async loadedGraphNode (node, app) {
+    // console.log('##loadedGraphNode', node)
+    if (node.type === 'ComparingTwoFrames_') {
+      // node.widgets[0].div.id = 'mix_comparingtowframes_' + node.id
+      // if (node.widgets_values && node.widgets_values[0]) {
+      //   node.widgets[0].div.innerHTML = ''
+
+      //   let slider = new juxtapose.JXSlider(
+      //     '#mix_comparingtowframes_' + node.id,
+      //     node.widgets_values,
+      //     {
+      //       animate: true,
+      //       showLabels: true,
+      //       showCredits: false,
+      //       startingPosition: '50%',
+      //       makeResponsive: false
+      //     }
+      //   )
+      // }
+    }
+  }
+})

web/javascript/layer_mixlab.js

@@ -1,8 +1,70 @@
 import { app } from '../../../scripts/app.js'
-// import { api } from '../../../scripts/api.js'
+import { api } from '../../../scripts/api.js'
 import { ComfyWidgets } from '../../../scripts/widgets.js'
 import { $el } from '../../../scripts/ui.js'
 
+function downloadJsonFile (jsonData, fileName = 'grid.json') {
+  const dataString = JSON.stringify(jsonData)
+  const blob = new Blob([dataString], { type: 'application/json' })
+  const url = URL.createObjectURL(blob)
+
+  const link = document.createElement('a')
+  link.href = url
+  link.download = fileName
+  link.click()
+
+  // 释放URL对象
+  setTimeout(() => {
+    URL.revokeObjectURL(url)
+  }, 0)
+}
+
+function createSelectWithOptions (options) {
+  const select = document.createElement('select')
+
+  options.forEach(option => {
+    const optionElement = document.createElement('option')
+    optionElement.text = option
+    optionElement.value = option
+    select.appendChild(optionElement)
+  })
+
+  select.style = `cursor: pointer;
+  font-weight: 300; 
+  height: 30px;
+  min-width: 122px;
+  position: absolute;
+  top: 24px;
+  left: 88px;
+  z-index: 999999999999999;
+       `
+
+  return select
+}
+
+function drawCanvasWithText (w, h, tag, color = 'rgba(255,255,255,0.4)') {
+  const canvas = document.createElement('canvas')
+  const ctx = canvas.getContext('2d')
+
+  // 设置画布大小
+  canvas.width = w
+  canvas.height = h
+
+  // 绘制白色背景
+  ctx.fillStyle = color
+  ctx.fillRect(0, 0, canvas.width, canvas.height)
+
+  // 绘制文字
+  ctx.fillStyle = '#000000'
+  ctx.font = '20px Arial'
+  ctx.fillText(tag, 50, 50)
+
+  // 导出为Base64
+  const base64 = canvas.toDataURL()
+
+  return base64
+}
+
 function get_position_style (ctx, widget_width, y, node_height) {
   const MARGIN = 4 // the margin around the html element
 
@@ -156,32 +218,29 @@ const parseSvg = async svgContent => {
   return { data, image: base64, svgElement }
 }
 
-
-function findImages(nodeId) {
+function findImages (nodeId) {
   // 检查当前节点是否有 imgs 字段
   const n = app.graph.getNodeById(nodeId)
   if (n.imgs) {
-    return n.imgs;
+    return n.imgs
   }
 
   // 检查当前节点的 inputs 是否有 image 字段
   if (n.inputs) {
     for (let i = 0; i < n.inputs.length; i++) {
-      if (n.inputs[i].name==='image'||n.inputs[i].name==='images') {
+      if (n.inputs[i].name === 'image' || n.inputs[i].name === 'images') {
         // 获取新的 nodeId，并递归调用 findImages 函数
-        var linkId = n.inputs[i]?.link;
+        var linkId = n.inputs[i]?.link
         var origin_id = app.graph.links[linkId].origin_id
-        return findImages(origin_id);
+        return findImages(origin_id)
       }
     }
   }
 
   // 如果没有找到 imgs 字段或者 image 字段，则返回 null
-  return null;
+  return null
 }
 
- 
-
 async function setArea (cw, ch, topBase64, base64, data, fn) {
   let displayHeight = Math.round(window.screen.availHeight * 0.8)
   let div = document.createElement('div')
@@ -353,6 +412,196 @@ async function setArea (cw, ch, topBase64, base64, data, fn) {
   }
 }
 
+async function setAreaTags (cw, ch, grids, fn) {
+  let base64 = drawCanvasWithText(cw, ch, '', 'white')
+  let displayHeight = Math.round(window.screen.availHeight * 0.8)
+  let div = document.createElement('div')
+  div.innerHTML = `
+    <div id='ml_overlay' style='position: absolute;top:0;background: #251f1fc4;
+    height: 100vh;
+    z-index:999999;
+    width: 100%;'>
+      <img id='ml_video' style='position: absolute; 
+      height: ${displayHeight}px;user-select: none; 
+      -webkit-user-drag: none;
+      outline: 2px solid #eaeaea;
+      box-shadow: 8px 9px 17px #575757;' />
+      ${Array.from(grids, g => {
+        const { label: tag, grid } = g
+        const [dx, dy, dw, dh] = grid
+        const base64Data = drawCanvasWithText(dw, dh, tag)
+
+        let x = 0,
+          y = 0,
+          width = (cw * displayHeight) / ch,
+          height = displayHeight
+
+        let imgWidth = cw
+        let imgHeight = ch
+
+        if (dw > 0 && dh > 0) {
+          // 相同尺寸窗口，恢复选区
+          x = (width * dx) / imgWidth
+          y = (height * dy) / imgHeight
+          width = (width * dw) / imgWidth
+          height = (height * dh) / imgHeight
+        }
+
+        return `<div class='ml_selection' 
+          data-tag="${tag}"
+          style='position:absolute; 
+          border: 2px dashed red;
+          pointer-events: none;
+          background-image: url("${base64Data}");
+                background-repeat: no-repeat;
+                background-size: cover;
+                left:${x}px;
+                top:${y}px;
+                width:${width}px;
+                height:${height}px;
+                '></div>`
+      })}
+      <div class="mx_close"> X </div>
+    </div>`
+  // document.body.querySelector('#ml_overlay')
+  document.body.appendChild(div)
+
+  const tags = Array.from(grids, g => g.label)
+  let select = createSelectWithOptions(tags)
+  document.body.appendChild(select)
+
+  let img = div.querySelector('#ml_video')
+  // let overlay = div.querySelector('#ml_overlay')
+  let selections = [...div.querySelectorAll('.ml_selection')]
+
+  let selection = selections.filter(
+    s => s.getAttribute('data-tag') === select.value
+  )[0]
+
+  select.addEventListener('change', e => {
+    selection = selections.filter(
+      s => s.getAttribute('data-tag') === select.value
+    )[0]
+  })
+
+  // console.log(select.value,selection)
+  let close = div.querySelector('.mx_close')
+  let startX, startY, endX, endY
+  let start = false
+  let setDone = false
+  // Set video source
+  img.src = base64
+  // canvas.toDataURL();
+  close.style = `cursor: pointer;
+  position: fixed;
+  left: 12px;
+  top: 12px;
+  z-index: 99999999;
+  background: black;
+  width: 44px;
+  height: 44px;
+  text-align: center;
+  line-height: 44px;`
+
+  // Add mouse events
+  img.addEventListener('mousedown', startSelection)
+  img.addEventListener('mousemove', updateSelection)
+  img.addEventListener('mouseup', endSelection)
+
+  const removeDiv = () => {
+    div.remove()
+    select?.remove()
+    close.removeEventListener('click', removeDiv)
+    img.removeEventListener('mousedown', startSelection)
+    img.removeEventListener('mousemove', updateSelection)
+    img.removeEventListener('mouseup', endSelection)
+    img.removeEventListener('mousedown', setDoneCheck)
+  }
+  close.addEventListener('click', removeDiv)
+
+  const setDoneCheck = event => {
+    console.log(setDone)
+    if (setDone) {
+      img.addEventListener('mousedown', startSelection)
+      img.addEventListener('mousemove', updateSelection)
+      img.addEventListener('mouseup', endSelection)
+      setDone = false
+      start = false
+      startX = event.clientX
+      startY = event.clientY
+    }
+  }
+  img.addEventListener('mousedown', setDoneCheck)
+
+  function remove () {
+    img.removeEventListener('mousedown', startSelection)
+    img.removeEventListener('mousemove', updateSelection)
+    img.removeEventListener('mouseup', endSelection)
+    setDone = true
+    // select?.remove()
+  }
+
+  function startSelection (event) {
+    if (start == false) {
+      startX = event.clientX
+      startY = event.clientY
+      updateSelection(event)
+      start = true
+    } else {
+    }
+  }
+
+  function updateSelection (event) {
+    endX = event.clientX
+    endY = event.clientY
+
+    // Calculate width, height, and coordinates
+    let width = Math.abs(endX - startX)
+    let height = Math.abs(endY - startY)
+    let left = Math.min(startX, endX)
+    let top = Math.min(startY, endY)
+
+    // Set selection style
+    selection.style.left = left + 'px'
+    selection.style.top = top + 'px'
+    selection.style.width = width + 'px'
+    selection.style.height = height + 'px'
+  }
+
+  function endSelection (event) {
+    endX = event.clientX
+    endY = event.clientY
+
+    // 获取img元素的真实宽度和高度
+    let imgWidth = img.naturalWidth
+    let imgHeight = img.naturalHeight
+
+    // 换算起始坐标
+    let realStartX = (startX / img.offsetWidth) * imgWidth
+    let realStartY = (startY / img.offsetHeight) * imgHeight
+
+    // 换算起始坐标
+    let realEndX = (endX / img.offsetWidth) * imgWidth
+    let realEndY = (endY / img.offsetHeight) * imgHeight
+
+    startX = realStartX
+    startY = realStartY
+    endX = realEndX
+    endY = realEndY
+    // Calculate width, height, and coordinates
+    let width = Math.round(Math.abs(endX - startX))
+    let height = Math.round(Math.abs(endY - startY))
+    let left = Math.round(Math.min(startX, endX))
+    let top = Math.round(Math.min(startY, endY))
+
+    if (width <= 0 && height <= 0) return remove()
+
+    if (!!fn) fn(select.value, left, top, width, height)
+
+    remove()
+  }
+}
+
 app.registerExtension({
   name: 'Mixlab.layer.ShowLayer',
   async getCustomWidgets (app) {
@@ -598,8 +847,8 @@ app.registerExtension({
             }
             try {
               console.log('this.inputs', this.id)
-              let imgs=findImages(this.id)
-              
+              let imgs = findImages(this.id)
+
               // let topLinkId = this.inputs[0].link
               // let topNodeId = app.graph.links[topLinkId].origin_id
               let topIm = imgs[0]
@@ -607,9 +856,9 @@ app.registerExtension({
               let linkId = this.inputs[3].link
               let nodeId = app.graph.links[linkId].origin_id
               // console.log(linkId,this.inputs)
-              let imgs2=findImages(nodeId)
+              let imgs2 = findImages(nodeId)
               let im = imgs2[0]
-              console.log(topIm,im)
+              console.log(topIm, im)
               // let src = im.src
               setArea(
                 im.naturalWidth,
@@ -641,3 +890,306 @@ app.registerExtension({
     }
   }
 })
+
+app.registerExtension({
+  name: 'Mixlab.layer.GridInput',
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeType.comfyClass == 'GridInput') {
+      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+      nodeType.prototype.onNodeCreated = async function () {
+        orig_nodeCreated?.apply(this, arguments)
+
+        const grids_widget = this.widgets.filter(w => w.name == 'grids')[0]
+
+        const widget = {
+          type: 'div',
+          name: 'upload',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, widget_width, y, node.size[1]),
+              {
+                justifyContent: 'flex-start'
+              }
+            )
+          }
+        }
+
+        widget.div = $el('div', {})
+
+        const addBtn = document.createElement('button')
+        addBtn.innerText = 'Add Box'
+        addBtn.style = `cursor: pointer;
+        font-weight: 300;
+        margin: 2px; 
+        color: var(--descrip-text);
+        background-color: var(--comfy-input-bg);
+        border-radius: 8px;
+        border-color: var(--border-color);
+        border-style: solid;height: 30px;min-width: 122px;
+       `
+
+        const vbtn = document.createElement('button')
+        vbtn.innerText = 'Set Box'
+        vbtn.style = `cursor: pointer;
+        font-weight: 300;
+        margin: 2px; 
+        color: var(--descrip-text);
+        background-color: var(--comfy-input-bg);
+        border-radius: 8px;
+        border-color: var(--border-color);
+        border-style: solid;height: 30px;min-width: 122px;
+       `
+
+        const btn = document.createElement('button')
+        btn.innerText = 'Upload JSON'
+
+        btn.style = `cursor: pointer;
+        font-weight: 300;
+        margin: 2px; 
+        color: var(--descrip-text);
+        background-color: var(--comfy-input-bg);
+        border-radius: 8px;
+        border-color: var(--border-color);
+        border-style: solid;height: 30px;min-width: 122px;
+       `
+
+        addBtn.addEventListener('click', () => {
+          const { width, height, grids } = JSON.parse(grids_widget.value)
+          grids.push({
+            label: 'background',
+            grid: [12, 12, width - 24, height - 24]
+          })
+          grids_widget.value = JSON.stringify(
+            {
+              width,
+              height,
+              grids
+            },
+            null,
+            2
+          )
+        })
+
+        vbtn.addEventListener('click', () => {
+          const { width, height, grids } = JSON.parse(grids_widget.value)
+
+          setAreaTags(width, height, grids, (tag, x, y, w, h) => {
+            grids_widget.value = JSON.stringify(
+              {
+                width,
+                height,
+                grids: Array.from(grids, g => {
+                  if (g.label === tag) {
+                    g.grid = [x, y, w, h]
+                  }
+                  return g
+                })
+              },
+              null,
+              2
+            )
+          })
+        })
+
+        btn.addEventListener('click', () => {
+          let inp = document.createElement('input')
+          inp.type = 'file'
+          inp.accept = '.json'
+          inp.click()
+          inp.addEventListener('change', event => {
+            // 获取选择的文件
+            const file = event.target.files[0]
+            this.title = file.name.split('.')[0]
+
+            // console.log(file.name.split('.')[0])
+            // 创建文件读取器
+            const reader = new FileReader()
+
+            // 定义读取完成事件的回调函数
+            reader.onload = event => {
+              // 读取完成后的文本内容
+              const fileContent = JSON.parse(event.target.result)
+              const grids = fileContent
+              grids_widget.value = JSON.stringify(grids, null, 2)
+              // widget.value = grids
+
+              inp.remove()
+            }
+
+            // 以文本方式读取文件
+            reader.readAsText(file)
+          })
+        })
+
+        widget.div.appendChild(addBtn)
+        widget.div.appendChild(vbtn)
+        widget.div.appendChild(btn)
+        document.body.appendChild(widget.div)
+        this.addCustomWidget(widget)
+
+        const onExecuted = nodeType.prototype.onExecuted
+        nodeType.prototype.onExecuted = function (message) {
+          const r = onExecuted?.apply?.(this, arguments)
+
+          let json = message.json
+          if (json) {
+            json = {
+              width: json[0],
+              height: json[1],
+              grids: json[2]
+            }
+            grids_widget.value = JSON.stringify(json, null, 2)
+            // widget.value = json
+          }
+
+          return r
+        }
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          widget.div.remove()
+          return onRemoved?.()
+        }
+
+        if (this.onResize) {
+          this.onResize(this.size)
+        }
+
+        this.serialize_widgets = true //需要保存参数
+      }
+    }
+  },
+  async loadedGraphNode (node, app) {
+    if (node.type === 'GridInput') {
+      try {
+        const grids_widget = node.widgets.filter(w => w.name == 'grids')[0]
+        const { width, height, grids } = JSON.parse(grids_widget.value)
+        console.log('#GridInput', node, grids)
+
+        const div = node.widgets.filter(w => w.name == 'upload')[0]
+        div.div.querySelector('select').innerHTML = Array.from(
+          grids,
+          g => `<option value="${g.label}">${g.label}</option>`
+        ).join('')
+      } catch (error) {}
+    }
+  }
+})
+
+app.registerExtension({
+  name: 'Mixlab.layer.GridDisplayAndSave',
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeType.comfyClass == 'GridDisplayAndSave') {
+      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+      nodeType.prototype.onNodeCreated = async function () {
+        orig_nodeCreated?.apply(this, arguments)
+
+        const grids_widget = this.widgets.filter(w => w.name == 'grids')[0]
+        console.log('GridDisplayAndSave', grids_widget)
+        const widget = {
+          type: 'div',
+          name: 'save_json',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, widget_width, y, node.size[1]),
+              {
+                justifyContent: 'flex-start',
+                flexDirection: 'column'
+              }
+            )
+          }
+        }
+
+        widget.div = $el('div', {})
+
+        const btn = document.createElement('button')
+        btn.innerText = 'Save JSON'
+
+        btn.style = `cursor: pointer;
+        font-weight: 300;
+        margin: 2px; 
+        color: var(--descrip-text);
+        background-color: var(--comfy-input-bg);
+        border-radius: 8px;
+        border-color: var(--border-color);
+        border-style: solid;height: 30px;min-width: 122px;
+        max-width: 122px;
+       `
+
+        btn.addEventListener('click', () => {
+          if (window._mixlab_grid)
+            downloadJsonFile(
+              window._mixlab_grid,
+              this.widgets.filter(w => w.name == 'filename_prefix')[0]?.value +
+                '_grid.json'
+            )
+        })
+
+        widget.div.appendChild(btn)
+        document.body.appendChild(widget.div)
+        this.addCustomWidget(widget)
+
+        const onExecuted = nodeType.prototype.onExecuted
+        nodeType.prototype.onExecuted = function (message) {
+          const r = onExecuted?.apply?.(this, arguments)
+          let save_json = this.widgets.filter(d => d.name == 'save_json')[0]
+          let div = save_json?.div
+          // console.log('Test',message)
+
+          let image = message.image[0]
+          let json = message.json
+          if (image) {
+            const { filename, subfolder, type } = image
+
+            if (!div.querySelector('img')) {
+              let im = new Image()
+              div.appendChild(im)
+              im.style.width = '100%'
+            }
+            div.querySelector('img').src = api.apiURL(
+              `/view?filename=${encodeURIComponent(
+                filename
+              )}&type=${type}&subfolder=${subfolder}${app.getPreviewFormatParam()}${app.getRandParam()}`
+            )
+
+            window._mixlab_grid = {
+              width: json[0],
+              height: json[1],
+              grids: json[2]
+            }
+            // console.log(src)
+          }
+
+          this.onResize?.(this.size)
+
+          return r
+        }
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          widget.div.remove()
+          return onRemoved?.()
+        }
+
+        if (this.onResize) {
+          this.onResize(this.size)
+        }
+
+        this.serialize_widgets = true //需要保存参数
+      }
+    }
+  },
+  async loadedGraphNode (node, app) {
+    if (node.type === 'GridDisplayAndSave') {
+      try {
+        let grids_widget = node.widgets.filter(w => w.name === 'grids')[0]
+        // let ks = getLocalData(`_mixlab_PromptSlide`)
+        let uploadWidget = node.widgets.filter(w => w.name == 'upload')[0]
+        // console.log('##widget', uploadWidget.value)
+        let grids = JSON.parse(uploadWidget.value)
+      } catch (error) {}
+    }
+  }
+})

web/javascript/saveTripoSrMesh.js

@@ -0,0 +1,203 @@
+import { app } from '../../../scripts/app.js'
+import { api } from '../../../scripts/api.js'
+import { $el } from '../../../scripts/ui.js'
+
+
+function get_position_style (ctx, widget_width, y, node_height) {
+  const MARGIN = 14 // the margin around the html element
+
+  /* Create a transform that deals with all the scrolling and zooming */
+  const elRect = ctx.canvas.getBoundingClientRect()
+  const transform = new DOMMatrix()
+    .scaleSelf(
+      elRect.width / ctx.canvas.width,
+      elRect.height / ctx.canvas.height
+    )
+    .multiplySelf(ctx.getTransform())
+    .translateSelf(MARGIN, MARGIN + y)
+
+  return {
+    transformOrigin: '0 0',
+    transform: transform,
+    left: `0`,
+    top: `0`,
+    cursor: 'pointer',
+    position: 'absolute',
+    maxWidth: `${widget_width - MARGIN * 2}px`,
+    // maxHeight: `${node_height - MARGIN * 2}px`, // we're assuming we have the whole height of the node
+    width: `${widget_width - MARGIN * 2}px`,
+    // height: `${node_height * 0.3 - MARGIN * 2}px`,
+    // background: '#EEEEEE',
+    // outline: '1px solid red',
+    display: 'flex',
+    flexDirection: 'column',
+    // alignItems: 'center',
+    justifyContent: 'space-around'
+  }
+}
+
+
+app.registerExtension({
+  name: 'Mixlab.3D.SaveTripoSRMesh',
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeType.comfyClass == 'SaveTripoSRMesh') {
+      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+      nodeType.prototype.onNodeCreated = async function () {
+        orig_nodeCreated?.apply(this, arguments)
+
+        const widget = {
+          type: 'div',
+          name: 'preview',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, widget_width, 88, node.size[1])
+            )
+          }
+          //   value: [],
+          //   async serializeValue (nodeId, widgetIndex) {
+          //     return widget.value
+          //   }
+        }
+
+        widget.div = $el('div', {})
+        widget.div.style.width = `120px`
+
+        document.body.appendChild(widget.div)
+
+        // preview.style = `margin-top: 12px;display: flex;
+        //   justify-content: center;
+        //   align-items: center;background-repeat: no-repeat;background-size: contain;`
+
+        this.addCustomWidget(widget)
+
+        const onResize = this.onResize
+        this.onResize = () => {
+          widget.div.style.width = `${this.size[0]}px`
+          widget.div.style.height = `${this.size[1] - 112}px`
+          let mvs = widget.div.querySelectorAll('model-viewer')
+          for (const m of mvs) {
+            m.style.height = `${Math.round(
+              (this.size[1] - 112) / mvs.length
+            )}px`
+            // console.log(m.style.height)
+          }
+          //   console.log('resize', this.size)
+          return onResize?.apply(this, arguments)
+        }
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          widget.div.remove()
+          return onRemoved?.()
+        }
+
+        if (this.onResize) {
+          this.onResize(this.size)
+        }
+        // this.isVirtualNode = true
+        this.serialize_widgets = false //需要保存参数
+      }
+
+      const onExecuted = nodeType.prototype.onExecuted
+      nodeType.prototype.onExecuted = function (message) {
+        const r = onExecuted?.apply?.(this, arguments)
+
+        let widget = this.widgets.filter(d => d.name == 'preview')[0]
+        console.log('Test', widget, message)
+
+        let meshes = message.mesh
+        widget.div.innerHTML = ''
+
+        for (const mesh of meshes) {
+          if (mesh) {
+            const { filename, subfolder, type } = mesh
+            const fileURL = api.apiURL(
+              `/view?filename=${encodeURIComponent(
+                filename
+              )}&type=${type}&subfolder=${subfolder}${app.getPreviewFormatParam()}${app.getRandParam()}`
+            )
+
+            let modelViewer = document.createElement('div')
+            modelViewer.innerHTML = `<model-viewer  src="${fileURL}" 
+              min-field-of-view="0deg" max-field-of-view="180deg"
+              shadow-intensity="1" 
+              camera-controls 
+              touch-action="pan-y"
+                  style="width:100%;margin:4px;min-height:88px"
+              >
+              
+              <div class="controls">
+            
+              <div><button class="export" style="
+              background-color: var(--comfy-input-bg);
+              border-radius: 8px;
+              border-color: var(--border-color);
+              border-style: solid;
+              color: var(--descrip-text);cursor: pointer;">Export GLB</button></div>
+              
+              </div></model-viewer>`
+            widget.div.appendChild(modelViewer)
+            let modelViewerVariants= modelViewer
+            .querySelector('model-viewer');
+            
+            modelViewer
+              .querySelector('.export')
+              .addEventListener('click', async e => {
+                e.preventDefault()
+                const glTF = await modelViewerVariants.exportScene()
+                const file = new File([glTF], filename)
+                const link = document.createElement('a')
+                link.download = file.name
+                link.href = URL.createObjectURL(file)
+                link.click()
+              })
+          }
+        }
+
+        // widget.value = [meshes]
+
+        this.onResize?.(this.size)
+
+        return r
+      }
+    }
+  },
+  async loadedGraphNode (node, app) {
+    const sleep = (t = 1000) => {
+      return new Promise((res, rej) => {
+        setTimeout(() => res(1), t)
+      })
+    }
+    // if (node.type === 'SaveTripoSRMesh') {
+    //   await sleep(0)
+    //   let widget = node.widgets.filter(w => w.name === 'preview')[0]
+    //   widget.div.innerHTML = ''
+
+    //   for (const mesh of widget.value) {
+    //     if (mesh) {
+    //       const { filename, subfolder, type } = mesh
+    //       const fileURL = api.apiURL(
+    //         `/view?filename=${encodeURIComponent(
+    //           filename
+    //         )}&type=${type}&subfolder=${subfolder}${app.getPreviewFormatParam()}${app.getRandParam()}`
+    //       )
+
+    //       let modelViewer = document.createElement('div')
+    //       modelViewer.innerHTML = `<model-viewer  src="${fileURL}"
+    // 	 min-field-of-view="0deg" max-field-of-view="180deg"
+    // 	 shadow-intensity="1"
+    // 	 camera-controls
+    // 	 touch-action="pan-y">
+
+    // 	 <div class="controls">
+
+    // 	  <div><button class="export">Export GLB</button></div>
+
+    // 	</div></model-viewer>`
+    //       widget.div.appendChild(modelViewer)
+    //     }
+    //   }
+    // }
+  }
+})

web/javascript/smart_connect.js

@@ -46,6 +46,18 @@ const smart_connect_config_input = [
     node_widget_name: 'image',
     inputNodeName: 'LoadImage',
     inputNode_output_name: 'IMAGE'
+  },
+  {
+    node_type: 'TripoSRSampler_',
+    node_widget_name: 'image',
+    inputNodeName: 'LoadImagesToBatch',
+    inputNode_output_name: 'IMAGE'
+  },
+  {
+    node_type: 'TripoSRSampler_',
+    node_widget_name: 'mask',
+    inputNodeName: 'RembgNode_Mix',
+    inputNode_output_name: 'masks'
   }
 ]
 
@@ -74,6 +86,18 @@ const smart_connect_config_output = [
     outputNodeName: 'SaveImage',
     outputNode_input_name: 'images'
   },
+  {
+    node_type: 'VAEDecode',
+    node_output_name: 'IMAGE',
+    outputNodeName: 'AppInfo',
+    outputNode_input_name: 'IMAGE'
+  },
+  {
+    node_type: 'VAEDecode',
+    node_output_name: 'IMAGE',
+    outputNodeName: 'SaveImageAndMetadata_',
+    outputNode_input_name: 'images'
+  },
   {
     node_type: 'Moondream',
     node_output_name: 'STRING',
@@ -181,7 +205,10 @@ export function smart_init () {
       ]
       let node_slotType = config[0]
       // 如果input没有，则创建
-      if (!node.inputs?.filter(inp => inp.name === widget.name)[0]||!node.inputs)
+      if (
+        !node.inputs?.filter(inp => inp.name === widget.name)[0] ||
+        !node.inputs
+      )
         convertToInput(node, widget, config)
       input_node.connectByType(inputNode_slot, node, node_slotType)
     }

web/javascript/video_mixlab.js

@@ -0,0 +1,515 @@
+import { app } from '../../../scripts/app.js'
+import { api } from '../../../scripts/api.js'
+import { ComfyWidgets } from '../../../scripts/widgets.js'
+
+import { $el } from '../../../scripts/ui.js'
+
+// The code is based on ComfyUI-VideoHelperSuite modification.
+
+
+
+function injectCSS (css) {
+  // 检查页面中是否已经存在具有相同内容的style标签
+  const existingStyle = document.querySelector('style')
+  if (existingStyle && existingStyle.textContent === css) {
+    return // 如果已经存在相同的样式，则不进行注入
+  }
+
+  // 创建一个新的style标签，并将CSS内容注入其中
+  const style = document.createElement('style')
+  style.textContent = css
+
+  // 将style标签插入到页面的head元素中
+  const head = document.querySelector('head')
+  head.appendChild(style)
+}
+
+injectCSS(`
+.hidden{
+  display:none !important
+}`)
+
+function get_position_style (ctx, widget_width, y, node_height) {
+  const MARGIN = 4 // the margin around the html element
+
+  /* Create a transform that deals with all the scrolling and zooming */
+  const elRect = ctx.canvas.getBoundingClientRect()
+  const transform = new DOMMatrix()
+    .scaleSelf(
+      elRect.width / ctx.canvas.width,
+      elRect.height / ctx.canvas.height
+    )
+    .multiplySelf(ctx.getTransform())
+    .translateSelf(MARGIN, MARGIN + y)
+
+  return {
+    transformOrigin: '0 0',
+    transform: transform,
+    left: `0`,
+    top: `0`,
+    cursor: 'pointer',
+    position: 'absolute',
+    maxWidth: `${widget_width - MARGIN * 2}px`,
+    // maxHeight: `${node_height - MARGIN * 2}px`, // we're assuming we have the whole height of the node
+    width: `${widget_width - MARGIN * 2}px`,
+    // height: `${node_height * 0.3 - MARGIN * 2}px`,
+    // background: '#EEEEEE',
+    display: 'flex',
+    flexDirection: 'column',
+    // alignItems: 'center',
+    justifyContent: 'space-around'
+  }
+}
+
+function videoUpload (node, inputName, inputData, app) {
+  const imageWidget = node.widgets.find(w => w.name === 'video')
+  let uploadWidget
+
+  const widget = {
+    type: 'div',
+    name: 'upload-preview',
+    draw (ctx, node, widget_width, y, widget_height) {
+      Object.assign(
+        this.div.style,
+        get_position_style(ctx, widget_width, 220, node.size[1]),
+        {
+          outline: '1px solid'
+        }
+      )
+    }
+  }
+
+  widget.div = $el('div', {})
+  widget.div.style.width = `120px`
+  document.body.appendChild(widget.div)
+  node.addCustomWidget(widget)
+  // console.log('#imageWidget', imageWidget)
+  const displayDiv = document.createElement('video')
+  displayDiv.controls = true
+  // displayDiv.style=`width:200px;height:200px`
+  imageWidget.callback = () => {
+    displayDiv.src = `/view?filename=${
+      imageWidget.value
+    }&type=input&subfolder=${''}&rand=${Math.random()}`
+
+    // displayDiv.onloadedmetadata = function () {
+    //   var frameCount = displayDiv.duration * displayDiv.webkitDecodedFrameCount
+    //   console.log('视频帧数：' + frameCount)
+    //   node.widgets.filter(w => w.name == 'video_segment_frames')[0].value =
+    //     frameCount
+    // }
+  }
+
+  if (imageWidget.value) {
+    // console.log(imageWidget.value)
+    displayDiv.src = `/view?filename=${
+      imageWidget.value
+    }&type=input&subfolder=${''}&rand=${Math.random()}`
+  }
+
+  widget.div.appendChild(displayDiv)
+
+  const onRemoved = node.onRemoved
+  node.onRemoved = () => {
+    widget.div.remove()
+    return onRemoved?.()
+  }
+
+  var default_value = imageWidget.value
+  Object.defineProperty(imageWidget, 'value', {
+    set: function (value) {
+      this._real_value = value
+    },
+
+    get: function () {
+      let value = ''
+      if (this._real_value) {
+        value = this._real_value
+      } else {
+        return default_value
+      }
+
+      if (value.filename) {
+        let real_value = value
+        value = ''
+        if (real_value.subfolder) {
+          value = real_value.subfolder + '/'
+        }
+
+        value += real_value.filename
+
+        if (real_value.type && real_value.type !== 'input')
+          value += ` [${real_value.type}]`
+      }
+      return value
+    }
+  })
+  async function uploadFile (file, updateNode, pasted = false) {
+    try {
+      // Wrap file in formdata so it includes filename
+      const body = new FormData()
+      body.append('image', file)
+      if (pasted) body.append('subfolder', 'pasted')
+      const resp = await api.fetchApi('/upload/image', {
+        method: 'POST',
+        body
+      })
+
+      if (resp.status === 200) {
+        const data = await resp.json()
+        // Add the file to the dropdown list and update the widget value
+        let path = data.name
+        if (data.subfolder) path = data.subfolder + '/' + path
+
+        if (!imageWidget.options.values.includes(path)) {
+          imageWidget.options.values.push(path)
+        }
+
+        if (updateNode) {
+          imageWidget.value = path
+        }
+
+        return `/view?filename=${path}&type=input&subfolder=${
+          pasted ? 'pasted' : ''
+        }&rand=${Math.random()}`
+      } else {
+        alert(resp.status + ' - ' + resp.statusText)
+      }
+    } catch (error) {
+      alert(error)
+    }
+  }
+
+  const fileInput = document.createElement('input')
+  Object.assign(fileInput, {
+    type: 'file',
+    accept: 'video/*,.mkv,video/webm,video/mp4,video/x-matroska,image/gif',
+    style: 'display: none',
+    onchange: async () => {
+      if (fileInput.files.length) {
+        let file = fileInput.files[0]
+
+        const url = await uploadFile(file, true)
+
+        // console.log('fileInput', file)
+        var reader = new FileReader()
+        reader.onload = function () {
+          displayDiv.src = url
+          displayDiv.onloadedmetadata = function () {
+            // var frameCount =
+            //   displayDiv.duration * displayDiv.webkitDecodedFrameCount
+            // console.log('视频帧数：' + frameCount)
+            // node.widgets.filter(
+            //   w => w.name == 'video_segment_frames'
+            // )[0].value = frameCount
+          }
+        }
+        reader.readAsDataURL(file)
+      }
+    }
+  })
+  document.body.append(fileInput)
+
+  // Create the button widget for selecting the files
+  uploadWidget = node.addWidget('button', 'upload file', 'video', () => {
+    fileInput.click()
+  })
+  uploadWidget.serialize = false
+  return { widget: uploadWidget }
+}
+ComfyWidgets.VIDEOUPLOAD_ = videoUpload
+
+app.registerExtension({
+  name: 'Mixlab.Video.LoadVideoAndSegment_',
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeData?.name == 'LoadVideoAndSegment_') {
+      nodeData.input.required.upload = ['VIDEOUPLOAD_']
+    }
+  },
+  async loadedGraphNode (node, app) {
+    if (node.type === 'LoadVideoAndSegment_') {
+      const imageWidget = node.widgets.find(w => w.name === 'video')
+      const uploadPreview = node.widgets.find(w => w.name === 'upload-preview')
+      if (imageWidget.value) {
+        // console.log(imageWidget.value)
+        uploadPreview.div.querySelector('video').src = `/view?filename=${
+          imageWidget.value
+        }&type=input&subfolder=${''}&rand=${Math.random()}`
+      }
+    }
+  }
+})
+
+
+function offsetDOMWidget(
+  widget,
+  ctx,
+  node,
+  widgetWidth,
+  widgetY,
+  height
+) {
+  const margin = 10
+  const elRect = ctx.canvas.getBoundingClientRect()
+  const transform = new DOMMatrix()
+    .scaleSelf(
+      elRect.width / ctx.canvas.width,
+      elRect.height / ctx.canvas.height
+    )
+    .multiplySelf(ctx.getTransform())
+    .translateSelf(0, widgetY + margin)
+
+  const scale = new DOMMatrix().scaleSelf(transform.a, transform.d)
+  Object.assign(widget.inputEl.style, {
+    transformOrigin: '0 0',
+    transform: scale,
+    left: `${transform.e}px`,
+    top: `${transform.d + transform.f}px`,
+    width: `${widgetWidth}px`,
+    height: `${(height || widget.parent?.inputHeight || 32) - margin}px`,
+    position: 'absolute',
+    background: !node.color ? '' : node.color,
+    color: !node.color ? '' : 'white',
+    zIndex: 5, //app.graph._nodes.indexOf(node),
+  })
+}
+
+export const hasWidgets = (node) => {
+  if (!node.widgets || !node.widgets?.[Symbol.iterator]) {
+    return false
+  }
+  return true
+}
+
+export const cleanupNode = (node) => {
+  if (!hasWidgets(node)) {
+    return
+  }
+
+  for (const w of node.widgets) {
+    if (w.canvas) {
+      w.canvas.remove()
+    }
+    if (w.inputEl) {
+      w.inputEl.remove()
+    }
+    // calls the widget remove callback
+    w.onRemoved?.()
+  }
+}
+
+const CreatePreviewElement = (name, val, format) => {
+  const [type] = format.split('/')  
+  const w = {
+      name,
+      type,
+      value: val,
+      draw: function (ctx, node, widgetWidth, widgetY, height) {
+        const [cw, ch] = this.computeSize(widgetWidth)
+        offsetDOMWidget(this, ctx, node, widgetWidth, widgetY, ch)
+      },
+      computeSize: function (_) {
+        const ratio = this.inputRatio || 1
+        const width = Math.max(220, this.parent.size[0])
+        return [width, (width / ratio + 10)]
+      },
+      onRemoved: function () {
+        if (this.inputEl) {
+          this.inputEl.remove()
+        }
+      },
+    }
+
+    w.inputEl = document.createElement(type === 'video' ? 'video' : 'img')
+    w.inputEl.src = w.value
+    if (type === 'video') {
+      w.inputEl.setAttribute('type', 'video/webm');
+      w.inputEl.autoplay = true
+      w.inputEl.loop = true
+      w.inputEl.controls = false;
+    }
+    w.inputEl.onload = function () {
+      w.inputRatio = w.inputEl.naturalWidth / w.inputEl.naturalHeight
+    }
+    document.body.appendChild(w.inputEl)
+    return w
+  }
+
+
+app.registerExtension({
+  name: 'Mixlab.Video.ImageListReplace',
+  async beforeRegisterNodeDef (nodeType, nodeData, app) {
+    if (nodeData?.name == 'ImageListReplace_') {
+      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+      nodeType.prototype.onNodeCreated = function () {
+        orig_nodeCreated?.apply(this, arguments)
+        const widget = {
+          type: 'div',
+          name: 'preview',
+          draw (ctx, node, widget_width, y, widget_height) {
+            Object.assign(
+              this.div.style,
+              get_position_style(ctx, widget_width, 188, node.size[1]),
+              {
+                outline: '1px solid',
+                display: 'flex',
+                flexWrap: 'wrap',
+                flexDirection: 'row',
+                justifyContent: 'flex-start'
+              }
+            )
+          }
+        }
+
+        widget.div = $el('div', {})
+        widget.div.style.width = `120px`
+        widget.div.className = 'hidden'
+        document.body.appendChild(widget.div)
+        this.addCustomWidget(widget)
+        // console.log('#ImageListReplace', widget)
+
+        const onRemoved = this.onRemoved
+        this.onRemoved = () => {
+          widget.div.remove()
+          return onRemoved?.()
+        }
+      }
+
+      const onExecuted = nodeType.prototype.onExecuted
+      nodeType.prototype.onExecuted = function (message) {
+        onExecuted?.apply(this, arguments)
+
+        // let _image_replace = message._image_replace[0]
+        // _image_replace = `/view?filename=${_image_replace.filename}&type=${
+        //   _image_replace.type
+        // }&subfolder=${_image_replace.subfolder}&rand=${Math.random()}`
+
+        let preview = this.widgets.filter(w => w.name == 'preview')[0]
+
+        if (message._images.length > 0) {
+          preview.div.className = ''
+          // console.log('#ImageListReplace', preview.div)
+        }
+
+        preview.div.innerHTML = ''
+        for (const img_ of message._images) {
+          let img = new Image()
+          img.style = `width: 100px;
+          margin: 4px;`
+          img.src = `/view?filename=${img_.filename}&type=${
+            img_.type
+          }&subfolder=${img_.subfolder}&rand=${Math.random()}`
+          preview.div.appendChild(img)
+        }
+
+        let start_index = this.widgets.filter(w => w.name == 'start_index')[0]
+        let end_index = this.widgets.filter(w => w.name == 'end_index')[0]
+        let invert = this.widgets.filter(w => w.name == 'invert')[0]
+        let _sc = start_index.callback.bind(start_index)
+        let _ec = end_index.callback.bind(end_index)
+
+        const selectImages = () => {
+          // console.log(v)
+          let s = start_index.value,
+            e = end_index.value
+          let imgs = preview.div.querySelectorAll('img')
+          for (let index = 0; index < imgs.length; index++) {
+            if (invert.value) {
+              imgs[index].style.outline =
+                index >= s && index <= e ? 'none' : '4px solid #cbd3fe'
+            } else {
+              imgs[index].style.outline =
+                index >= s && index <= e ? '4px solid #cbd3fe' : 'none'
+            }
+          }
+        }
+
+        selectImages()
+
+        start_index.callback = v => {
+          let s = v,
+            e = end_index.value
+          let imgs = preview.div.querySelectorAll('img')
+          for (let index = 0; index < imgs.length; index++) {
+            if (invert.value) {
+              imgs[index].style.outline =
+                index >= s && index <= e ? 'none' : '4px solid #cbd3fe'
+            } else {
+              imgs[index].style.outline =
+                index >= s && index <= e ? '4px solid #cbd3fe' : 'none'
+            }
+          }
+
+          _sc(v)
+        }
+
+        end_index.callback = v => {
+          let s = start_index.value,
+            e = v
+          let imgs = preview.div.querySelectorAll('img')
+          for (let index = 0; index < imgs.length; index++) {
+            if (invert.value) {
+              imgs[index].style.outline =
+                index >= s && index <= e ? 'none' : '4px solid #cbd3fe'
+            } else {
+              imgs[index].style.outline =
+                index >= s && index <= e ? '4px solid #cbd3fe' : 'none'
+            }
+          }
+
+          _ec(v)
+        }
+
+        invert.callback = v => {
+          selectImages()
+        }
+
+        try {
+        } catch (error) {}
+      }
+    }
+
+    if (nodeData?.name == 'VideoCombine_Adv') {
+      const onExecuted = nodeType.prototype.onExecuted
+      nodeType.prototype.onExecuted = function (message) {
+        const prefix = 'vhs_gif_preview_'
+        const r = onExecuted ? onExecuted.apply(this, message) : undefined
+
+        if (this.widgets) {
+          const pos = this.widgets.findIndex(w => w.name === `${prefix}_0`)
+          if (pos !== -1) {
+            for (let i = pos; i < this.widgets.length; i++) {
+              this.widgets[i].onRemoved?.()
+            }
+            this.widgets.length = pos
+          }
+          if (message?.gifs) {
+            message.gifs.forEach((params, i) => {
+              const previewUrl = api.apiURL(
+                '/view?' + new URLSearchParams(params).toString()
+              )
+              const w = this.addCustomWidget(
+                CreatePreviewElement(
+                  `${prefix}_${i}`,
+                  previewUrl,
+                  params.format || 'image/gif'
+                )
+              )
+              w.parent = this
+            })
+          }
+          const onRemoved = this.onRemoved
+          this.onRemoved = () => {
+            cleanupNode(this)
+            return onRemoved?.()
+          }
+        }
+        this.setSize([
+          this.size[0],
+          this.computeSize([this.size[0], this.size[1]])[1]
+        ])
+        return r
+      }
+    }
+  }
+})

web/lib/juxtapose.css

@@ -0,0 +1,347 @@
+/* juxtapose - v1.2.2 - 2020-09-03
+ * Copyright (c) 2020 Alex Duner and Northwestern University Knight Lab
+ */
+div.juxtapose {
+	width: 100%;
+	font-family: Helvetica, Arial, sans-serif;
+}
+
+div.jx-slider {
+	width: 100%;
+	height: 100%;
+	position: relative;
+	overflow: hidden;
+	cursor: pointer;
+	color: #f3f3f3;
+}
+
+
+div.jx-handle {
+	position: absolute;
+	height: 100%;
+	width: 40px;
+	cursor: col-resize;
+	z-index: 15;
+	margin-left: -20px;
+}
+
+.vertical div.jx-handle {
+	height: 40px;
+	width: 100%;
+	cursor: row-resize;
+	margin-top: -20px;
+	margin-left: 0;
+}
+
+div.jx-control {
+	height: 100%;
+	margin-right: auto;
+	margin-left: auto;
+	width: 3px;
+	background-color: currentColor;
+}
+
+.vertical div.jx-control {
+	height: 3px;
+	width: 100%;
+	background-color: currentColor;
+	position: relative;
+	top: 50%;
+	transform: translateY(-50%);
+}
+
+div.jx-controller {
+	position: absolute;
+	margin: auto;
+	top: 0;
+	bottom: 0;
+	height: 60px;
+	width: 9px;
+	margin-left: -3px;
+	background-color: currentColor;
+}
+
+.vertical div.jx-controller {
+	height: 9px;
+	width: 100px;
+	margin-left: auto;
+	margin-right: auto;
+	top: -3px;
+	position: relative;
+}
+
+div.jx-arrow {
+	position: absolute;
+	margin: auto;
+	top: 0;
+	bottom: 0;
+	width: 0;
+	height: 0;
+	transition: all .2s ease;
+}
+
+.vertical div.jx-arrow {
+	position: absolute;
+	margin: 0 auto;
+	left: 0;
+	right: 0;
+	width: 0;
+	height: 0;
+	transition: all .2s ease;
+}
+
+
+div.jx-arrow.jx-left {
+	left: 2px;
+	border-style: solid;
+	border-width: 8px 8px 8px 0;
+	border-color: transparent currentColor transparent transparent;
+}
+
+div.jx-arrow.jx-right {
+	right: 2px;
+	border-style: solid;
+	border-width: 8px 0 8px 8px;
+	border-color: transparent transparent transparent currentColor;
+}
+
+.vertical div.jx-arrow.jx-left {
+	left: 0px;
+	top: 2px;
+	border-style: solid;
+	border-width: 0px 8px 8px 8px;
+	border-color: transparent transparent currentColor transparent;
+}
+
+.vertical div.jx-arrow.jx-right {
+	right: 0px;
+	top: auto;
+	bottom: 2px;
+	border-style: solid;
+	border-width: 8px 8px 0 8px;
+	border-color: currentColor transparent transparent transparent;
+}
+
+div.jx-handle:hover div.jx-arrow.jx-left,
+div.jx-handle:active div.jx-arrow.jx-left {
+	left: -1px;
+}
+
+div.jx-handle:hover div.jx-arrow.jx-right,
+div.jx-handle:active div.jx-arrow.jx-right {
+	right: -1px;
+}
+
+.vertical div.jx-handle:hover div.jx-arrow.jx-left,
+.vertical div.jx-handle:active div.jx-arrow.jx-left {
+	left: 0px;
+	top: 0px;
+}
+
+.vertical div.jx-handle:hover div.jx-arrow.jx-right,
+.vertical div.jx-handle:active div.jx-arrow.jx-right {
+	right: 0px;
+	bottom: 0px;
+}
+
+
+div.jx-image {
+	position: absolute;
+	height: 100%;
+	display: inline-block;
+	top: 0;
+	overflow: hidden;
+	-webkit-backface-visibility: hidden;
+}
+
+.vertical div.jx-image {
+	width: 100%;
+	left: 0;
+	top: auto;
+}
+
+div.jx-image img {
+	height: 100%;
+	width: auto;
+	z-index: 5;
+	position: absolute;
+	margin-bottom: 0;
+
+	max-height: none;
+	max-width: none;
+	max-height: initial;
+	max-width: initial;
+}
+
+.vertical div.jx-image img {
+	height: auto;
+	width: 100%;
+}
+
+div.jx-image.jx-left {
+	left: 0;
+	background-position: left;
+}
+
+div.jx-image.jx-left img {
+	left: 0;
+}
+
+div.jx-image.jx-right {
+	right: 0;
+	background-position: right;
+}
+
+div.jx-image.jx-right img {
+	right: 0;
+	bottom: 0;
+}
+
+
+.veritcal div.jx-image.jx-left {
+	top: 0;
+	background-position: top;
+}
+
+.veritcal div.jx-image.jx-left img {
+	top: 0;
+}
+
+.vertical div.jx-image.jx-right {
+	bottom: 0;
+	background-position: bottom;
+}
+
+.veritcal div.jx-image.jx-right img {
+	bottom: 0;
+}
+
+
+div.jx-image div.jx-label {
+	font-size: 1em;
+	padding: .25em .75em;
+	position: relative;
+	display: inline-block;
+	top: 0;
+	background-color: #000; /* IE 8 */
+	background-color: rgba(0,0,0,.7);
+	color: white;
+	z-index: 10;
+	white-space: nowrap;
+	line-height: 18px;
+	vertical-align: middle;
+}
+
+div.jx-image.jx-left div.jx-label {
+	float: left;
+	left: 0;
+}
+
+div.jx-image.jx-right div.jx-label {
+	float: right;
+	right: 0;
+}
+
+.vertical div.jx-image div.jx-label {
+	display: table;
+	position: absolute;
+}
+
+.vertical div.jx-image.jx-right div.jx-label {
+	left: 0;
+	bottom: 0;
+	top: auto;
+}
+
+div.jx-credit {
+	line-height: 1.1;
+	font-size: 0.75em;
+}
+
+div.jx-credit em {
+	font-weight: bold;
+	font-style: normal;
+}
+
+
+/* Animation */
+
+div.jx-image.transition {
+	transition: width .5s ease;
+}
+
+div.jx-handle.transition {
+	transition: left .5s ease;
+}
+
+.vertical div.jx-image.transition {
+	transition: height .5s ease;
+}
+
+.vertical div.jx-handle.transition {
+	transition: top .5s ease;
+}
+
+/* Knight Lab Credit */
+a.jx-knightlab {
+	background-color: #000; /* IE 8 */
+	background-color: rgba(0,0,0,.25);
+	bottom: 0;
+	display: table;
+	height: 14px;
+	line-height: 14px;
+	padding: 1px 4px 1px 5px;
+	position: absolute;
+	right: 0;
+	text-decoration: none;
+	z-index: 10;
+}
+
+a.jx-knightlab div.knightlab-logo {
+	display: inline-block;
+	vertical-align: middle;
+	height: 8px;
+	width: 8px;
+	background-color: #c34528;
+	transform: rotate(45deg);
+	-ms-transform: rotate(45deg);
+	-webkit-transform: rotate(45deg);
+	top: -1.25px;
+	position: relative;
+	cursor: pointer;
+}
+
+a.jx-knightlab:hover {
+	background-color: #000; /* IE 8 */
+	background-color: rgba(0,0,0,.35);
+}
+a.jx-knightlab:hover div.knightlab-logo {
+	background-color: #ce4d28;
+}
+
+a.jx-knightlab span.juxtapose-name {
+	display: table-cell;
+	margin: 0;
+	padding: 0;
+	font-family: Helvetica, Arial, sans-serif;
+	font-weight: 300;
+	color: white;
+	font-size: 10px;
+	padding-left: 0.375em;
+	vertical-align: middle;
+	line-height: normal;
+	text-shadow: none;
+}
+
+/* keyboard accessibility */
+div.jx-controller:focus,
+div.jx-image.jx-left div.jx-label:focus,
+div.jx-image.jx-right div.jx-label:focus,
+a.jx-knightlab:focus {
+	background: #eae34a;
+	color: #000;
+}
+a.jx-knightlab:focus span.juxtapose-name{
+	color: #000;
+	border: none;
+}

web/lib/miniPaint-4.14.2/dist/bundle.js.LICENSE.txt

@@ -0,0 +1,47 @@
+/*
+object-assign
+(c) Sindre Sorhus
+@license MIT
+*/
+
+/*!
+
+pica
+https://github.com/nodeca/pica
+
+*/
+
+/*!
+ * Block below copied from Protovis: http://mbostock.github.com/protovis/
+ * Copyright 2010 Stanford Visualization Group
+ * Licensed under the BSD License: http://www.opensource.org/licenses/bsd-license.php
+ * @license
+ */
+
+/*!
+ * jQuery JavaScript Library v3.7.1
+ * https://jquery.com/
+ *
+ * Copyright OpenJS Foundation and other contributors
+ * Released under the MIT license
+ * https://jquery.org/license
+ *
+ * Date: 2023-08-28T13:37Z
+ */
+
+/*!
+ * quantize.js Copyright 2008 Nick Rabinowitz.
+ * Licensed under the MIT license: http://www.opensource.org/licenses/mit-license.php
+ * @license
+ */
+
+/*! alertifyjs - v1.13.1 - Mohammad Younes <Mohammad@alertifyjs.com> (http://alertifyjs.com) */
+
+/*! regenerator-runtime -- Copyright (c) 2014-present, Facebook, Inc. -- license (MIT): https://github.com/facebook/regenerator/blob/main/LICENSE */
+
+/**
+ * hermite-resize - Canvas image resize/resample using Hermite filter with JavaScript.
+ * @version v2.2.10
+ * @link https://github.com/viliusle/miniPaint
+ * @license MIT
+ */

web/lib/miniPaint-4.14.2/dist/bundle.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"bundle.js","sources":["webpack://miniPaint/bundle.js"],"mappings":";AAAA","sourceRoot":""}

web/lib/miniPaint-4.14.2/images/favicon.svg

@@ -0,0 +1,13 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 18.1.1, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 298.73 298.73" style="enable-background:new 0 0 298.73 298.73;" xml:space="preserve">
+<g>
+	<path style="fill:#010002;" d="M264.959,9.35H33.787C15.153,9.35,0,24.498,0,43.154v212.461c0,18.634,15.153,33.766,33.787,33.766
+		h231.171c18.634,0,33.771-15.132,33.771-33.766V43.154C298.73,24.498,283.593,9.35,264.959,9.35z M193.174,59.623
+		c18.02,0,32.634,14.615,32.634,32.634s-14.615,32.634-32.634,32.634c-18.025,0-32.634-14.615-32.634-32.634
+		S175.149,59.623,193.174,59.623z M254.363,258.149H149.362H49.039c-9.013,0-13.027-6.521-8.964-14.566l56.006-110.93
+		c4.058-8.044,11.792-8.762,17.269-1.605l56.316,73.596c5.477,7.158,15.05,7.767,21.386,1.354l13.777-13.951
+		c6.331-6.413,15.659-5.619,20.826,1.762l35.675,50.959C266.487,252.16,263.376,258.149,254.363,258.149z"/>
+</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/animation.svg

@@ -0,0 +1,13 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 18.1.1, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 30.051 30.051" style="enable-background:new 0 0 30.051 30.051;" xml:space="preserve">
+<g>
+	<path d="M19.982,14.438l-6.24-4.536c-0.229-0.166-0.533-0.191-0.784-0.062c-0.253,0.128-0.411,0.388-0.411,0.669v9.069
+		c0,0.284,0.158,0.543,0.411,0.671c0.107,0.054,0.224,0.081,0.342,0.081c0.154,0,0.31-0.049,0.442-0.146l6.24-4.532
+		c0.197-0.145,0.312-0.369,0.312-0.607C20.295,14.803,20.177,14.58,19.982,14.438z"/>
+	<path d="M15.026,0.002C6.726,0.002,0,6.728,0,15.028c0,8.297,6.726,15.021,15.026,15.021c8.298,0,15.025-6.725,15.025-15.021
+		C30.052,6.728,23.324,0.002,15.026,0.002z M15.026,27.542c-6.912,0-12.516-5.601-12.516-12.514c0-6.91,5.604-12.518,12.516-12.518
+		c6.911,0,12.514,5.607,12.514,12.518C27.541,21.941,21.937,27.542,15.026,27.542z"/>
+</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/arrow-down.svg

@@ -0,0 +1,5 @@
+<svg width="124" height="150" viewBox="0 0 124 150" xmlns="http://www.w3.org/2000/svg">
+<rect x="55" y="14" width="14" height="63"/>
+<rect x="116" width="14" height="61" transform="rotate(90 116 0)"/>
+<path d="M62 150L8.30643 75L115.694 75L62 150Z"/>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/blur.svg

@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 19.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<svg version="1.1" id="Layer_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 512 512" style="enable-background:new 0 0 512 512;" xml:space="preserve">
+<g>
+	<g>
+		<path d="M264.574,4.675C262.697,1.761,259.467,0,256,0c-3.467,0-6.697,1.761-8.574,4.675
+			c-6.532,10.14-159.966,249.362-159.966,338.784C87.459,436.393,163.066,512,256,512s168.541-75.607,168.541-168.541
+			C424.541,254.037,271.106,14.815,264.574,4.675z M256,491.602c-81.686,0-148.142-66.456-148.142-148.143
+			c0-34.037,26.926-101.269,77.865-194.427C213.83,97.626,242.219,51.324,256,29.29c13.77,22.016,42.123,68.259,70.223,119.64
+			c50.976,93.212,77.92,160.478,77.92,194.529C404.142,425.146,337.686,491.602,256,491.602z"/>
+	</g>
+</g>
+<g>
+	<g>
+		<path d="M375.907,332.939c-5.633,0-10.199,4.566-10.199,10.199c0,43.197-25.482,82.521-64.919,100.181
+			c-5.141,2.301-7.442,8.335-5.14,13.476c1.695,3.788,5.416,6.034,9.314,6.034c1.393,0,2.809-0.287,4.163-0.893
+			c46.764-20.941,76.981-67.572,76.981-118.797C386.106,337.505,381.54,332.939,375.907,332.939z"/>
+	</g>
+</g>
+<g>
+	<g>
+		<path d="M281.818,460.702c-0.729-5.586-5.85-9.519-11.435-8.791c-4.736,0.619-9.574,0.933-14.383,0.933
+			c-5.633,0-10.199,4.566-10.199,10.199c0,5.633,4.566,10.199,10.199,10.199c5.69,0,11.419-0.372,17.028-1.106
+			C278.613,471.407,282.548,466.287,281.818,460.702z"/>
+	</g>
+</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/bold.svg

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<svg width="1em" height="1em" viewBox="0 0 16 16" fill="currentColor" xmlns="http://www.w3.org/2000/svg">
+	<path d="M8.21 13c2.106 0 3.412-1.087 3.412-2.823 0-1.306-.984-2.283-2.324-2.386v-.055a2.176 2.176 0 0 0 1.852-2.14c0-1.51-1.162-2.46-3.014-2.46H3.843V13H8.21zM5.908 4.674h1.696c.963 0 1.517.451 1.517 1.244 0 .834-.629 1.32-1.73 1.32H5.908V4.673zm0 6.788V8.598h1.73c1.217 0 1.88.492 1.88 1.415 0 .943-.643 1.449-1.832 1.449H5.907z"/>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/brush.svg

@@ -0,0 +1 @@
+<svg height="443pt" viewBox="0 0 443.06138 443" width="443pt" xmlns="http://www.w3.org/2000/svg"><path d="m431.328125 25.894531-14.136719-14.136719c-8.070312-8.078124-19.210937-12.320312-30.613281-11.6601558-11.402344.6601558-21.976563 6.1640628-29.058594 15.1249998l-117.023437 164.839844c-4.089844 5.761719-10.511719 9.425781-17.554688 10.019531-7.039062.59375-13.984375-1.945312-18.980468-6.941406l-34.234376-34.207031c-9.480468-9.109375-24.46875-9.109375-33.949218 0l-11.296875 11.304687 158.398437 158.398438 11.304688-11.304688c4.527344-4.488281 7.070312-10.597656 7.070312-16.96875 0-6.375-2.542968-12.484375-7.070312-16.972656l-34.222656-34.234375c-4.996094-4.996094-7.535157-11.941406-6.941407-18.980469.59375-7.042969 4.257813-13.464843 10.019531-17.554687l165-117.183594c8.890626-7.109375 14.332032-17.671875 14.960938-29.035156.628906-11.367188-3.617188-22.464844-11.671875-30.507813zm-24 43.808594c-9.375 9.371094-24.570313 9.371094-33.945313 0-9.371093-9.375-9.371093-24.574219 0-33.945313 9.496094-9.0625 24.441407-9.0625 33.9375 0 9.375 9.371094 9.378907 24.570313.007813 33.945313zm0 0"/><path d="m390.351562 44.734375c-3.234374 0-6.152343 1.949219-7.390624 4.9375-1.234376 2.988281-.550782 6.429687 1.734374 8.71875 3.160157 3.03125 8.148438 3.03125 11.304688 0 2.289062-2.289063 2.972656-5.730469 1.734375-8.71875s-4.15625-4.9375-7.390625-4.9375zm0 0"/><path d="m135.792969 420.429688 22.65625 22.65625 113.085937-113.167969-158.398437-158.402344-113.136719 113.121094 22.65625 22.65625 84.847656-84.855469c2.007813-2.078125 4.984375-2.914062 7.78125-2.183594 2.796875.734375 4.980469 2.917969 5.710938 5.714844.734375 2.796875-.101563 5.773438-2.179688 7.78125l-84.847656 84.855469 22.632812 22.625 50.902344-50.90625c3.140625-3.03125 8.128906-2.988281 11.214844.097656s3.128906 8.078125.097656 11.214844l-50.90625 50.921875 22.617188 22.621094 62.214844-62.222657c3.125-3.125 8.191406-3.128906 11.316406-.003906 3.128906 3.125 3.128906 8.191406.003906 11.316406l-62.222656 62.222657 22.625 22.625 73.535156-73.535157c3.140625-3.03125 8.128906-2.988281 11.214844.097657 3.085937 3.085937 3.128906 8.074218.097656 11.214843zm0 0"/></svg>

web/lib/miniPaint-4.14.2/images/icons/bulge_pinch.svg

@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 19.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 402.56 402.56" style="enable-background:new 0 0 402.56 402.56;" xml:space="preserve">
+<g>
+	<g>
+		<polygon points="38.613,234.88 38.4,274.56 96.64,274.56 116.693,265.173 9.6,372.48 39.68,402.56 147.2,295.253 137.813,316.587 
+			137.813,373.76 177.493,374.187 177.493,234.88 		"/>
+	</g>
+</g>
+<g>
+	<g>
+		<polygon points="306.987,128.213 285.653,137.6 392.96,30.08 362.88,0 255.573,107.307 264.96,87.04 264.96,28.8 225.28,29.013 
+			225.28,167.893 364.587,167.893 364.16,128.213 		"/>
+	</g>
+</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/clone.svg

@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 19.1.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 540.721 540.721" style="enable-background:new 0 0 540.721 540.721;" xml:space="preserve">
+<g>
+	<g>
+		<path d="M521.858,465.271H18.862c-7.545,0-12.575,5.03-12.575,12.575v50.3c0,7.545,5.03,12.575,12.575,12.575h502.996
+			c7.545,0,12.575-5.03,12.575-12.575v-50.3C534.433,470.301,529.403,465.271,521.858,465.271z"/>
+		<path d="M227.606,98.084c5.03-5.03,20.12-12.575,30.18-12.575c7.545,0,12.575-5.03,12.575-12.575c0-7.545-5.03-10.06-12.575-10.06
+			l0,0c-17.605,0-37.725,7.545-47.785,17.605c-12.575,10.06-22.635,27.665-22.635,47.785c0,7.545,5.03,10.06,12.575,10.06l0,0
+			c7.545,0,12.575-2.515,12.575-10.06C212.516,118.204,220.061,103.114,227.606,98.084z"/>
+		<path d="M18.862,440.121h502.996c5.03,0,7.545-2.515,10.06-5.03c2.515-2.515,2.515-7.545,0-12.575l-47.785-100.599
+			c0-5.03-5.03-7.545-10.06-7.545H333.235v-77.964c25.15-20.12,60.359-60.36,60.359-110.659C393.594,55.33,338.265,0,270.36,0
+			S147.126,57.845,147.126,125.749c-2.515,47.785,35.21,88.024,60.36,110.659v77.964H66.647c-5.03,0-10.06,2.515-10.06,7.545
+			L8.802,422.517c-2.515,5.03-2.515,7.545,0,12.575C11.317,440.121,13.832,440.121,18.862,440.121z M169.761,125.749
+			c0-55.33,45.27-100.599,98.084-100.599c52.815,0,98.084,45.27,98.084,100.599c0,37.725-27.665,75.449-55.33,93.054
+			c-2.515,2.515-5.03,5.03-5.03,10.06v123.234c0,2.515-12.575,10.06-22.635,10.06h-37.725c-7.545,0-15.09-7.545-15.09-10.06V228.863
+			c0-5.03-2.515-7.545-5.03-10.06C207.486,206.228,169.761,168.504,169.761,125.749z"/>
+	</g>
+</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/crop.svg

@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 16.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
+<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 width="956.815px" height="956.815px" viewBox="0 0 956.815 956.815" style="enable-background:new 0 0 956.815 956.815;"
+	 xml:space="preserve">
+<g>
+	<path d="M137.621,162.622H20c-11.046,0-20,8.954-20,20v72.919c0,11.046,8.954,20,20,20h117.622L137.621,162.622L137.621,162.622z"
+		/>
+	<path d="M774.193,956.815c11.046,0,20-8.954,20-20V819.193H681.274v117.621c0,11.046,8.954,20,20,20L774.193,956.815
+		L774.193,956.815z"/>
+	<path d="M794.193,656.275V182.622c0-11.046-8.954-20-20-20H300.54v112.919h380.734v380.734H794.193z"/>
+	<path d="M936.814,681.275H794.193H681.274H275.54V275.541V162.622V20c0-11.046-8.954-20-20-20h-72.918c-11.046,0-20,8.954-20,20
+		v142.622v112.919v498.653c0,11.046,8.954,20,20,20h498.653h112.918h142.622c11.045,0,20-8.954,20-20v-72.918
+		C956.814,690.229,947.86,681.275,936.814,681.275z"/>
+</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/delete.svg

@@ -0,0 +1,6 @@
+<svg height="365pt" viewBox="0 0 365.71733 365" width="365pt" xmlns="http://www.w3.org/2000/svg">
+	<g fill="#f44336">
+		<path d="m356.339844 296.347656-286.613282-286.613281c-12.5-12.5-32.765624-12.5-45.246093 0l-15.105469 15.082031c-12.5 12.503906-12.5 32.769532 0 45.25l286.613281 286.613282c12.503907 12.5 32.769531 12.5 45.25 0l15.082031-15.082032c12.523438-12.480468 12.523438-32.75.019532-45.25zm0 0"/>
+		<path d="m295.988281 9.734375-286.613281 286.613281c-12.5 12.5-12.5 32.769532 0 45.25l15.082031 15.082032c12.503907 12.5 32.769531 12.5 45.25 0l286.632813-286.59375c12.503906-12.5 12.503906-32.765626 0-45.246094l-15.082032-15.082032c-12.5-12.523437-32.765624-12.523437-45.269531-.023437zm0 0"/>
+	</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/desaturate.svg

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Svg Vector Icons : http://www.onlinewebfonts.com/icon -->
+<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
+<svg version="1.1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px" viewBox="0 0 1000 1000" enable-background="new 0 0 1000 1000" xml:space="preserve">
+<metadata> Svg Vector Icons : http://www.onlinewebfonts.com/icon </metadata>
+<g><g transform="translate(0.000000,511.000000) scale(0.100000,-0.100000)"><path d="M2646.5,3789.8c-1040.2-119.8-2106-953.5-2445-1917.3l-91.8-260.1L102,355.5c-5.1-1081-2.5-1279.9,33.1-1415c155.5-609.4,578.8-986.7,1282.4-1139.7c300.9-66.3,670.5-96.9,1534.8-127.5c400.3-12.7,831.2-43.3,961.2-63.7c464-76.5,836.3-229.5,1448.1-599.2c698.6-420.7,1063.2-550.7,1682.7-599.2c234.6-17.9,645,20.4,866.9,81.6c494.6,132.6,1007.1,476.8,1341.1,902.6c318.7,405.4,484.4,749.6,589,1223.8c53.5,255,56.1,313.6,58.6,1527.2c0,1155-5.1,1274.8-45.9,1415c-186.1,596.6-609.4,917.9-1374.2,1045.3c-135.1,22.9-550.7,51-963.7,63.7c-1560.3,53.5-1828,112.2-2957.5,662.9c-685.8,334-928,415.6-1328.3,458.9C2944.8,3820.4,2919.3,3820.4,2646.5,3789.8z M3378.2,3165.2c255-51,461.5-135.1,1014.7-413c543.1-270.3,821-382.4,1213.6-487c311-84.2,798-155.5,1042.8-155.5c142.8,0,158.1-5.1,234.6-91.8c112.2-124.9,214.2-163.2,420.7-163.2c193.8,0,359.5,61.2,423.2,153c40.8,56.1,48.5,58.6,328.9,43.3c894.9-51,1249.3-328.9,1246.7-971.4c0-117.3-7.6-247.3-17.8-288.1c-30.6-147.9-94.3-331.4-168.3-494.6l-76.5-168.3l-158.1-12.7c-379.9-30.6-563.4-270.3-372.2-487l73.9-86.7l-81.6-76.5c-211.6-198.9-527.8-359.5-843.9-428.3c-160.6-35.7-272.8-43.3-560.9-30.6c-555.8,20.4-777.6,99.4-1494,525.2C5111.9-174.8,4857-52.4,4482.2,70c-410.5,137.7-685.8,175.9-1402.3,206.5c-1287.5,51-1596,86.7-1891.8,219.3c-191.2,84.1-379.9,288.1-443.6,471.7c-66.3,193.8-58.6,471.7,17.9,696l61.2,175.9l155.5,5.1c226.9,5.1,402.8,119.8,433.4,280.5c20.4,107.1,0,168.3-84.1,249.9l-66.3,66.3l61.2,66.3c249.9,267.7,808.2,571.1,1200.8,657.8C2745.9,3213.6,3143.6,3213.6,3378.2,3165.2z M3391-391.5c221.8-112.2,397.7-351.8,362-492.1c-51-211.6-418.1-183.6-657.8,51C2972.8-712.8,2921.8-610.8,2932-496c5.1,79.1,22.9,102,96.9,137.7C3130.9-307.4,3240.5-317.6,3391-391.5z M1328.3-414.5c140.2-86.7,221.8-209.1,221.8-323.8c0-216.7-392.6-153-573.7,91.8C792.9-394.1,1040.2-230.9,1328.3-414.5z M5550.4-1273.7c135.1-84.1,244.8-257.5,221.8-349.3c-23-89.2-145.3-132.6-267.7-96.9c-137.7,43.4-293.2,163.2-341.7,267.7c-56.1,117.3-53.5,160.6,10.2,224.4C5244.5-1156.4,5387.3-1174.2,5550.4-1273.7z M2682.2-1339.9c293.2-165.7,341.6-476.8,73.9-476.8c-196.3,0-466.6,232-466.6,397.8C2289.5-1278.7,2503.7-1238,2682.2-1339.9z M6659.5-2446.5c209-107.1,372.2-308.5,372.2-461.5c0-68.8-79-150.4-163.2-173.4c-168.3-40.8-517.6,158.1-617,351.8C6111.3-2459.2,6353.6-2290.9,6659.5-2446.5z"/><path d="M3013.6,2639.9c-135.1-45.9-221.8-130-221.8-219.3c0-160.6,181-270.3,453.8-270.3c418.1-2.5,599.2,288.1,280.5,453.8C3408.8,2665.4,3138.5,2683.3,3013.6,2639.9z"/><path d="M2213,1490.1c-168.3-48.5-295.8-132.6-334-226.9c-28-68.8-25.5-89.2,17.9-168.3c188.7-351.9,1106.5-288.1,1106.5,79c0,48.5-22.9,114.7-51,150.4C2837.7,1472.2,2450.2,1556.4,2213,1490.1z"/><path d="M5693.2,1492.6c-119.8-40.8-216.7-145.3-216.7-232c0-267.7,594-372.2,846.5-147.9c201.4,181-17.9,410.5-387.5,407.9C5843.6,1520.7,5736.6,1507.9,5693.2,1492.6z"/><path d="M8051.6,1278.5c-124.9-53.5-178.5-109.6-178.5-188.7c0-137.7,165.7-229.5,415.6-229.5c175.9,0,341.7,73.9,377.3,168.3C8737.4,1227.5,8321.8,1395.8,8051.6,1278.5z"/><path d="M4321.5,1013.3c-147.9-28-234.6-102-234.6-204c0-124.9,96.9-186.1,313.6-198.9c204-12.7,346.7,35.7,392.6,132.6C4869.7,916.4,4609.6,1064.3,4321.5,1013.3z"/><path d="M6662.1,225.5c-196.3-48.4-344.2-181-344.2-311c0-323.8,775.1-453.8,1070.8-175.9C7671.7,3.7,7161.8,345.3,6662.1,225.5z"/></g></g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/erase.svg

@@ -0,0 +1,14 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 16.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
+<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 width="548.157px" height="548.157px" viewBox="0 0 548.157 548.157" style="enable-background:new 0 0 548.157 548.157;"
+	 xml:space="preserve">
+<g>
+	<path d="M545.027,112.765c-3.046-6.471-7.57-11.657-13.565-15.555c-5.996-3.9-12.614-5.852-19.846-5.852H292.351
+		c-11.04,0-20.175,4.184-27.408,12.56L9.13,396.279c-4.758,5.328-7.661,11.56-8.708,18.698c-1.049,7.139-0.144,13.941,2.712,20.417
+		c3.044,6.468,7.564,11.652,13.561,15.553c5.997,3.898,12.612,5.853,19.845,5.853h219.268c11.042,0,20.177-4.179,27.41-12.56
+		l255.813-292.363c4.75-5.33,7.655-11.561,8.699-18.699C548.788,126.039,547.877,119.238,545.027,112.765z M255.811,420.254H36.54
+		l95.93-109.632h219.27L255.811,420.254z"/>
+</g>
+</svg>

web/lib/miniPaint-4.14.2/images/icons/fill.svg

@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="iso-8859-1"?>
+<!-- Generator: Adobe Illustrator 19.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
+<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 512 512" style="enable-background:new 0 0 512 512;" xml:space="preserve">
+<g>
+	<path d="M494.412,245.154L371.59,122.507l93.792-93.934L437.07,0.304L343.28,94.237l-56.613-56.532L37.714,286.659
+		l207.754,207.459c11.699,11.699,27.069,17.549,42.436,17.549c15.367,0,30.735-5.849,42.436-17.549l164.082-164.082
+		c11.389-11.388,17.632-26.53,17.578-42.635C511.946,271.446,505.704,256.446,494.412,245.154z M466.131,301.745L302.048,465.828
+		c-7.799,7.8-20.489,7.8-28.3-0.01L94.314,286.638L286.687,94.266l28.324,28.283L211.272,226.445l28.312,28.269L343.322,150.82
+		l122.811,122.636c3.761,3.761,5.842,8.761,5.859,14.079C472.009,292.902,469.928,297.948,466.131,301.745z"/>
+</g>
+<g>
+	<path d="M95.137,398.966c-10.179-15.198-20.245-27.482-20.669-27.997l-15.455-18.808l-15.455,18.808
+		c-0.424,0.516-10.49,12.799-20.669,27.997C2.372,429.596,0,444.293,0,452.684c0,32.54,26.474,59.012,59.012,59.012
+		c32.539,0,59.012-26.473,59.012-59.012C118.025,444.293,115.652,429.596,95.137,398.966z M59.012,471.688
+		c-10.479,0-19.004-8.525-19.005-18.956c0.004-0.08,0.526-8.271,16.291-31.757c0.907-1.35,1.813-2.678,2.714-3.972
+		c0.899,1.294,1.806,2.622,2.714,3.972c15.491,23.079,16.265,31.388,16.29,31.745C77.996,463.182,69.48,471.688,59.012,471.688z"/>
+</g>
+</svg>