Update ComfyUI. Support HYDiT-1.1/1.2, Strandard/Diffusers Wrapper and model weights trained via KoHya or official training scripts.

Author: xuhuaren

README.md

@@ -25,6 +25,7 @@ This repo contains PyTorch model definitions, pre-trained weights and inference/
 > [**DialogGen: Multi-modal Interactive Dialogue System for Multi-turn Text-to-Image Generation**](https://arxiv.org/abs/2403.08857) <br>
 
 ## 🔥🔥🔥 News!!
+* Jul 15, 2024: :tada: Update ComfyUI to support standardized workflows and compatibility with weights from t2i module and Lora training for versions 1.1/1.2, as well as those trained by Kohya or the official script. See [ComfyUI](./comfyui-hydit) for details.
 * Jul 08, 2024: :tada: HYDiT-v1.2 version is released. Please check [HunyuanDiT-v1.2](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2) and [Distillation-v1.2](https://huggingface.co/Tencent-Hunyuan/Distillation-v1.2) for more details.
 * Jul 03, 2024: :tada: Kohya-hydit version now available for v1.1 and v1.2 models, with GUI for inference. Official Kohya version is under review. See [kohya](./kohya_ss-hydit) for details.
 * Jun 27, 2024: :art: Hunyuan-Captioner is released, providing fine-grained caption for training data. See [mllm](./mllm) for details.
@@ -648,43 +649,13 @@ We list some more useful configurations for easy usage:
 
 ### Using ComfyUI
 
-We provide several commands to quick start: 
-
-```shell
-# Download comfyui code
-git clone https://github.com/comfyanonymous/ComfyUI.git
-
-# Install torch, torchvision, torchaudio
-pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu117
+- Support two workflows: Standard ComfyUI and Diffusers Wrapper, with the former being recommended.
+- Support HunyuanDiT-v1.1 and v1.2.
+- Support module, lora and clip lora models trained by Kohya.
+- Support module, lora models trained by HunyunDiT official training scripts.
+- ControlNet is coming soon.
 
-# Install Comfyui essential python package.
-cd ComfyUI
-pip install -r requirements.txt
-
-# ComfyUI has been successfully installed!
-
-# Download model weight as before or link the existing model folder to ComfyUI.
-python -m pip install "huggingface_hub[cli]"
-mkdir models/hunyuan
-huggingface-cli download Tencent-Hunyuan/HunyuanDiT --local-dir ./models/hunyuan/ckpts
-
-# Move to the ComfyUI custom_nodes folder and copy comfyui-hydit folder from HunyuanDiT Repo.
-cd custom_nodes
-cp -r ${HunyuanDiT}/comfyui-hydit ./
-cd comfyui-hydit
-
-# Install some essential python Package.
-pip install -r requirements.txt
-
-# Our tool has been successfully installed!
-
-# Go to ComfyUI main folder
-cd ../..
-# Run the ComfyUI Lauch command
-python main.py --listen --port 80
-
-# Running ComfyUI successfully!
-```
+![Workflow](comfyui-hydit/img/workflow_v1.2_lora.png)
 More details can be found in [./comfyui-hydit](comfyui-hydit/README.md)
 
 ### Using Kohya

comfyui-hydit/LICENSE

@@ -1,128 +1,186 @@
 # comfyui-hydit
+The ComfyUI code is under review in the official repository. Meanwhile, a temporary version is available below for immediate community use.
+We welcome users to try our workflow and appreciate any inquiries or suggestions.
 
-This repository houses a tailored node and workflow designed specifically for HunYuan DIT. The official tests conducted on DDPM, DDIM, and DPMMS have consistently yielded results that align with those obtained through the Diffusers library. However, it's important to note that we cannot assure the consistency of results from other ComfyUI native samplers with the Diffusers inference. We cordially invite users to explore our workflow and are open to receiving any inquiries or suggestions you may have.
 
 ## Overview
+- Support two workflows: Standard ComfyUI and Diffusers Wrapper, with the former being recommended.
+- Support HunyuanDiT-v1.1 and v1.2.
+- Support module, lora and clip lora models trained by Kohya.
+- Support module, lora models trained by HunyunDiT official training scripts.
+- ControlNet is coming soon.
 
+### Standard Workflow (Recommended)
+- [HunyuanDiT-v1.2](workflow/workflow_v1.2_lora.json)
+![Workflow](img/workflow_v1.2_lora.png)
 
-### Workflow text2image
+- [HunyuanDiT-v1.1](workflow/workflow_v1.1_lora.json)
+![Workflow](img/workflow_v1.1_lora.png)
 
-![Workflow](img/work_diffusers.png)
-[workflow_diffusers](workflow/workflow_diffusers.json) file for HunyuanDiT txt2image with diffusers backend.  
-![Workflow](img/workflow_ksampler.png)
-[workflow_ksampler](workflow/workflow_ksampler.json) file for HunyuanDiT txt2image with ksampler backend.
+### Diffusers Wrapper
+- [HunyuanDiT-v1.1-diffusers](workflow/workflow_lora_controlnet.json)
 ![Workflow](img/workflow_lora_controlnet.png)
-[workflow_lora_controlnet_diffusers](workflow/workflow_lora_controlnet.json) file for HunyuanDiT lora and controlnet model with diffusers backend.
 
 
 ## Usage
+### Dependencies
 
-Make sure you run the following command inside [ComfyUI](https://github.com/comfyanonymous/ComfyUI) project with our [comfyui-hydit](.) and have correct conda environment.
-
+1. Official ComfyUI Environment Setup
 ```shell
 # Please use python 3.10 version with cuda 11.7
+conda create --name comfyui-hydit python=3.10
+conda activate comfyui-hydit
+
 # Download comfyui code
 git clone https://github.com/comfyanonymous/ComfyUI.git
+cd ${ComfyUI}
+git reset --hard 90389b3b8a69c08c3ed0bcc9d87a92246578a8e3
 
 # Install torch, torchvision, torchaudio
 pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu117 --default-timeout=100 future
 
 # Install Comfyui essential python package
-cd ComfyUI
 pip install -r requirements.txt
+```
 
-# ComfyUI has been successfully installed!
-
-
+2. HunyuanDiT Environment Setup in ComfyUI
+```shell
 # Move to the ComfyUI custom_nodes folder and copy comfyui-hydit folder from HunyuanDiT Repo.
-cd custom_nodes
+cd ${ComfyUI}/custom_nodes
 git clone https://github.com/Tencent/HunyuanDiT.git
 cp -r HunyuanDiT/comfyui-hydit ./
 rm -rf HunyuanDiT
-cd comfyui-hydit
+cd ${ComfyUI}/custom_nodes/comfyui-hydit
 
-# !!! If using windows system !!!
-cd custom_nodes
+# Install some essential python Package.
+pip install -r requirements.txt
+```
+
+3. (Optional) Deployment on Windows environment
+```shell
+cd ${ComfyUI}/custom_nodes
 git clone https://github.com/Tencent/HunyuanDiT.git
 xcopy /E /I HunyuanDiT\comfyui-hydit comfyui-hydit
 rmdir /S /Q HunyuanDiT
-cd comfyui-hydit
-
+cd ${ComfyUI}/custom_nodes/comfyui-hydit
 # Install some essential python Package.
 pip install -r requirements.txt
+```
 
-# Our tool has been successfully installed!
-
+4. Running ComfyUI successfully
+```shell
 # Go to ComfyUI main folder
-cd ../..
+cd ${ComfyUI}
 # Run the ComfyUI Lauch command
 python main.py --listen --port 80
+```
+
+### Standard workflow (Recommended)
+
+1. Preparing Model Weights
+
+Use the command below to download the file to the specified folder. If already downloaded, create a symbolic link to that folder.
+```shell
+# clip encoder
+wget -O ${ComfyUI}/models/clip/pytorch_model.bin https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/blob/main/t2i/clip_text_encoder/pytorch_model.bin
+# mt5
+mkdir ${ComfyUI}/models/t5
+wget -O ${ComfyUI}/models/t5/pytorch_model.bin https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/blob/main/t2i/mt5/pytorch_model.bin
+# vae
+wget -O ${ComfyUI}/models/vae/diffusion_pytorch_model.bin https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/blob/main/t2i/sdxl-vae-fp16-fix/diffusion_pytorch_model.bin
+# base model
+wget -O ${ComfyUI}/models/checkpoints/pytorch_model_ema.pt https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/blob/main/t2i/model/pytorch_model_ema.pt
+```
+Put module weights trained through Kohya or the official script in `${ComfyUI}/models/checkpoints/` to switch model weights in ComfyUI.
+
+2. Preparing LoRa Weights
 
-# Running ComfyUI successfully!
+```shell
+# Put LoRa weights trained by Kohya in ComfyUI/models/loras
+cp ${HunyuanDiT}/kohya_ss/outputs/last-step{xxxx}.safetensors ${ComfyUI}/models/loras
+
+# (Optional) Put LoRa weights trained by official scripts in ComfyUI/models/loras
+python convert_hunyuan_to_coimfyui_lora.py \
+      --lora_path ${HunyuanDiT}/log_EXP/001-lora_porcelain_ema_rank64/checkpoints/0000100.pt/adapter_model.safetensors \
+      --save_lora_path ${ComfyUI}/models/loras/adapter_model_convert.safetensors
+
+# update the `lora.py` file
+cp ${ComfyUI}/custom_nodes/comfyui-hydit/lora.py ${ComfyUI}/comfy/lora.py
 ```
 
-## Download weights for diffusers mode
+### Diffusers Wrapper
+1. Preparing Model Weights
 
 ```shell
 python -m pip install "huggingface_hub[cli]"
 mkdir models/hunyuan
-huggingface-cli download Tencent-Hunyuan/HunyuanDiT-v1.1 --local-dir ./models/hunyuan/ckpts
-huggingface-cli download Tencent-Hunyuan/Distillation-v1.1 pytorch_model_distill.pt --local-dir ./models/hunyuan/ckpts/t2i/model
+huggingface-cli download Tencent-Hunyuan/HunyuanDiT-v1.2 --local-dir ./models/hunyuan/ckpts
+huggingface-cli download Tencent-Hunyuan/HunyuanDiT-v1.2 t2i/model/pytorch_model_ema.pt --local-dir ./models/hunyuan/ckpts/t2i/model
 ```
 
-## Download weights for ksampler mode
-Download the [clip encoder](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT/blob/main/t2i/clip_text_encoder/pytorch_model.bin) and place it in `ComfyUI/models/clip`  
-Download the [mt5](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT/blob/main/t2i/mt5/pytorch_model.bin) and place it in `ComfyUI/models/t5`  
-Download the [base model](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT/blob/main/t2i/model/pytorch_model_ema.pt) and place it in `ComfyUI/models/checkpoints`  
-Download the [sdxl vae](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT/blob/main/t2i/sdxl-vae-fp16-fix/diffusion_pytorch_model.bin) and place it in `ComfyUI/models/vae`   
+2. Preparing LoRa Weights
 
+```shell
+# Put LoRa weights trained by Kohya in ComfyUI/models/loras
+cp ${HunyuanDiT}/kohya_ss/outputs/adapter_model.safetensors ${ComfyUI}/models/loras
+
+# (Optional) Put LoRa weights trained by official scripts in ComfyUI/models/loras
+# The PEFT diffuser format needs to be converted into the standard ComfyUI format
+python convert_hunyuan_to_coimfyui_lora.py \
+      --lora_path ${HunyuanDiT}/log_EXP/001-lora_porcelain_ema_rank64/checkpoints/0000100.pt/adapter_model.safetensors \
+      --save_lora_path ${ComfyUI}/models/loras/adapter_model_convert.safetensors
+
+# update the `lora.py` file
+cp ${ComfyUI}/custom_nodes/comfyui-hydit/lora.py ${ComfyUI}/comfy/lora.py
+```
 
 ## Custom Node
 Below I'm trying to document all the nodes, thanks for some good work[[1]](#1)[[2]](#2).
-#### HunYuan Pipeline Loader
+### HunYuan Pipeline Loader
 - Loads the full stack of models needed for HunYuanDiT.  
 - **pipeline_folder_name** is the official weight folder path for hunyuan dit including clip_text_encoder, model, mt5, sdxl-vae-fp16-fix and tokenizer.
 - **lora** optional to load lora weight.
 
-#### HunYuan Checkpoint Loader
+### HunYuan Checkpoint Loader
 - Loads the base model for HunYuanDiT in ksampler backend.  
 - **model_name** is the weight list of comfyui checkpoint folder.
+- **version** two option, v1.1 and v1.2.
 
 
-#### HunYuan CLIP Loader
+### HunYuan CLIP Loader
 - Loads the clip and mt5 model for HunYuanDiT in ksampler backend.  
 - **text_encoder_path** is the weight list of comfyui clip model folder.
 - **t5_text_encoder_path** is the weight list of comfyui t5 model folder.
 
-#### HunYuan VAE Loader
+### HunYuan VAE Loader
 - Loads the vae model for HunYuanDiT in ksampler backend.  
 - **model_name** is the weight list of comfyui vae model folder.
 
-#### HunYuan Scheduler Loader
+### HunYuan Scheduler Loader
 - Loads the scheduler algorithm for HunYuanDiT.  
 - **Input** is the algorithm name including ddpm, ddim and dpmms.
 - **Output** is the instance of diffusers.schedulers.
 
-#### HunYuan Model Makeup
+### HunYuan Model Makeup
 - Assemble the models and scheduler module.  
 - **Input** is the instance of StableDiffusionPipeline and diffusers.schedulers.
 - **Output** is the updated instance of StableDiffusionPipeline.
 
-#### HunYuan Clip Text Encode
+### HunYuan Clip Text Encode
 - Assemble the models and scheduler module.  
 - **Input** is the string of positive and negative prompts.
 - **Output** is the converted string for model.
 
-#### HunYuan Sampler
+### HunYuan Sampler
 - Similar with KSampler in ComfyUI.  
 - **Input** is the instance of StableDiffusionPipeline and some hyper-parameters for sampling.
 - **Output** is the generated image.
 
-#### HunYuan Lora Loader
+### HunYuan Lora Loader
 - Loads the lora model for HunYuanDiT in diffusers backend.  
 - **lora_name** is the weight list of comfyui lora folder.
 
-#### HunYuan ControNet Loader
+### HunYuan ControNet Loader
 - Loads the controlnet model for HunYuanDiT in diffusers backend.  
 - **controlnet_path** is the weight list of comfyui controlnet folder.
 

comfyui-hydit/README.md

@@ -3,16 +3,20 @@
 import comfy.model_patcher
 import comfy.model_base
 import comfy.utils
-from .hydit.modules.text_encoder import MT5Embedder
+from .hydit_v1_1.modules.text_encoder import MT5Embedder
 from transformers import BertModel, BertTokenizer, AutoTokenizer
 import torch
 import os
 from transformers import T5Config, T5EncoderModel, BertConfig, BertModel
 from transformers import AutoTokenizer, modeling_utils
+from comfy import model_management
 #import pdb
 
 class CLIP:
-    def __init__(self, root, text_encoder_path = None, t5_text_encoder_path = None):
+    def __init__(self,no_init=False, root = None, text_encoder_path = None, t5_text_encoder_path = None):
+        if no_init:
+            return
+
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         if text_encoder_path == None:
             text_encoder_path = os.path.join(root,"clip_text_encoder")
@@ -53,7 +57,10 @@ def __init__(self, root, text_encoder_path = None, t5_text_encoder_path = None):
 
 
         self.cond_stage_model = clip_text_encoder
-        
+        load_device = model_management.text_encoder_device()
+        offload_device = model_management.text_encoder_offload_device()
+        self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
+        self.layer_idx = None
 
     def tokenize(self, text):
         tokens = self.tokenizer.tokenize(text)
@@ -64,7 +71,9 @@ def tokenize(self, text):
     def tokenize_t5(self, text):
         return self.tokenizer_t5.tokenize(text)
 
-    def encode_from_tokens(self, tokens, return_pooled=False):
+    def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False):
+        #self.cond_stage_model.reset_clip_options()
+        self.load_model()
         attention_mask = tokens['attention_mask'].to(self.device)
         with torch.no_grad():
             prompt_embeds = self.cond_stage_model(
@@ -83,14 +92,44 @@ def encode_from_tokens(self, tokens, return_pooled=False):
         addit_embeds = {
                 "t5_embeds": t5_embeds,
                 "attention_mask": attention_mask.float(),
-                "t5_attention_mask": t5_attention_mask.float()
+                "t5_attention_mask": t5_attention_mask.float() 
             }
         prompt_embeds.addit_embeds = addit_embeds
 
         if return_pooled:
             return prompt_embeds, None
         else:
             return prompt_embeds
+        
+    def clone(self):
+        n = CLIP(no_init=True)
+        n.patcher = self.patcher.clone()
+        n.cond_stage_model = self.cond_stage_model
+        n.tokenizer = self.tokenizer
+        n.layer_idx = self.layer_idx
+        n.tokenizer_t5 = self.tokenizer_t5
+        n.embedder_t5 = self.embedder_t5
+        n.device = self.device 
+        return n
+    
+    def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
+        #import pdb
+        #pdb.set_trace()
+        return self.patcher.add_patches(patches, strength_patch, strength_model)
+    
+    def load_sd(self, sd):
+        return self.cond_stage_model.load_sd(sd)
+
+    def get_sd(self):
+        return self.cond_stage_model.state_dict()
+
+    def load_model(self):
+        #if self.load_device != "cpu":
+        model_management.load_model_gpu(self.patcher)
+        return self.patcher
+
+    def get_key_patches(self):
+        return self.patcher.get_key_patches()
 
 class HyBertTokenizer:
     def __init__(self, tokenizer_path=None, max_length=77, truncation=True, return_attention_mask=True, device='cpu'):