增加ComparingTwoFrames、右键image-to-text

Update README.md
2024-05-14 09:59:39 +08:00 · 2024-05-14 09:57:51 +08:00 · 2024-05-14 09:55:34 +08:00 · 2024-05-14 09:54:27 +08:00 · 2024-05-14 09:54:24 +08:00 · 2024-05-14 09:39:57 +08:00
39 changed files with 5437 additions and 329 deletions
@@ -1,30 +1,50 @@
-> 适配了最新版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-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)
@@ -32,59 +52,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 反馈问题）
+> 如果遇到上传图片不成功，请检查下：局域网或者是云服务，请使用 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 、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 
+
+> 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) -->

@@ -98,61 +150,65 @@ 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

 #### LoadImagesToBatch
-> Upload multiple images for batch input into the IP adapter. 
+
+> 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. 
+
+> 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
+
 > 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) 
+
+> 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) 
-
+> StyleAligned , Modified from [style_aligned_comfy](https://github.com/brianfitzgerald/style_aligned_comfy)

 ### 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

 ![main](./assets/all-workflow.svg)
@@ -160,33 +216,27 @@ https://github.com/shadowcz007/comfyui-mixlab-nodes/assets/12645064/e7e77f90-e43

 [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
+**_ 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.
@@ -197,18 +247,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

@@ -224,40 +277,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
@@ -277,4 +325,3 @@ File / LoadImagesFromPath SaveImageToLocal LoadImagesFromURL
    src="https://api.star-history.com/svg?repos=shadowcz007/comfyui-mixlab-nodes&type=Date"
  />
 </picture>
-
@@ -7,9 +7,28 @@ import urllib
 import hashlib
 import datetime
 import folder_paths
-
+import logging
+from comfy.cli_args import args
 python = sys.executable

+#修复 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 server import PromptServer

@@ -42,7 +61,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:
@@ -263,7 +282,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 = [{
@@ -273,7 +292,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:
@@ -413,37 +433,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}")
@@ -454,7 +523,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

 # 创建路由表
@@ -549,13 +617,19 @@ async def nodes_map_hander(request):
 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)
-    
-    names = folder_paths.get_filename_list(t)

+    try:
+        if data['type']=='llamafile':
+            names=get_llama_models()
+    except:
+        print("llamafile none")
+    
    return web.json_response({"names":names,"types":list(folder_paths.folder_names_and_paths.keys())})


@@ -577,34 +651,143 @@ async def post_prompt_result(request):
    return web.json_response({"result":res})


+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']


-# 扩展api接口
-# from server import PromptServer
-# from aiohttp import web
+    chat_format="chatml"
+
+    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)
+
+    app = create_app(
+                server_settings=server_settings,
+                model_settings=[
+                    ModelSettings(
+                    model=model,
+                    model_alias=os.path.basename(model),
+                    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)
+    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/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 GLIGENTextBoxApply_Advanced,EmbeddingPrompt,RandomPrompt,PromptSlide,PromptSimplification,PromptImage,JoinWithDelimiter
-from .nodes.ImageNode import 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.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 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 MaskListReplace,MaskListMerge,OutlineMask,FeatheredMask
+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
+

 # 要导出的所有节点及其名称的字典
 # 注意：名称应全局唯一
@@ -651,6 +834,7 @@ NODE_CLASS_MAPPINGS = {
    # "VAELoaderConsistencyDecoder":VAELoader,
    "SaveImageToLocal":SaveImageToLocal,
    "SaveImageAndMetadata_":SaveImageAndMetadata,
+    "ComparingTwoFrames_":ComparingTwoFrames,
    # "VAEDecodeConsistencyDecoder":VAEDecode,
    "ScreenShare":ScreenShareNode,
    "FloatingVideo":FloatingVideo,
@@ -688,7 +872,11 @@ NODE_CLASS_MAPPINGS = {
    "MaskListReplace_":MaskListReplace,
    "ImageListReplace_":ImageListReplace,
    "VAEEncodeForInpaint_Frames":VAEEncodeForInpaint_Frames,
-    "IncrementingListNode_":IncrementingListNode
+    "IncrementingListNode_":IncrementingListNode,
+    "PreviewMask_":PreviewMask_,
+     "LoadTripoSRModel_": LoadTripoSRModel,
+    "TripoSRSampler_": TripoSRSampler,
+    "SaveTripoSRMesh": SaveTripoSRMesh
    # "GamePal":GamePal
 }

@@ -701,89 +889,96 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "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",
+    "PromptImage":"Output Prompt and Image ♾️Mixlab",
    "SplitLongMask":"Splitting a long image into sections",
    "VAELoaderConsistencyDecoder":"Consistency Decoder Loader",
    "VAEDecodeConsistencyDecoder":"Consistency Decoder Decode",
    "ScreenShare":"Screen Share ♾️Mixlab",
    "FloatingVideo":"FloatingVideo ♾️Mixlab",
-    "ChatGPTOpenAI":"ChatGPT ♾️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",
+    "CompositeImages_":"Composite Images ♾️Mixlab",
    "DynamicDelayProcessor":"DynamicDelayByText ♾️Mixlab",
    "LaMaInpainting":"LaMaInpainting ♾️Mixlab",
    "PromptSlide":"Prompt Slide ♾️Mixlab",
    "PromptGenerate_Mix":"Prompt Generate ♾️Mixlab",
    "ChinesePrompt_Mix":"Chinese Prompt ♾️Mixlab",
    "GamePal":"GamePal ♾️Mixlab",
-    "RembgNode_Mix":"Remove Background",
-    "LoraNames_":"LoraName",
-    "ApplyVisualStylePrompting_":"Apply VisualStyle Prompting",
-    "StyleAlignedReferenceSampler_": "StyleAligned Reference Sampler",
-    "StyleAlignedSampleReferenceLatents_": "StyleAligned Sample Reference Latents",
-    "StyleAlignedBatchAlign_": "StyleAligned Batch Align",
-    "LoadVideoAndSegment_":"Load Video And Segment",
-    "VideoCombine_Adv":"Video Combine",
-    "MaskListMerge_":"MaskList to Mask",
-    "ListSplit_":"Split List",
-    "MaskListReplace_":"MaskList Replace",
-    "ImageListReplace_":"ImageList Replace",
-    "SwitchByIndex":"List Switch By Index",
+    "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",
-    "GridInput":"Grid Input",
-    "GridOutput":"Grid Output",
-    "GetImageSize_":"Get Image Size",
-    "VAEEncodeForInpaint_Frames":"VAE Encode For Inpaint Frames",
-    "IncrementingListNode_":"Create Incrementing Number List",
-    "LoadImagesToBatch":"Load Images To Batch"
+    "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')
@@ -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 (
@@ -4,8 +4,17 @@ import urllib.error
 import re,json,os,string,random
 import folder_paths
 import hashlib
-import codecs
-from zhipuai import ZhipuAI
+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())
@@ -48,13 +57,110 @@ def openai_client(key,url):
    base_url=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 ):

@@ -62,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
@@ -74,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)
@@ -97,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}),
@@ -107,16 +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-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"], 
-                          {"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}),
            },
@@ -139,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()
@@ -162,6 +282,9 @@ class ChatGPTNode:
            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')
@@ -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):
@@ -16,23 +16,135 @@ import math,glob
 from .Watcher import FolderWatcher
 import hashlib

-def composite_images(foreground, background, mask):
+
+
+# 将PIL图片转换为OpenCV格式
+def pil_to_opencv(image):
+    open_cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    return open_cv_image
+
+# 将OpenCV格式图片转换为PIL格式
+def opencv_to_pil(image):
+    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+    return pil_image
+
+
+def composite_images(foreground, background, mask,is_multiply_blend=False,position="overall"):
    width,height=foreground.size
 
    bg_image=background
+
+    bwidth,bheight=bg_image.size
+
    # 按z-index排序
-    layer = {
-            "x":0,
-            "y":0,
-            "width":width,
-            "height":height,
-            "z_index":88,
-            "scale_option":'overall',
-            "image":foreground,
-            "mask":mask
-    }
+    if position=="overall":
+        layer = {
+                "x":0,
+                "y":0,
+                "width":bwidth,
+                "height":bheight,
+                "z_index":88,
+                "scale_option":'overall',
+                "image":foreground,
+                "mask":mask
+        }
+
+    elif position=='center_bottom':
+        
+        scale = int(bwidth*0.25) / width
+        new_height = int(height * scale)
+
+        layer = {
+                "x":int(bwidth*0.75*0.5),
+                "y":bheight-new_height-24,
+                "width":int(bwidth*0.25),
+                "height":int(bheight*0.25),
+                "z_index":88,
+                "scale_option":'width',
+                "image":foreground,
+                "mask":mask
+        }
+
+    elif position=='right_bottom':
+        
+        scale = int(bwidth*0.25) / width
+        new_height = int(height * scale)
+
+        layer = {
+                "x":bwidth-int(bwidth*0.25)-24,
+                "y":bheight-new_height-24,
+                "width":int(bwidth*0.25),
+                "height":int(bheight*0.25),
+                "z_index":88,
+                "scale_option":'width',
+                "image":foreground,
+                "mask":mask
+        }
+
+
+    elif position=='center_top':
+        
+        scale = int(bwidth*0.25) / width
+        new_height = int(height * scale)
+
+        layer = {
+                "x":int( bwidth*0.75*0.5),
+                "y":24,
+                "width":int(bwidth*0.25),
+                "height":int(bheight*0.25),
+                "z_index":88,
+                "scale_option":'width',
+                "image":foreground,
+                "mask":mask
+        }
+
+    elif position=='right_top':
+        
+        scale = int(bwidth*0.25) / width
+        new_height = int(height * scale)
+
+        layer = {
+                "x":bwidth-int(bwidth*0.25)-24,
+                "y":24,
+                "width":int(bwidth*0.25),
+                "height":int(bheight*0.25),
+                "z_index":88,
+                "scale_option":'width',
+                "image":foreground,
+                "mask":mask
+        }
+    elif position=='left_top':
+        
+        scale = int(bwidth*0.25) / width
+        new_height = int(height * scale)
+
+        layer = {
+                "x":24,
+                "y":24,
+                "width":int(bwidth*0.25),
+                "height":int(bheight*0.25),
+                "z_index":88,
+                "scale_option":'width',
+                "image":foreground,
+                "mask":mask
+        }
+    elif position=='left_bottom':
+        
+        scale = int(bwidth*0.25) / width
+        new_height = int(height * scale)
+
+        layer = {
+                "x":24,
+                "y":bheight-new_height-24,
+                "width":int(bwidth*0.25),
+                "height":int(bheight*0.25),
+                "z_index":88,
+                "scale_option":'width',
+                "image":foreground,
+                "mask":mask
+        }
                
-    width, height = bg_image.size
+    # width, height = bg_image.size

    layer_image=layer['image']
    layer_mask=layer['mask']
@@ -40,12 +152,12 @@ def composite_images(foreground, background, mask):
    bg_image=merge_images(bg_image,
                        layer_image,
                        layer_mask,
-                                        layer['x'],
-                                        layer['y'],
-                                        layer['width'],
-                                        layer['height'],
-                                        layer['scale_option']
-                                        )
+                        layer['x'],
+                        layer['y'],
+                        layer['width'],
+                        layer['height'],
+                        layer['scale_option'],
+                        is_multiply_blend )
                   
    bg_image=bg_image.convert('RGB')
    
@@ -620,7 +732,7 @@ def detect_faces(image):

 def areaToMask(x,y,w,h,image):
    # 创建一个与原图片大小相同的空白图片
-    mask = Image.new('1', image.size)
+    mask = Image.new('L', image.size)

    # 创建一个可用于绘制的对象
    draw = ImageDraw.Draw(mask)
@@ -653,7 +765,45 @@ def areaToMask(x,y,w,h,image):
 #     # bg_image.save("output.jpg")
 #     return bg_image

-def merge_images(bg_image, layer_image, mask, x, y, width, height, scale_option):
+
+import cv2
+import numpy as np
+
+# ps的正片叠底
+# 可以基于https://www.cnblogs.com/jsxyhelu/p/16947810.html ，用gpt写python代码
+def multiply_blend(image1, image2):
+    image1=pil_to_opencv(image1)
+    image2=pil_to_opencv(image2)
+    # 将图像转换为浮点型
+    image1 = image1.astype(float)
+    image2 = image2.astype(float)
+    if image1.shape != image2.shape:
+        image1 = cv2.resize(image1, (image2.shape[1], image2.shape[0]))
+        
+    # 归一化图像
+    image1 /= 255.0
+    image2 /= 255.0
+
+    # 正片叠底混合
+    blended = image1 * image2
+
+    # 将图像还原为8位无符号整数
+    blended = (blended * 255).astype(np.uint8)
+
+    blended=opencv_to_pil(blended)
+    return blended
+
+# # 读取图像
+# image1 = cv2.imread('1.png')
+# image2 = cv2.imread('3.png')
+
+# # 进行正片叠底混合
+# result = multiply_blend(image1, image2)
+
+# cv2.imwrite('result.jpg', result)
+
+
+def merge_images(bg_image, layer_image, mask, x, y, width, height, scale_option,is_multiply_blend=False):
    # 打开底图
    bg_image = bg_image.convert("RGBA")

@@ -678,12 +828,55 @@ def merge_images(bg_image, layer_image, mask, x, y, width, height, scale_option)
        # 整体缩放
        layer_image = layer_image.resize((width, height))

+    elif scale_option == "longest":
+        original_width, original_height = layer_image.size
+        if original_width > original_height:
+            new_width=width
+            scale = width / original_width
+            new_height = int(original_height * scale)
+            x=0
+            y=int((height-new_height)*0.5)
+        else:
+            new_height=height
+            scale = height / original_height
+            new_width = int(original_height * scale)
+            x=int((width-new_width)*0.5)
+            y=0
+    # elif side == "shortest":
+    #         if width < height:
+    #              
+    #         else:
+    #              
+    
+
    # 调整mask的大小
    nw, nh = layer_image.size
    mask = mask.resize((nw, nh))

-    # 在底图上粘贴图层
-    bg_image.paste(layer_image, (x, y), mask=mask)
+    # # 分离出a通道
+    # r, g, b, alpha = layer_image.split()
+    # alpha = ImageOps.invert(alpha)
+    # # 创建一个新的RGB图像
+    # new_rgb_image = Image.new("RGB", layer_image.size)
+    # # 将透明通道粘贴到新的RGB图像上
+    # new_rgb_image.paste(layer_image, (0, 0), mask=alpha)
+   
+    # new_rgb_image.paste(layer_image, (x, y), mask=mask)
+    # mask=new_rgb_image.convert('L')
+    # mask = ImageOps.invert(mask)
+
+    if is_multiply_blend:
+        bg_image_white=Image.new("RGB", bg_image.size,(255, 255, 255))
+
+        bg_image_white.paste(layer_image, (x, y), mask=mask)
+        bg_image=multiply_blend(bg_image_white,bg_image)
+        bg_image=bg_image.convert("RGBA")
+    else:
+        transparent_img = Image.new("RGBA",layer_image.size, (255, 255, 255, 0))
+        transparent_img.paste(layer_image,(0, 0), mask)
+        # transparent_img.save('test.png')
+        bg_image.paste(transparent_img, (x, y), transparent_img)
+

    # 输出合成后的图片
    return bg_image
@@ -731,6 +924,31 @@ def resize_image(layer_image, scale_option, width, height,color="white"):
        resized_image = resized_image.convert("RGB")
        resized_image=resize_2(resized_image)
        return resized_image
+    elif scale_option == "longest":
+    #暂时不用，  
+        if original_width > original_height:
+            new_width=width
+            scale = width / original_width
+            new_height = int(original_height * scale)
+            x=0
+            y=int((new_height-height)*0.5)
+            resized_image = Image.new("RGB", (new_width, new_height), color=color)
+            resized_image.paste(layer_image.resize((new_width, new_height)), (x,y))
+            resized_image = resized_image.convert("RGB")
+            resized_image=resize_2(resized_image)
+            return resized_image
+        else:
+            new_height=height
+            scale = height / original_height
+            new_width = int(original_height * scale)
+            x=int((new_width-width)*0.5)
+            y=0
+            resized_image = Image.new("RGB", (new_width, new_height), color=color)
+            resized_image.paste(layer_image.resize((new_width, new_height)), (x,y))
+            resized_image = resized_image.convert("RGB")
+            resized_image=resize_2(resized_image)
+            return resized_image
+
    
    layer_image=resize_2(layer_image)
    return layer_image
@@ -1632,10 +1850,15 @@ class CompositeImages:
    def INPUT_TYPES(s):
        return {
            "required": { 
-            "foreground": ("IMAGE",),
-            "mask":("MASK",),
-            "background": ("IMAGE",),
-            },
+                    "foreground": (any_type,),
+                    "mask":("MASK",),
+                    "background": ("IMAGE",),
+                    },
+             "optional":{
+                           
+                  "is_multiply_blend":  ("BOOLEAN", {"default": False}),
+                  "position":  (['overall',"center_bottom","center_top","right_bottom","left_bottom","right_top","left_top"],),    
+                    }
                }
    
    RETURN_TYPES = ("IMAGE",)
@@ -1647,11 +1870,11 @@ class CompositeImages:

    # OUTPUT_IS_LIST = (True,)

-    def run(self, foreground,mask,background):
+    def run(self, foreground,mask,background,is_multiply_blend,position):
        foreground= tensor2pil(foreground)
        mask= tensor2pil(mask)
        background= tensor2pil(background)
-        res=composite_images(foreground,background,mask)
+        res=composite_images(foreground,background,mask,is_multiply_blend,position)
        
        return (pil2tensor(res),)

@@ -1753,7 +1976,7 @@ class NewLayer:
                    "display": "number" # Cosmetic only: display as "number" or "slider"
                }),
                "scale_option": (["width","height",'overall'],),
-                "image": ("IMAGE",),
+                "image": (any_type,),
            },
             "optional":{
                    "mask": ("MASK",{"default": None}),
@@ -1810,21 +2033,20 @@ def createMask(image,x,y,w,h):
    # mask.save("mask.png")
    return mask

-
 def splitImage(image, num):
    width, height = image.size

    num_rows = int(num ** 0.5)
    num_cols = int(num / num_rows)
    
-    grid_width = width // num_cols
-    grid_height = height // num_rows
+    grid_width = int(width // num_cols)
+    grid_height = int(height // num_rows)

    grid_coordinates = []
    for i in range(num_rows):
        for j in range(num_cols):
-            x = j * grid_width
-            y = i * grid_height
+            x = int(j * grid_width)
+            y = int(i * grid_height)
            grid_coordinates.append((x, y, grid_width, grid_height))

    return grid_coordinates
@@ -2156,12 +2378,16 @@ class GridOutput:
    def INPUT_TYPES(s):
        return { 
            "required": {
-                "grid": ("_GRID",)
-            }
+                "grid": ("_GRID",),
+                
+            },
+            "optional":{
+                  "bg_image":("IMAGE",)
+                }
                }
    
-    RETURN_TYPES = ("INT","INT","INT","INT",)
-    RETURN_NAMES = ("x","y","width","height",)
+    RETURN_TYPES = ("INT","INT","INT","INT","MASK",)
+    RETURN_NAMES = ("x","y","width","height","mask",)

    FUNCTION = "run"

@@ -2170,9 +2396,26 @@ class GridOutput:
    INPUT_IS_LIST = False
    # OUTPUT_IS_LIST = (True,)

-    def run(self,grid):
+    def run(self,grid,bg_image=None):
        x,y,w,h=grid
-        return (x,y,w,h,)
+        x=int(x)
+        y=int(y)
+        w=int(w)
+        h=int(h)
+
+        masks=[]
+        if bg_image!=None:
+            for i in range(len(bg_image)):
+                im=bg_image[i]
+                #增加输出mask
+                im=tensor2pil(im)
+                mask=areaToMask(x,y,w,h,im)
+                mask=pil2tensor(mask)
+                masks.append(mask)
+        out=None
+        if len(masks)>0:
+            out = torch.cat(masks, dim=0)
+        return (x,y,w,h,out,)



@@ -2267,10 +2510,14 @@ class MergeLayers:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { 
-            "layers": ("LAYER",),
-            "images": ("IMAGE",),
-                             },
-
+                            "layers": ("LAYER",),
+                            "images": ("IMAGE",),
+                                            },
+            "optional":{
+                           
+                            "is_multiply_blend":  ("BOOLEAN", {"default": False}),
+                            
+                    }
                }
    
    RETURN_TYPES = ("IMAGE","MASK",)
@@ -2283,11 +2530,12 @@ class MergeLayers:
    INPUT_IS_LIST = True
    # OUTPUT_IS_LIST = (False,)

-    def run(self,layers,images):
+    def run(self,layers,images,is_multiply_blend):

        bg_images=[]
        masks=[]
- 
+        
+        is_multiply_blend=is_multiply_blend[0]
        # print(len(images),images[0].shape)
        # 1  torch.Size([2, 512, 512, 3])
        # 4  torch.Size([1, 1024, 768, 3])
@@ -2318,6 +2566,8 @@ class MergeLayers:
                    
                    layer_image=tensor2pil(image)
                    layer_mask=tensor2pil(mask)
+                    # t=layer_image.convert("RGBA")
+                    # t.save('test.png') 如果layerimage传入的是rgba，则是透明的
                    bg_image=merge_images(bg_image,
                                        layer_image,
                                        layer_mask,
@@ -2325,7 +2575,8 @@ class MergeLayers:
                                        layer['y'],
                                        layer['width'],
                                        layer['height'],
-                                        layer['scale_option']
+                                        layer['scale_option'],
+                                        is_multiply_blend
                                        )
                    
                    final_mask=merge_images(final_mask,
@@ -2732,6 +2983,65 @@ class SaveImageAndMetadata:

        return { "ui": { "images": results } }

+class ComparingTwoFrames:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = "ComparingTwoFrames"
+        self.compress_level = 4
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": 
+                    {"before_image": ("IMAGE", ),
+                    "after_image": ("IMAGE", )
+                     }, 
+                }
+
+    RETURN_TYPES = ()
+    FUNCTION = "comparingImages"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "♾️Mixlab/Output"
+
+    def comparingImages(self, before_image,after_image):
+        filename_prefix = self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
+            filename_prefix, self.output_dir, after_image[0].shape[1], after_image[0].shape[0])
+        
+        bresults = list()
+        
+        for bimage in before_image:
+            i = 255. * bimage.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+         
+            file = f"{filename}_{counter:05}_.png"
+            img.save(os.path.join(full_output_folder, file), pnginfo=None, compress_level=self.compress_level)
+            bresults.append({
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type
+            })
+            counter += 1
+
+
+        results = list()
+        for aimage in after_image:
+            i = 255. * aimage.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+         
+            file = f"{filename}_{counter:05}_.png"
+            img.save(os.path.join(full_output_folder, file), pnginfo=None, compress_level=self.compress_level)
+            results.append({
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type
+            })
+            counter += 1
+
+        return { "ui": { "after_images": results,"before_images":bresults } }
+
 class ImageColorTransfer:
    @classmethod
    def INPUT_TYPES(s):
@@ -2752,35 +3062,45 @@ class ImageColorTransfer:
    CATEGORY = "♾️Mixlab/Color"

    # 输入是否为列表
-    INPUT_IS_LIST = True
+    # INPUT_IS_LIST = True

    # 输出是否为列表
-    OUTPUT_IS_LIST = (True,)
+    # OUTPUT_IS_LIST = (True,)

    def run(self,source,target,weight):

        res=[]

-        target=target[0][0]
-        # print(target.shape)
-        target=tensor2pil(target)
+        #batch-list
+        source_list = [source[i:i + 1, ...] for i in range(source.shape[0])]
+        target_list = [target[i:i + 1, ...] for i in range(target.shape[0])]

-        for ims in source:
-            for im in ims:
-                image=tensor2pil(im)
-                image_res=color_transfer(image,target)
+        # 长度纠正为相等
+        if len(target_list) != len(source_list):
+            target_list = target_list * (len(source_list) // len(target_list)) + target_list[:len(source_list) % len(target_list)]
+        
+        for i in range(len(source_list)):
+            target=target_list[i]
+            source=source_list[i]
+            target=tensor2pil(target)

-                # weight Blend image # contributors:@ning
-                blend_mask = Image.new(mode="L", size=image.size,
-                                    color=(round(weight[0] * 255)))
-                blend_mask = ImageOps.invert(blend_mask)
-                img_result = Image.composite(image, image_res, blend_mask)
+            image=tensor2pil(source)

-                del image, image_res, blend_mask
+            image_res=color_transfer(image,target)

-                img_result=pil2tensor(img_result)
-                res.append(img_result)
+            # weight Blend image # contributors:@ning
+            blend_mask = Image.new(mode="L", size=image.size,
+                                    color=(round(weight * 255)))
+            blend_mask = ImageOps.invert(blend_mask)
+            img_result = Image.composite(image, image_res, blend_mask)
+            del image, image_res, blend_mask
+            
+            img_result=pil2tensor(img_result)

+            res.append(img_result)
+
+        # list - batch
+        res=torch.cat(res, dim=0)

        return (res,)

@@ -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")
@@ -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):
@@ -71,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
@@ -18,19 +18,25 @@ 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)
@@ -0,0 +1,173 @@
+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_filename_list, 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
+
+triposr_model_path=path.join(models_dir,'triposr/model.ckpt')
+
+
+# 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}}
+
+
+
@@ -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
@@ -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
@@ -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()
@@ -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
@@ -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
@@ -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,
+        )
@@ -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
@@ -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)
@@ -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
@@ -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
@@ -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
@@ -7,6 +7,12 @@ openai
 simple-lama-inpainting
 clip-interrogator==0.6.0
 transformers>=4.36.0
-zhipuai
 lark-parser
-imageio-ffmpeg
+imageio-ffmpeg
+rembg[gpu]
+omegaconf==2.3.0
+Pillow>=9.5.0
+einops==0.7.0
+trimesh>=4.0.5
+huggingface-hub
+scikit-image
@@ -247,10 +247,10 @@

        .card textarea {
            width: 100%;
-            /* height: 200px; */
+            height:'fit-content';
            /* min-width: 300px; */
            margin-top: 12px;
-            resize: vertical;
+            resize: none;
            overflow: hidden;
        }

@@ -278,8 +278,11 @@
        }

        button:hover {
-            border-color: yellow !important;
-            color: yellow !important;
+            /* border-color: yellow !important; */
+            color: #ffffff !important;
+            /* font-weight: 400; */
+            background: #232222;
+            cursor: pointer;
        }

        .disabled {
@@ -400,6 +403,50 @@
        .hidden-caption-content {
            display: none;
        }
+
+        /* 定义滚动条轨道的背景颜色 */
+
+        ::-webkit-scrollbar-track {
+            background-color: #f1f1f1;
+            /* 轨道背景颜色 */
+        }
+
+        /* 定义滚动条滑块的颜色 */
+
+        ::-webkit-scrollbar-thumb {
+            background-color: #888;
+            /* 滑块颜色 */
+            border-radius: 4px;
+            /* 滑块圆角 */
+        }
+
+        /* 鼠标悬停时滚动条滑块的颜色 */
+
+        ::-webkit-scrollbar-thumb:hover {
+            background-color: #555;
+            /* 悬停时滑块颜色 */
+        }
+
+        /* 定义滚动条角落的颜色 */
+
+        ::-webkit-scrollbar-corner {
+            background-color: transparent;
+            /* 角落颜色 */
+        }
+
+/* 
+        .dynamic_prompt::after {
+            content: attr(title);
+            position: absolute;
+            color: black;
+            padding: 4px;
+            padding-left: 25px;
+            border-radius: 4px;
+        } */
+
+        summary {
+            user-select: none;
+        }
    </style>
    <!-- <script src="../../../scripts/api.js" type="module"></script> -->
    <link href="/extensions/comfyui-mixlab-nodes/lib/photoswipe.min.css" rel="stylesheet">
@@ -440,8 +487,39 @@
    <script type="module">
        import PhotoSwipeLightbox from '/extensions/comfyui-mixlab-nodes/lib/photoswipe-lightbox.esm.min.js'
        // console.log(Lightbox)
-
        import { api } from "../../../scripts/api.js";
+
+
+        // ComfyUI\web\extensions\core\dynamicPrompts.js
+        // 官方实现修改
+        // Allows for simple dynamic prompt replacement
+        // Inputs in the format {a|b} will have a random value of a or b chosen when the prompt is queued.
+
+        /*
+         * Strips C-style line and block comments from a string
+         */
+        function stripComments(str) {
+            return str.replace(/\/\*[\s\S]*?\*\/|\/\/.*/g, '');
+        }
+
+        function dynamicPrompts(prompt) {
+            prompt = stripComments(prompt);
+            while (prompt.replace("\\{", "").includes("{") && prompt.replace("\\}", "").includes("}")) {
+                const startIndex = prompt.replace("\\{", "00").indexOf("{");
+                const endIndex = prompt.replace("\\}", "00").indexOf("}");
+
+                const optionsString = prompt.substring(startIndex + 1, endIndex);
+                const options = optionsString.split("|");
+
+                const randomIndex = Math.floor(Math.random() * options.length);
+                const randomOption = options[randomIndex];
+
+                prompt = prompt.substring(0, startIndex) + randomOption + prompt.substring(endIndex + 1);
+            }
+            return prompt
+        }
+
+
        // console.log('api', api)
        const base64Df =
            'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAwAAAAMCAYAAABWdVznAAAAAXNSR0IArs4c6QAAALZJREFUKFOFkLERwjAQBPdbgBkInECGaMLUQDsE0AkRVRAYWqAByxldPPOWHwnw4OBGye1p50UDSoA+W2ABLPN7i+C5dyC6R/uiAUXRQCs0bXoNIu4QPQzAxDKxHoALOrZcqtiyR/T6CXw7+3IGHhkYcy6BOR2izwT8LptG8rbMiCRAUb+CQ6WzQVb0SNOi5Z2/nX35DRyb/ENazhpWKoGwrpD6nICp5c2qogc4of+c7QcrhgF4Aa/aoAFHiL+RAAAAAElFTkSuQmCC'
@@ -771,6 +849,21 @@
            // 随机seed
            promptWorkflow = randomSeed(seed, promptWorkflow);

+            // //动态提示,改为输入的时候，手动触发
+            // for (const id in promptWorkflow) {
+            //     let node = promptWorkflow[id]
+            //     if (["TextInput_", "CLIPTextEncode", "PromptSimplification", "ChinesePrompt_Mix"].includes(
+            //         node.class_type
+            //     )) {
+            //         if (node.class_type == "PromptSimplification") {
+            //             promptWorkflow[id].inputs.prompt = dynamicPrompts(node.inputs.prompt);
+            //         } else {
+            //             promptWorkflow[id].inputs.text = dynamicPrompts(node.inputs.text);
+            //         }
+            //         console.log('#动态提示', promptWorkflow[id].inputs)
+            //     }
+            // }
+
            let url = get_url()
            const data = JSON.stringify({ prompt: promptWorkflow, client_id });
            fetch(`${url}/prompt`, {
@@ -859,6 +952,10 @@


        function createOutputs(outputData, link) {
+            const url = new URL(window.location.href);
+            const params = new URLSearchParams(url.search);
+            const innerApp = params.get("innerApp");
+
            const container = document.createElement('div');
            container.className = "output";

@@ -867,11 +964,11 @@

            const copyHTML = document.createElement('button');
            copyHTML.innerText = 'copy as html'
-            action.appendChild(copyHTML)
+            if (!innerApp) action.appendChild(copyHTML)

            const copyImage = document.createElement('button');
            copyImage.innerText = 'copy image'
-            action.appendChild(copyImage)
+            if (!innerApp) action.appendChild(copyImage)
            copyImage.style.marginLeft = '18px';

            let isURL = false;
@@ -971,18 +1068,21 @@
                    "PreviewImage",
                    "PromptImage",
                    "Image Save",
-                    "SaveImageAndMetadata_"].includes(node.class_type)) {
+                    "SaveImageAndMetadata_",
+                    "TransparentImage"].includes(node.class_type)) {

+                    const url = node.options?.defaultImage || window._appData?.icon || base64Df;
                    let a = document.createElement('a');
                    a.id = `output_${node.id}`
                    a.setAttribute('data-pswp-width', "200");
                    a.setAttribute('data-pswp-height', "200");
                    a.setAttribute('target', "_blank");
-                    a.setAttribute('href', base64Df);
+                    a.setAttribute('href', url);
+                    a.setAttribute('title', node.title);

                    let img = new Image();
                    // img;
-                    img.src = base64Df;
+                    img.src = url;
                    a.appendChild(img)
                    output_card.appendChild(a);
                    isShowImageFn = true;
@@ -1532,7 +1632,7 @@
                        textInput.value = data.inputs.prompt;
                    } else {
                        textInput.value = data.inputs.text;
-                    }
+                    };

                    // uploadImageInput.type = "text";
                    let json = localStorage.getItem(`t_${data.id}`)
@@ -1548,7 +1648,6 @@
                            window._appData.data[data.id].inputs.text = textInput.value;
                        }

-
                    } catch (error) {

                    }
@@ -1556,6 +1655,27 @@
                    uploadContainer.appendChild(textInput);
                    // autoResize(textInput);

+
+                    //动态提示功能
+                    const dynamicPromptsBtn = document.createElement('button');
+                    dynamicPromptsBtn.className = "dynamic_prompt"
+                    dynamicPromptsBtn.innerText = 'dynamic';
+                    dynamicPromptsBtn.style.width = '88px'
+                    uploadContainer.appendChild(dynamicPromptsBtn);
+                    dynamicPromptsBtn.addEventListener('click', e => {
+                        e.preventDefault();
+                        e.stopPropagation();
+                        let prompt = dynamicPrompts(textInput.value)
+                        textInput.setAttribute('title', prompt)
+                        dynamicPromptsBtn.setAttribute('title', prompt)
+                        if (data.class_type == "PromptSimplification") {
+                            window._appData.data[data.id].inputs.prompt = prompt;
+                        } else {
+                            window._appData.data[data.id].inputs.text = prompt;
+                        }
+                    })
+
+
                    function autoResize(textarea) {
                        textarea.style.height = 'auto';
                        textarea.style.height = textarea.scrollHeight + 'px';
@@ -2156,7 +2276,8 @@
            // leftDiv.appendChild(des);
            leftDiv.appendChild(statusDiv);
            leftDiv.appendChild(input1);
-            mainDiv.appendChild(submitButton);
+
+            if (typeof (data.data) == 'object') mainDiv.appendChild(submitButton);

            rightDiv.appendChild(output);

@@ -2513,7 +2634,7 @@
            ui.title.update(appData.name || 'Mixlab APP');

            // 更新应用图标
-            ui.icon.update(appData.icon || base64Df);
+            ui.icon.update(appData.icon || appData.output[0]?.options?.defaultImage || base64Df);

            ui.des.update(appData.description || '-');

@@ -2,6 +2,24 @@ 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'

@@ -189,9 +207,27 @@ async function extractInputAndOutputData (
        // 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 (
@@ -267,7 +303,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], //用于分享的功能扩展
@@ -288,6 +327,13 @@ async function save (json, download = false, showInfo = true) {
  try {
    let data = await app.graphToPrompt()

+    //从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,
@@ -350,11 +396,11 @@ async function save (json, download = false, showInfo = true) {

 function getInputsAndOutputs () {
  const inputs =
-      `LoadImage ImagesPrompt_ 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,Image Save,SaveImageAndMetadata_`.split(
+      `SaveTripoSRMesh,PreviewImage,SaveImage,TransparentImage,ShowTextForGPT,VHS_VideoCombine,VideoCombine_Adv,Image Save,SaveImageAndMetadata_,ClipInterrogator`.split(
        ','
      )

@@ -379,6 +425,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
@@ -0,0 +1,106 @@
+async function* completion (url, messages, controller) {
+  let data = {
+    model: 'gpt-3.5-turbo-16k',
+    messages,
+    temperature: 0.6,
+    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)
+  }
+}
@@ -3,7 +3,7 @@ import { app } from '../../../scripts/app.js'
 const repoOwner = 'shadowcz007' // 替换为仓库的所有者
 const repoName = 'comfyui-mixlab-nodes' // 替换为仓库的名称

-const version = 'v0.21.0'
+const version = 'v0.26.0'

 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;
@@ -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();
@@ -673,7 +705,9 @@ app.registerExtension({

  async beforeRegisterNodeDef (nodeType, nodeData, app) {
    if (nodeType.comfyClass == 'LoadImagesToBatch') {
+
      const orig_nodeCreated = nodeType.prototype.onNodeCreated
+
      nodeType.prototype.onNodeCreated = function () {
        orig_nodeCreated?.apply(this, arguments)

@@ -712,7 +746,9 @@ app.registerExtension({
          const file = e.target.files[0]
          const reader = new FileReader()
          reader.onload = async event => {
-            const base64 = event.target.result
+            let base64 = event.target.result
+            //压缩图片，控制1024以内
+            base64 = await loadImageToCanvas(base64)
            // console.log(base64)
            if (!imagesWidget.value) imagesWidget.value = { base64: [] }
            imagesWidget.value.base64.push(base64)
@@ -765,6 +801,33 @@ app.registerExtension({
        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') {
@@ -780,3 +843,150 @@ app.registerExtension({
    }
  }
 })
+
+// 如何引入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
+      //     }
+      //   )
+      // }
+    }
+  }
+})
@@ -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)
+    //     }
+    //   }
+    // }
+  }
+})
@@ -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)
    }
@@ -9,6 +9,148 @@ import {

 import { smart_init, addSmartMenu } from './smart_connect.js'

+import { completion_ } from './chat.js'
+
+function showTextByLanguage (key, json) {
+  // 获取浏览器语言
+  var language = navigator.language
+  // 判断是否为中文
+  if (
+    language.indexOf('zh') !== -1 ||
+    (language.indexOf('cn') !== -1 && json[key])
+  ) {
+    return json[key]
+  } else {
+    return key
+  }
+}
+
+//系统prompt
+const systemPrompt = `You are a prompt creator, your task is to create prompts for the user input request, the prompts are image descriptions that include keywords for (an adjective, type of image, framing/composition, subject, subject appearance/action, environment, lighting situation, details of the shoot/illustration, visuals aesthetics and artists), brake keywords by comas, provide high quality, non-verboose, coherent, brief, concise, and not superfluous prompts, the subject from the input request must be included verbatim on the prompt,the prompt is english`
+
+if (!localStorage.getItem('_mixlab_system_prompt')) {
+  localStorage.setItem('_mixlab_system_prompt', systemPrompt)
+}
+
+// 获取llama 模型
+async function get_llamafile_models () {
+  try {
+    const response = await fetch('/mixlab/folder_paths', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json'
+      },
+      body: JSON.stringify({
+        type: 'llamafile'
+      })
+    })
+
+    const data = await response.json()
+    console.log(data)
+    return data.names
+  } catch (error) {
+    console.error(error)
+  }
+}
+// 运行llama
+async function start_llama (model = 'Phi-3-mini-4k-instruct-Q5_K_S.gguf') {
+  let n_gpu_layers = -1
+  try {
+    n_gpu_layers = parseInt(localStorage.getItem('_mixlab_llama_n_gpu'))
+  } catch (error) {}
+
+  try {
+    const response = await fetch('/mixlab/start_llama', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json'
+      },
+      body: JSON.stringify({
+        model,
+        n_gpu_layers
+      })
+    })
+
+    const data = await response.json()
+    if (data.llama_cpp_error) {
+      return
+    }
+
+    return {
+      url: `http://${window.location.hostname}:${data.port}`,
+      model: data.model,
+      chat_format: data.chat_format
+    }
+  } catch (error) {
+    console.error(error)
+  }
+}
+
+function resizeImage (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 = 768
+      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 createMenu () {
+  const menu = document.querySelector('.comfy-menu')
+  const separator = document.createElement('div')
+  separator.style = `margin: 20px 0px;
+  width: 100%;
+  height: 1px;
+  background: var(--border-color);
+  `
+  menu.append(separator)
+
+  if (!menu.querySelector('#mixlab_chatbot_by_llamacpp')) {
+    const appsButton = document.createElement('button')
+    appsButton.id = 'mixlab_chatbot_by_llamacpp'
+    appsButton.textContent = '♾️Mixlab'
+
+    // appsButton.onclick = () =>
+    appsButton.onclick = async () => {
+      if (window._mixlab_llamacpp) {
+        //显示运行的模型
+        createModelsModal([
+          window._mixlab_llamacpp.url,
+          window._mixlab_llamacpp.model
+        ])
+      } else {
+        let ms = await get_llamafile_models()
+        ms = ms.filter(m => !m.match('-mmproj-'))
+        if (ms.length > 0) createModelsModal(ms)
+      }
+    }
+    menu.append(appsButton)
+  }
+}
+
 let isScriptLoaded = {}

 function loadExternalScript (url) {
@@ -300,14 +442,16 @@ async function get_my_app (filename = null, category = '') {

    for (const res of result.data) {
      let { app, workflow } = res.data
-      if (app.filename)
+      if (app?.filename)
        data.push({
          ...app,
          data: workflow,
          date: res.date
        })
    }
-  } catch (error) {}
+  } catch (error) {
+    console.log(error)
+  }
  return data
 }

@@ -343,6 +487,33 @@ injectCSS(`::-webkit-scrollbar {
  width: 2px;
 }

+#mixlab_chatbot_by_llamacpp{
+  font-size:14px
+}
+
+#mixlab_chatbot_by_llamacpp::before {
+  content: attr(title);
+  position: absolute;
+  margin-top: 24px;
+  font-size: 10px;
+}
+
+.mix_tag{
+  padding:8px;cursor: pointer;font-size: 14px;
+    color: var(--input-text);
+    background-color: var(--comfy-input-bg);
+    border-radius: 8px;
+    border-color: var(--border-color);
+    border-style: solid;
+    margin-top: 2px;
+    margin-bottom: 14px;
+}
+
+.mix_tag:hover{
+  background-color: #101c19;
+  color: aquamarine;
+}
+
@keyframes loading_mixlab {
  0% {
    background-color: green;
@@ -493,13 +664,28 @@ app.showMissingNodesError = async function (
  // console.log('#nodesMap', nodesMap)
  // console.log('###MIXLAB', missingNodeTypes, hasAddedNodes)
  this.ui.dialog.show(
-    `When loading the graph, the following node types were not found: <ul>${missingNodeGithub(
-      missingNodeTypes,
-      nodesMap
-    ).join('')}</ul>${
-      hasAddedNodes
-        ? 'Nodes that have failed to load will show as red on the graph.'
-        : ''
+    `<a style="color: white;
+    font-size: 18px;
+    font-weight: 800;
+    letter-spacing: 2px;
+    font-family: sans-serif;
+  }"
+  href="https://discord.gg/cXs9vZSqeK"  target="_blank">${showTextByLanguage(
+    'Welcome to Mixlab nodes discord, seeking help.',
+    {
+      'Welcome to Mixlab nodes discord, seeking help.':
+        '寻求帮助，加入Mixlab nodes交流频道'
+    }
+  )}</a><br><br>${showTextByLanguage(
+      'When loading the graph, the following node types were not found:',
+      {
+        'When loading the graph, the following node types were not found:':
+          '缺少以下节点：'
+      }
+    )}
+  
+   <ul>${missingNodeGithub(missingNodeTypes, nodesMap).join('')}</ul>${
+      hasAddedNodes ? '' : ''
    }`
  )
  this.logging.addEntry('Comfy.App', 'warn', {
@@ -585,13 +771,245 @@ async function fetchReadmeContent (url) {
  }
 }

+function createModelsModal (models) {
+  var div =
+    document.querySelector('#model-modal') || document.createElement('div')
+  div.id = 'model-modal'
+  div.innerHTML = ''
+  div.style.cssText = `
+    width: 100%;
+    z-index: 9990;
+    height: 100vh;
+    display: flex;
+    color: var(--descrip-text);
+    position: fixed;
+    top: 0;
+    left: 0;
+    background: #000000a8;
+    `
+
+  var modal = document.createElement('div')
+
+  div.addEventListener('click', e => {
+    e.stopPropagation()
+    div.remove()
+  })
+
+  div.appendChild(modal)
+  modal.classList.add('modal-body')
+  // Set modal styles
+  modal.style.cssText = `
+  color: var(--descrip-text);
+    background-color: var(--comfy-menu-bg);
+    position: fixed;
+    overflow:hidden;
+    top: 50%;
+    left: 50%;
+    transform: translate(-50%, -50%);
+    z-index: 9999;
+    border-radius: 4px;
+    box-shadow: 4px 4px 14px rgba(255,255,255,0.2);
+  `
+
+  // Create modal header
+  const headerElement = document.createElement('div')
+  headerElement.classList.add('modal-header')
+  headerElement.style.cssText = `
+      display: flex;
+      padding: 20px 24px 8px 24px;
+      justify-content: space-between;
+    `
+
+  const headTitleElement = document.createElement('a')
+  headTitleElement.classList.add('header-title')
+  headTitleElement.style.cssText = `
+      color: var(--descrip-text);
+      font-size: 18px;
+      display: flex;
+      align-items: center;
+      flex: 1;
+      overflow: hidden;
+      text-decoration: none;
+      font-weight: bold;
+      justify-content: space-between;
+      padding: 20px;
+      cursor: pointer;
+      user-select: none;
+    `
+
+  // headTitleElement.href = 'https://github.com/shadowcz007/comfyui-mixlab-nodes'
+  // headTitleElement.target = '_blank'
+  const linkIcon = document.createElement('small')
+  linkIcon.textContent = showTextByLanguage('Auto Open', {
+    'Auto Open': '自动开启'
+  })
+  linkIcon.style.padding = '4px'
+
+  const n_gpu = document.createElement('input')
+  n_gpu.type = 'number'
+  n_gpu.setAttribute('min', -1)
+  n_gpu.setAttribute('max', 9999)
+
+  n_gpu.style = `color: var(--input-text);
+  background-color: var(--comfy-input-bg);
+  border-radius: 8px;
+  border-color: var(--border-color);
+  height: 26px;
+  padding: 4px 10px;
+  width: 48px;
+  margin-left: 12px;`
+  if (localStorage.getItem('_mixlab_llama_n_gpu')) {
+    n_gpu.value = parseInt(localStorage.getItem('_mixlab_llama_n_gpu'))
+  } else {
+    n_gpu.value = -1
+    localStorage.setItem('_mixlab_llama_n_gpu', -1)
+  }
+
+  const n_gpu_p = document.createElement('p')
+  n_gpu_p.innerText = 'n_gpu_layers'
+
+  const n_gpu_div = document.createElement('div')
+  n_gpu_div.style = `display: flex;
+  justify-content: center;
+  align-items: center;
+  font-size: 12px;`
+  n_gpu_div.appendChild(n_gpu_p)
+  n_gpu_div.appendChild(n_gpu)
+
+  const title = document.createElement('p')
+  title.innerText = 'Models'
+  title.style = `font-size: 18px;
+  margin-right: 8px;`
+
+  const left_d = document.createElement('div')
+  left_d.style = `display: flex;
+  justify-content: center;
+  align-items: center;
+  font-size: 12px;`
+  left_d.appendChild(title)
+
+  left_d.appendChild(linkIcon)
+
+  headTitleElement.appendChild(left_d)
+
+  headTitleElement.appendChild(n_gpu_div)
+
+  //重启
+  const reStart = document.createElement('small')
+  reStart.textContent = showTextByLanguage('restart', {
+    restart: '重启'
+  })
+
+  reStart.style.padding = '4px'
+
+  headTitleElement.appendChild(reStart)
+
+  if (localStorage.getItem('_mixlab_auto_llama_open')) {
+    linkIcon.style.backgroundColor = '#66ff6c'
+    linkIcon.style.color = 'black'
+  }
+  linkIcon.addEventListener('click', e => {
+    e.stopPropagation()
+    if (localStorage.getItem('_mixlab_auto_llama_open')) {
+      localStorage.setItem('_mixlab_auto_llama_open', '')
+      linkIcon.style.backgroundColor = ''
+      linkIcon.style.color = 'var(--descrip-text)'
+    } else {
+      localStorage.setItem('_mixlab_auto_llama_open', 'true')
+      linkIcon.style.backgroundColor = '#66ff6c'
+      linkIcon.style.color = 'black'
+    }
+  })
+
+  reStart.addEventListener('click', e => {
+    e.stopPropagation()
+    div.remove()
+    fetch('mixlab/re_start', {
+      method: 'POST'
+    })
+  })
+
+  n_gpu.addEventListener('click', e => {
+    e.stopPropagation()
+    localStorage.setItem('_mixlab_llama_n_gpu', n_gpu.value)
+  })
+
+  modal.appendChild(headTitleElement)
+
+  // Create modal content area
+  var modalContent = document.createElement('div')
+  modalContent.classList.add('modal-content')
+
+  var input = document.createElement('textarea')
+  input.className = 'comfy-multiline-input'
+  input.style = `    height: 260px;
+  width: 480px;
+  font-size: 16px;
+  padding: 18px;`
+  input.value = localStorage.getItem('_mixlab_system_prompt')
+
+  input.addEventListener('change', e => {
+    e.stopPropagation()
+    localStorage.setItem('_mixlab_system_prompt', input.value)
+  })
+
+  input.addEventListener('click', e => {
+    e.stopPropagation()
+  })
+
+  modalContent.appendChild(input)
+
+  for (const m of models) {
+    let d = document.createElement('div')
+    d.innerText = m
+    d.className = `mix_tag`
+
+    if (!window._mixlab_llamacpp) {
+      d.addEventListener('click', async e => {
+        e.stopPropagation()
+        div.remove()
+        let res = await start_llama(m)
+        window._mixlab_llamacpp = res
+
+        localStorage.setItem('_mixlab_llama_select', res.model)
+
+        if (document.body.querySelector('#mixlab_chatbot_by_llamacpp')) {
+          document.body
+            .querySelector('#mixlab_chatbot_by_llamacpp')
+            .setAttribute('title', window._mixlab_llamacpp.url)
+        }
+      })
+    }
+
+    modalContent.appendChild(d)
+  }
+  modal.appendChild(modalContent)
+
+  const helpInfo = document.createElement('a')
+  helpInfo.textContent = showTextByLanguage('Help', {
+    Help: '寻求帮助'
+  })
+  helpInfo.style = `text-align: center;
+  display: block;
+  padding: 8px;
+  cursor: pointer;
+  font-size: 12px;
+  color: white;`
+  helpInfo.href="https://discord.gg/cXs9vZSqeK"
+  helpInfo.target="_blank"
+  modal.appendChild(helpInfo)
+
+  document.body.appendChild(div)
+}
+
 function createModal (url, markdown, title) {
  // Create modal element
  var div =
    document.querySelector('#mix-modal') || document.createElement('div')
  div.id = 'mix-modal'
  div.innerHTML = ''
-  div.style.cssText = `width: 100%;
+  div.style.cssText = `
+    width: 100%;
    z-index: 9990;
    height: 100vh;
    display: flex;
@@ -894,9 +1312,46 @@ function drawBadge (node, orig, restArgs) {
  return r
 }

+function convertImageUrlToBase64 (imageUrl) {
+  return fetch(imageUrl)
+    .then(response => response.blob())
+    .then(blob => {
+      return new Promise((resolve, reject) => {
+        const reader = new FileReader()
+        reader.onloadend = () => resolve(reader.result)
+        reader.onerror = reject
+        reader.readAsDataURL(blob)
+      })
+    })
+}
+
+async function getSelectImageNode () {
+  var nodes = app.canvas.selected_nodes
+  let imageNode = null
+  if (Object.keys(app.canvas.selected_nodes).length == 0) return
+  for (var id in nodes) {
+    if (nodes[id].imgs) {
+      let base64 = await convertImageUrlToBase64(nodes[id].imgs[0].currentSrc)
+      imageNode = await resizeImage(base64)
+    }
+  }
+  return imageNode
+}
+
 app.registerExtension({
  name: 'Comfy.Mixlab.ui',
  init () {
+    //是否要自动加载模型
+    if (localStorage.getItem('_mixlab_auto_llama_open')) {
+      let model = localStorage.getItem('_mixlab_llama_select')
+      start_llama(model).then(res => {
+        window._mixlab_llamacpp = res
+        document.body
+          .querySelector('#mixlab_chatbot_by_llamacpp')
+          .setAttribute('title', res.url)
+      })
+    }
+
    LGraphCanvas.prototype.helpAboutNode = async function (node) {
      nodesMap =
        nodesMap && Object.keys(nodesMap).length > 0
@@ -921,75 +1376,169 @@ app.registerExtension({

    smart_init()

-    const getNodeMenuOptions = LGraphCanvas.prototype.getNodeMenuOptions // store the existing method
-    LGraphCanvas.prototype.getNodeMenuOptions = function (node) {
-      // replace it
-      const options = getNodeMenuOptions.apply(this, arguments) // start by calling the stored one
-      node.setDirtyCanvas(true, true) // force a redraw of (foreground, background)
-      console.log('getNodeMenuOptions', node.type == 'CLIPTextEncode')
+    LGraphCanvas.prototype.text2text = async function (node) {
+      // console.log(node)
+      let widget = node.widgets.filter(
+        w => w.name === 'text' && typeof w.value == 'string'
+      )[0]
+      if (widget) {
+        app.canvas.centerOnNode(node)

-      let opts = [
-        {
-          content: 'Help ♾️Mixlab', // with a name
-          callback: () => {
-            LGraphCanvas.prototype.helpAboutNode(node)
-          } // and the callback
-        },
-        {
-          content: 'Fix node v2', // with a name
-          callback: () => {
-            LGraphCanvas.prototype.fixTheNode(node)
+        let controller = new AbortController()
+        let ends = [] //TODO 判断终止 <|im_start|>
+        let userInput = widget.value
+        widget.value = widget.value.trim()
+        widget.value += '\n'
+        try {
+          await completion_(
+            window._mixlab_llamacpp.url + '/v1/chat/completions',
+            [
+              {
+                role: 'system',
+                content: localStorage.getItem('_mixlab_system_prompt')
+              },
+              { role: 'user', content: userInput }
+            ],
+            controller,
+            t => {
+              // console.log(t)
+              widget.value += t
+            }
+          )
+        } catch (error) {
+          //是否要自动加载模型
+          if (localStorage.getItem('_mixlab_auto_llama_open')) {
+            let model = localStorage.getItem('_mixlab_llama_select')
+            start_llama(model).then(async res => {
+              window._mixlab_llamacpp = res
+              document.body
+                .querySelector('#mixlab_chatbot_by_llamacpp')
+                .setAttribute('title', res.url)
+
+              await completion_(
+                window._mixlab_llamacpp.url + '/v1/chat/completions',
+                [
+                  {
+                    role: 'system',
+                    content: localStorage.getItem('_mixlab_system_prompt')
+                  },
+                  { role: 'user', content: userInput }
+                ],
+                controller,
+                t => {
+                  // console.log(t)
+                  widget.value += t
+                }
+              )
+            })
          }
        }
-      ]

-      if (node.widgets) {
-        // let text_widget = node.widgets.filter(
-        //   w => w.name === 'text' && typeof w.value == 'string'
-        // )
-        // if (text_widget && text_widget.length == 1) {
-        //   opts.push({
-        //     content: 'Text-to-Text ♾️Mixlab', // with a name
-        //     callback: () => {
-        //       LGraphCanvas.prototype.text2text(node)
-        //     } // and the callback
-        //   })
-        // }
+        widget.value = widget.value.trim()
      }
+    }

-      opts = addSmartMenu(opts, node)
+    LGraphCanvas.prototype.image2text = async function (node) {
+      let imageBase64 = await getSelectImageNode()

-      // if (node.type == 'CLIPTextEncode') {
-      //   // 则出现 randomPrompt
-      //   // CLIPTextEncode 的widget ，name== 'text'
-      //   let node_widget_name = 'text'
-      //   const widget = node.widgets.filter(w => w.name === node_widget_name)[0]
+      if (imageBase64) {
+        // console.log('image2text')
+        // 添加note 节点
+        const NoteNode = LiteGraph.createNode('Note')
+        NoteNode.title = `Image-to-Text ${node.id}`
+        NoteNode.size = [NoteNode.size[0] + 100, NoteNode.size[1]]
+        let widget = NoteNode.widgets[0]
+        widget.value = ''

-      //  let mixlab_nodes_smart_connect= [{node_type:'CLIPTextEncode',
-      //   node_widget_name:'text',
-      //   inputNodeName:'RandomPrompt',
-      //   inputNode_output_type:'STRING'}]
+        NoteNode.pos = [node.pos[0] + node.size[0] + 24, node.pos[1] - 48]

-      //   if (widget) {
-      //     opts = [
-      //       {
-      //         content: 'RandomPrompt',
-      //         callback: () => {
-      //           LGraphCanvas.prototype._createNodeForInput(
-      //             node, //当前node
-      //             widget,//当前node里需要自动连线的widget
-      //             'RandomPrompt',//作为input的node type
-      //             'STRING'// 作为input的node的outputs的type. the input slot type of the target node
-      //           )
-      //         }
-      //       },
-      //       null,
-      //       ...opts
-      //     ]
-      //   }
-      // }
+        app.canvas.graph.add(NoteNode, false)
+        app.canvas.centerOnNode(NoteNode)

-      return [...opts, null, ...options] // and return the options
+        let controller = new AbortController()
+        let ends = []
+        let userInput = widget.value
+        widget.value = widget.value.trim()
+        widget.value += '\n'
+
+        try {
+          await completion_(
+            window._mixlab_llamacpp.url + '/v1/chat/completions',
+            [
+              {
+                role: 'system',
+                content: localStorage.getItem('_mixlab_system_prompt')
+              },
+              // { role: 'user', content: userInput }
+
+              {
+                role: 'user',
+                content: [
+                  {
+                    type: 'image_url',
+                    image_url: {
+                      url: imageBase64
+                    }
+                  },
+                  { type: 'text', text: 'What’s in this image?' }
+                ]
+              }
+            ],
+            controller,
+            t => {
+              // console.log(t)
+              widget.value += t
+
+              NoteNode.size[1] = widget.element.scrollHeight + 20
+              widget.computedHeight = NoteNode.size[1]
+              app.canvas.centerOnNode(NoteNode)
+            }
+          )
+        } catch (error) {
+          //是否要自动加载模型
+          if (localStorage.getItem('_mixlab_auto_llama_open')) {
+            let model = localStorage.getItem('_mixlab_llama_select')
+            start_llama(model).then(async res => {
+              window._mixlab_llamacpp = res
+              document.body
+                .querySelector('#mixlab_chatbot_by_llamacpp')
+                .setAttribute('title', res.url)
+
+              await completion_(
+                window._mixlab_llamacpp.url + '/v1/chat/completions',
+                [
+                  {
+                    role: 'system',
+                    content: localStorage.getItem('_mixlab_system_prompt')
+                  },
+                  {
+                    role: 'user',
+                    content: [
+                      {
+                        type: 'image_url',
+                        image_url: {
+                          url: imageBase64
+                        }
+                      },
+                      { type: 'text', text: 'What’s in this image?' }
+                    ]
+                  }
+                ],
+                controller,
+                t => {
+                  // console.log(t)
+                  widget.value += t
+                  NoteNode.size[1] = widget.element.scrollHeight + 20
+                  widget.computedHeight = NoteNode.size[1]
+                  app.canvas.centerOnNode(NoteNode)
+                }
+              )
+            })
+          }
+        }
+
+        widget.value = widget.value.trim()
+      }
    }

    const getGroupMenuOptions = LGraphCanvas.prototype.getGroupMenuOptions // store the existing method
@@ -1137,6 +1686,120 @@ app.registerExtension({
        (this.canvas.height * 0.5) / (this.ds.scale * dpr) // 考虑设备像素比
      this.setDirty(true, true)
    }
+
+    const getNodeMenuOptions = LGraphCanvas.prototype.getNodeMenuOptions
+    LGraphCanvas.prototype.getNodeMenuOptions = function (node) {
+      // replace it
+      const options = getNodeMenuOptions.apply(this, arguments) // start by calling the stored one
+      node.setDirtyCanvas(true, true) // force a redraw of (foreground, background)
+
+      let opts = [
+        {
+          content: 'Help ♾️Mixlab', // with a name
+          callback: () => {
+            LGraphCanvas.prototype.helpAboutNode(node)
+          } // and the callback
+        },
+        {
+          content: 'Fix node v2', // with a name
+          callback: () => {
+            LGraphCanvas.prototype.fixTheNode(node)
+          }
+        }
+      ]
+
+      if (node.widgets) {
+        let text_widget = node.widgets.filter(
+          w => w.name === 'text' && typeof w.value == 'string'
+        )
+
+        let text_input = node.inputs?.filter(
+          inp => inp.name == 'text' && inp.type == 'STRING'
+        )
+
+        if (
+          text_input &&
+          text_input.length == 0 &&
+          text_widget &&
+          text_widget.length == 1 &&
+          window._mixlab_llamacpp &&
+          node.type != 'ShowTextForGPT'
+        ) {
+          opts.push({
+            content: 'Text-to-Text ♾️Mixlab', // with a name
+            callback: () => {
+              LGraphCanvas.prototype.text2text(node)
+            } // and the callback
+          })
+        }
+
+        if (
+          node.imgs &&
+          node.imgs.length > 0 &&
+          window._mixlab_llamacpp &&
+          window._mixlab_llamacpp.chat_format === 'llava-1-5'
+        ) {
+          opts.push({
+            content: 'Image-to-Text ♾️Mixlab', // with a name
+            callback: () => {
+              LGraphCanvas.prototype.image2text(node)
+            } // and the callback
+          })
+        }
+      }
+
+      return [...opts, null, ...options] // and return the options
+    }
+
+    // 支持app模式的json
+    const loadAppJson = async data => {
+      let workflow
+      try {
+        let w = JSON.parse(data)
+        if (w.app && w.output) workflow = w.workflow
+      } catch (err) {}
+
+      if (workflow && workflow.version && workflow.nodes && workflow.extra) {
+        await app.loadGraphData(workflow)
+      }
+    }
+
+    if (!window._mixlab_app_paste_listener) {
+      window._mixlab_app_paste_listener = true
+      //粘贴json的事件
+      document.addEventListener('paste', async e => {
+        // ctrl+shift+v is used to paste nodes with connections
+        // this is handled by litegraph
+        if (this.shiftDown) return
+
+        let data = e.clipboardData || window.clipboardData
+
+        // No image found. Look for node data
+        data = data.getData('text/plain')
+
+        loadAppJson(data)
+      })
+
+      // 把json往里 拖
+      document.addEventListener('drop', async event => {
+        event.preventDefault()
+        event.stopPropagation()
+
+        // Dragging from Chrome->Firefox there is a file but its a bmp, so ignore that
+        if (
+          event.dataTransfer.files.length &&
+          event.dataTransfer.files[0].type == 'application/json'
+        ) {
+          const reader = new FileReader()
+          reader.onload = async () => {
+            loadAppJson(reader.result)
+          }
+          reader.readAsText(event.dataTransfer.files[0])
+        }
+      })
+    }
+
+    createMenu()
  },
  setup () {
    setTimeout(async () => {
@@ -1146,6 +1809,8 @@ app.registerExtension({
      const apps = await get_my_app()
      if (!apps) return

+      console.log('apps', apps)
+
      let apps_map = { 0: [] }

      for (const app of apps) {
@@ -1511,14 +2176,16 @@ app.registerExtension({
                document.body.appendChild(div)
            }
          },
-          apps_opts.length>0?{
-            content: 'Workflow App ♾️Mixlab',
-            has_submenu: true,
-            disabled: false,
-            submenu: {
-              options: apps_opts
-            }
-          }:null
+          apps_opts.length > 0
+            ? {
+                content: 'Workflow App ♾️Mixlab',
+                has_submenu: true,
+                disabled: false,
+                submenu: {
+                  options: apps_opts
+                }
+              }
+            : null
        )

        return options
@@ -1587,3 +2254,41 @@ app.registerExtension({
    } catch (error) {}
  }
 })
+
+//获取当前显存
+function fetchSystemStats () {
+  return new Promise(async (resolve, reject) => {
+    try {
+      const response = await fetch('/system_stats')
+      const data = await response.json()
+      resolve(data)
+    } catch (error) {
+      reject(error)
+    }
+  })
+}
+//清理显存
+function postFreeData () {
+  return new Promise(async (resolve, reject) => {
+    try {
+      const postData = {
+        unload_models: true,
+        free_memory: true
+      }
+      const response = await fetch('/free', {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json'
+        },
+        body: JSON.stringify(postData)
+      })
+      if (response.ok) {
+        resolve()
+      } else {
+        reject(new Error('Request failed'))
+      }
+    } catch (error) {
+      reject(error)
+    }
+  })
+}
@@ -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;
+}
Author	SHA1	Message	Date
shadowcz007	6cb9df700b	增加ComparingTwoFrames、右键image-to-text	2024-05-14 09:59:39 +08:00
shadowcz007	4fdda537b9	Update README.md	2024-05-14 09:57:51 +08:00
shadowcz007	cb6f32465a	Update README.md	2024-05-14 09:55:34 +08:00
shadowcz007	c235e36cb4	Merge branch 'main' of https://github.com/shadowcz007/comfyui-mixlab-nodes	2024-05-14 09:54:27 +08:00
shadowcz007	aaca440a94	Update README.md	2024-05-14 09:54:24 +08:00
shadow	1f57950a29	Update README.md	2024-05-14 09:39:57 +08:00
shadow	3d8855ec72	Update README.md	2024-05-14 09:39:45 +08:00
shadowcz007	ac9231d9f3	Update image_mixlab.js	2024-05-13 21:43:20 +08:00
shadowcz007	37c0a56d89	add ComparingTwoFrames	2024-05-13 21:33:44 +08:00
shadowcz007	010915dac4	add help	2024-05-13 11:37:17 +08:00
shadowcz007	83a8b3b970	Update ui_mixlab.js	2024-05-13 10:58:36 +08:00
shadowcz007	f130202aa1	resizeImage	2024-05-12 22:25:08 +08:00
shadowcz007	c8f6800bcd	add image-to-text :llava-phi-3-mini-gguf	2024-05-12 17:34:53 +08:00
shadowcz007	078f9f5dd4	Update ui_mixlab.js	2024-05-12 00:03:11 +08:00
shadowcz007	c0de178c7d	add re_start	2024-05-11 23:58:05 +08:00
shadowcz007	1c767b538d	set n_gpu_layers	2024-05-11 17:53:18 +08:00
shadowcz007	51aab44b5d	Update ui_mixlab.js	2024-05-11 14:43:35 +08:00
shadowcz007	b8a0d4a67b	Update ui_mixlab.js	2024-05-11 14:17:01 +08:00
shadowcz007	cd0dcfbb8c	v0.25.1	2024-05-11 14:12:00 +08:00
shadow	fcc9e30eae	Update __init__.py	2024-05-11 12:55:14 +08:00
shadowcz007	5f66218a43	修复 sys.stdout.isatty() object has no attribute 'isatty'	2024-05-11 12:28:14 +08:00
shadowcz007	61ef4f9a0f	Merge branch 'main' of https://github.com/shadowcz007/comfyui-mixlab-nodes	2024-05-11 12:13:34 +08:00
shadowcz007	f0a8734b42	修复 sys.stdout.isatty() object has no attribute 'isatty'	2024-05-11 12:13:31 +08:00
shadow	4fe95ef4ec	Update README.md	2024-05-11 08:53:33 +08:00
shadowcz007	2a5148845b	yaml	2024-05-10 14:38:32 +08:00
shadowcz007	8fa562caaf	Update install.bat	2024-05-09 09:55:34 +08:00
shadowcz007	cab5620cd5	Update index.html	2024-05-08 23:42:48 +08:00
shadowcz007	be38d36677	Update index.html	2024-05-08 23:32:20 +08:00
shadowcz007	69236fca89	Update __init__.py	2024-05-08 22:55:42 +08:00
shadowcz007	4a4f376bfd	Update __init__.py	2024-05-08 22:53:39 +08:00
shadowcz007	fd9718fe24	Update __init__.py	2024-05-08 22:50:30 +08:00
shadowcz007	26a6e11212	llama_cpp	2024-05-08 22:42:34 +08:00
shadowcz007	de1a669f6e	Update README.md	2024-05-08 10:11:16 +08:00
shadowcz007	e482c9e5c4	0.25.0 text-to-text for prompt	2024-05-07 21:26:47 +08:00
shadowcz007	4e96a77a41	Update README.md	2024-05-07 21:12:27 +08:00
shadowcz007	3346290e5c	Update ImageNode.py	2024-05-07 20:59:22 +08:00
shadowcz007	dcac593efe	Update install.bat	2024-05-07 12:49:03 +08:00
shadowcz007	7248d0de02	update	2024-05-07 12:14:42 +08:00
shadowcz007	8ad3ce632c	Update index.html	2024-05-07 00:08:35 +08:00
shadowcz007	9398b02562	Update ImageNode.py	2024-05-06 22:13:09 +08:00
shadowcz007	164e4da99d	Update ImageNode.py	2024-05-05 18:04:00 +08:00
shadowcz007	0025ea6119	output defaultImage	2024-05-05 13:16:14 +08:00
shadowcz007	eef53a5165	Update ImageNode.py	2024-05-05 12:59:21 +08:00
shadowcz007	9ea066d948	composite_images add position	2024-05-05 12:20:57 +08:00
shadowcz007	4bd900c4a1	Update __init__.py	2024-05-04 09:44:01 +08:00
shadowcz007	736cd2bebd	兼容旧版comfyui	2024-05-04 09:37:25 +08:00
shadowcz007	cd658c2a60	Merge branch 'main' of https://github.com/shadowcz007/comfyui-mixlab-nodes	2024-05-03 22:08:43 +08:00
shadowcz007	9730658f21	Update __init__.py	2024-05-03 22:08:40 +08:00
shadow	3fa107acb1	Update ui_mixlab.js	2024-05-03 18:35:39 +08:00
shadowcz007	7da22179a0	text-to-text	2024-05-03 00:30:16 +08:00
shadowcz007	fc7b71ee78	Update ui_mixlab.js	2024-05-02 19:02:33 +08:00
shadowcz007	b79573bf1f	Update __init__.py	2024-05-02 19:02:28 +08:00
shadowcz007	a01db6f7c0	Update README.md	2024-05-02 12:09:47 +08:00
shadowcz007	117d58c58e	v0.24.0	2024-05-02 12:06:19 +08:00
shadowcz007	ff6626ed89	add llama.cpp & Local LLM -Phi-3 & llama3 Phi-3 llama3	2024-05-02 12:02:47 +08:00
shadowcz007	10c798a440	Update index.html	2024-05-02 10:54:20 +08:00
shadowcz007	9097d87819	Update index.html	2024-05-02 10:46:09 +08:00
shadowcz007	d901f503d1	fixbug	2024-05-02 10:06:23 +08:00
shadowcz007	65f2b6ce6f	update https	2024-05-02 09:07:28 +08:00
shadowcz007	d949fe8bf1	Update __init__.py	2024-05-01 22:26:46 +08:00
shadowcz007	15cfb48550	Update __init__.py	2024-05-01 20:52:50 +08:00
shadowcz007	447dc6d4c3	fixbug	2024-05-01 15:55:43 +08:00
shadowcz007	c92e43b920	Update checkVersion_mixlab.js	2024-04-29 00:16:55 +08:00
shadowcz007	be6c32b0e0	fixbug	2024-04-29 00:15:47 +08:00
shadowcz007	3d2062e810	add TripoSRModel	2024-04-29 00:15:32 +08:00
shadowcz007	dd816e95cd	Update README.md	2024-04-27 23:43:40 +08:00
shadowcz007	6d1b51890d	Update checkVersion_mixlab.js	2024-04-27 23:38:00 +08:00
shadowcz007	11f03ec99a	支持正片叠底	2024-04-25 21:23:40 +08:00
shadowcz007	bd192f43e7	优化	2024-04-24 23:10:03 +08:00
shadowcz007	36e4b11983	gridout can export mask	2024-04-23 12:14:50 +08:00
shadowcz007	97397ba8c2	Update ImageNode.py	2024-04-22 12:32:48 +08:00
shadowcz007	052eee4111	Update __init__.py	2024-04-22 12:31:40 +08:00
shadowcz007	e319496044	v0.22.0 - 优化ImageColorTransfer - 支持动态提示 - 添加更多节点支持，AppInfo自动填充id - LoadImagesToBatch加载 - zhipuai 按需安装	2024-04-22 09:52:45 +08:00
shadowcz007	b83b63c362	添加支持的节点自动填充id	2024-04-21 21:51:07 +08:00
shadowcz007	4d6b1675bb	Load Images to Batch加载文件最大宽度1024	2024-04-21 21:50:56 +08:00
shadowcz007	13110fab39	支持动态提示	2024-04-21 20:58:44 +08:00
shadowcz007	74ea509848	fixbug	2024-04-20 23:05:59 +08:00
shadowcz007	192bff9d2c	optimize ImageColorTransfer and support batching,	2024-04-20 22:23:34 +08:00
shadowcz007	44ed8812dc	output add TransparentImage	2024-04-19 22:27:53 +08:00
shadowcz007	200696ba21	Update ui_mixlab.js	2024-04-19 20:44:27 +08:00
shadowcz007	51cf3b0c04	install Zhipuai as needed	2024-04-19 11:41:56 +08:00
shadowcz007	6a4831c83b	add nodes map for appinfo	2024-04-18 18:24:58 +08:00
shadowcz007	c5e7ed95a3	Update index.html	2024-04-18 16:10:02 +08:00
shadowcz007	42a97fa4d9	Update image_mixlab.js	2024-04-18 16:00:43 +08:00
shadowcz007	45240d0012	Update image_mixlab.js	2024-04-18 16:00:18 +08:00
shadowcz007	8ed085febd	Update image_mixlab.js	2024-04-18 15:57:34 +08:00
shadowcz007	37803ea61b	Update smart_connect.js	2024-04-18 09:06:42 +08:00
shadowcz007	acd416952c	Update image_mixlab.js	2024-04-18 07:48:50 +08:00
shadowcz007	6ec46cbc44	Update ui_mixlab.js	2024-04-17 22:55:16 +08:00
shadowcz007	a9d971e476	Update __init__.py	2024-04-17 22:55:12 +08:00
shadowcz007	9fe064675d	paste appinfo data / 支持粘贴appinfo导出的数据	2024-04-17 16:28:32 +08:00
shadowcz007	c84fa467d0	Update requirements.txt	2024-04-17 15:36:49 +08:00
shadowcz007	3d7a55f6d3	SaveImageAndMetadata支持格式化文件名@bakkhos8	2024-04-17 15:01:05 +08:00