add hunyuan dit txt2image

.gitignore

@@ -5,7 +5,7 @@ __pycache__/
 !/input/example.png
 /models/
 /temp/
-/custom_nodes/
+#/custom_nodes/
 !custom_nodes/example_node.py.example
 extra_model_paths.yaml
 /.vs

custom_nodes/comfyui-hydit_min/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Limitex
+
+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.

custom_nodes/comfyui-hydit_min/README.md

@@ -0,0 +1,91 @@
+# comfyui-hydit
+
+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
+
+
+### Workflow text2image
+
+
+![Workflow](img/workflow_ksampler.png)
+[workflow_ksampler](workflow/workflow_ksampler.json) file for HunyuanDiT txt2image with ksampler backend.
+
+
+## Usage
+
+Make sure you run the following command inside [ComfyUI](https://github.com/comfyanonymous/ComfyUI) project with our [comfyui-hydit](.) and have correct conda environment.
+
+```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 --default-timeout=100 future
+
+# Install Comfyui essential python package
+cd ComfyUI
+pip install -r requirements.txt
+
+# ComfyUI has been successfully installed!
+
+
+# Move to the ComfyUI custom_nodes folder and copy comfyui-hydit folder from HunyuanDiT Repo.
+cd custom_nodes
+git clone https://github.com/Tencent/HunyuanDiT.git
+cp -r HunyuanDiT/comfyui-hydit ./
+rm -rf HunyuanDiT
+cd comfyui-hydit
+
+# If using windows system !!
+cd 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
+
+# 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!
+```
+
+
+## Download weights for ksampler mode
+Download the [clip encoder](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/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-v1.2/blob/main/t2i/mt5/pytorch_model.bin) and place it in `ComfyUI/models/t5`  
+Download the [base model](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/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-v1.2/blob/main/t2i/sdxl-vae-fp16-fix/diffusion_pytorch_model.bin) and place it in `ComfyUI/models/vae`   
+
+
+## Custom Node
+Below I'm trying to document all the nodes, thanks for some good work[[1]](#1)[[2]](#2).
+
+#### HunYuan Checkpoint Loader
+- Loads the base model for HunYuanDiT in ksampler backend.  
+- **model_name** is the weight list of comfyui checkpoint folder.
+
+
+#### 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
+- Loads the vae model for HunYuanDiT in ksampler backend.  
+- **model_name** is the weight list of comfyui vae model folder.
+
+
+
+## Reference 
+<a id="1">[1]</a> 
+https://github.com/Limitex/ComfyUI-Diffusers  
+<a id="2">[2]</a>
+https://github.com/Tencent/HunyuanDiT/pull/59

custom_nodes/comfyui-hydit_min/__init__.py

@@ -0,0 +1,4 @@
+from .nodes import *
+#aa = DiffusersSampler()
+#print(aa)
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

custom_nodes/comfyui-hydit_min/clip.py

@@ -0,0 +1,150 @@
+import comfy.supported_models_base
+import comfy.latent_formats
+import comfy.model_patcher
+import comfy.model_base
+import comfy.utils
+from .hydit.modules.text_encoder import MT5Embedder
+from transformers import BertModel, BertTokenizer, AutoTokenizer
+import torch
+import os
+from transformers import BertConfig, BertModel
+from transformers import AutoTokenizer, modeling_utils
+#import pdb
+
+class CLIP:
+    def __init__(self, root, text_encoder_path = None, t5_text_encoder_path = None):
+        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")
+            clip_text_encoder = BertModel.from_pretrained(str(text_encoder_path), False, revision=None).to(self.device)
+        else:         
+            config = BertConfig.from_json_file(os.path.join(os.path.dirname(os.path.realpath(__file__)),"config_clip.json"))
+            with modeling_utils.no_init_weights():
+                clip_text_encoder = BertModel(config) 
+            sd = comfy.utils.load_torch_file(text_encoder_path)
+            prefix = "bert."
+            state_dict = {}
+            for key in sd:
+                nkey = key
+                if key.startswith(prefix):
+                        nkey = key[len(prefix):]
+                state_dict[nkey] = sd[key]
+
+            m, e = clip_text_encoder.load_state_dict(state_dict, strict=False)
+            if len(m) > 0 or len(e) > 0:
+                print(f"HYDiT: clip missing {len(m)} keys ({len(e)} extra)")
+                #clip_text_encoder.load_state_dict(sd, strict=False)
+            clip_text_encoder.eval().to(self.device)    
+
+        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "tokenizer_clip")
+        self.tokenizer = HyBertTokenizer(tokenizer_path)
+        #assert(0)
+        
+        if t5_text_encoder_path == None:
+            t5_text_encoder_path = os.path.join(root,'mt5')
+            embedder_t5 = MT5Embedder(t5_text_encoder_path, torch_dtype=torch.float16, max_length=256)
+            self.tokenizer_t5 = HyT5Tokenizer(embedder_t5.tokenizer, max_length=embedder_t5.max_length)
+            self.embedder_t5 = embedder_t5.model
+        else:
+
+            embedder_t5 = MT5Embedder(t5_text_encoder_path, torch_dtype=torch.float16, max_length=256, ksampler=True)
+            self.tokenizer_t5 = HyT5Tokenizer(embedder_t5.tokenizer, max_length=embedder_t5.max_length)
+            self.embedder_t5 = embedder_t5.model
+        
+
+        self.cond_stage_model = clip_text_encoder
+        
+
+    def tokenize(self, text):
+        tokens = self.tokenizer.tokenize(text)
+        t5_tokens = self.tokenizer_t5.tokenize(text)
+        tokens.update(t5_tokens)
+        return tokens
+    
+    def tokenize_t5(self, text):
+        return self.tokenizer_t5.tokenize(text)
+
+    def encode_from_tokens(self, tokens, return_pooled=False):
+        attention_mask = tokens['attention_mask'].to(self.device)
+        with torch.no_grad():
+            prompt_embeds = self.cond_stage_model(
+                    tokens['text_input_ids'].to(self.device),
+                    attention_mask=attention_mask
+                )
+            prompt_embeds = prompt_embeds[0]
+        t5_attention_mask = tokens['t5_attention_mask'].to(self.device)
+        with torch.no_grad():
+            t5_prompt_cond = self.embedder_t5(
+                    tokens['t5_text_input_ids'].to(self.device),
+                    attention_mask=t5_attention_mask
+                )
+            t5_embeds = t5_prompt_cond[0]
+
+        addit_embeds = {
+                "t5_embeds": t5_embeds,
+                "attention_mask": 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
+
+class HyBertTokenizer:
+    def __init__(self, tokenizer_path=None, max_length=77, truncation=True, return_attention_mask=True, device='cpu'):
+        self.tokenizer = BertTokenizer.from_pretrained(str(tokenizer_path))
+        self.max_length = self.tokenizer.model_max_length or max_length
+        self.truncation = truncation
+        self.return_attention_mask = return_attention_mask
+        self.device = device
+
+    def tokenize(self, text:str):
+        text_inputs = self.tokenizer(
+                text,
+                padding="max_length",
+                max_length=self.max_length,
+                truncation=self.truncation,
+                return_attention_mask=self.return_attention_mask,
+			    add_special_tokens = True,
+                return_tensors="pt",
+            )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        tokens = {
+            'text_input_ids': text_input_ids,
+            'attention_mask': attention_mask
+        }
+        return tokens
+    
+class HyT5Tokenizer:
+    def __init__(self, tokenizer, tokenizer_path = None, max_length=77, truncation=True, return_attention_mask=True, device='cpu'):
+        if tokenizer_path:
+            self.tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_path))
+        else:
+            self.tokenizer = tokenizer
+
+        self.max_length = max_length
+        self.truncation = truncation
+        self.return_attention_mask = return_attention_mask
+        self.device = device
+
+    def tokenize(self, text:str):
+        text_inputs = self.tokenizer(
+                text,
+                padding="max_length",
+                max_length=self.max_length,
+                truncation=self.truncation,
+                return_attention_mask=self.return_attention_mask,
+			    add_special_tokens = True,
+                return_tensors="pt",
+            )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        tokens = {
+            't5_text_input_ids': text_input_ids,
+            't5_attention_mask': attention_mask
+        }
+        return tokens
+

custom_nodes/comfyui-hydit_min/config_clip.json

@@ -0,0 +1,34 @@
+{
+  "_name_or_path": "hfl/chinese-roberta-wwm-ext-large",
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "directionality": "bidi",
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "pooler_fc_size": 768,
+  "pooler_num_attention_heads": 12,
+  "pooler_num_fc_layers": 3,
+  "pooler_size_per_head": 128,
+  "pooler_type": "first_token_transform",
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.22.1",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 47020
+}

custom_nodes/comfyui-hydit_min/config_mt5.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "/home/patrick/t5/mt5-xl",
+  "architectures": [
+    "MT5ForConditionalGeneration"
+  ],
+  "d_ff": 5120,
+  "d_kv": 64,
+  "d_model": 2048,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "mt5",
+  "num_decoder_layers": 24,
+  "num_heads": 32,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "tokenizer_class": "T5Tokenizer",
+  "transformers_version": "4.10.0.dev0",
+  "use_cache": true,
+  "vocab_size": 250112
+}

custom_nodes/comfyui-hydit_min/constant.py

@@ -0,0 +1,10 @@
+import os
+from .hydit.constants import SAMPLER_FACTORY
+
+base_path = os.path.dirname(os.path.realpath(__file__))
+HUNYUAN_PATH = os.path.join(base_path, "..", "..", "models",  "hunyuan")
+SCHEDULERS_hunyuan = list(SAMPLER_FACTORY.keys())
+
+
+T5_PATH = os.path.join(base_path, "..", "..", "models",  "t5")
+LORA_PATH = os.path.join(base_path, "..", "..", "models",  "loras")

custom_nodes/comfyui-hydit_min/dit.py

@@ -0,0 +1,145 @@
+import comfy.supported_models_base
+import comfy.latent_formats
+import comfy.model_patcher
+import comfy.model_base
+import comfy.utils
+from comfy import model_management
+from .supported_dit_models import HunYuan_DiT, HYDiT_Model, ModifiedHunYuanDiT
+from .clip import CLIP
+import os
+from .hydit.config_comfyui import get_args
+
+sampling_settings = {
+	"beta_schedule" : "linear",
+	"linear_start"  : 0.00085,
+	"linear_end"    : 0.03,
+	"timesteps"     : 1000,
+}
+
+hydit_conf = {
+	"G/2": { # Seems to be the main one
+		"unet_config": {
+			"depth"       :   40,
+			"num_heads"   :   16,
+			"patch_size"  :    2,
+			"hidden_size" : 1408,
+			"mlp_ratio" : 4.3637,
+			"input_size": (1024//8, 1024//8),
+			#"disable_unet_model_creation": True,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+}
+
+def load_dit(model_path, output_clip=True, output_model=True, output_vae=True, MODEL_PATH = None, VAE_PATH = None):
+	if MODEL_PATH:
+		state_dict = comfy.utils.load_torch_file(MODEL_PATH)
+	else:
+		state_dict = comfy.utils.load_torch_file(os.path.join(model_path, "t2i", "model", "pytorch_model_ema.pt"))
+
+
+	state_dict = state_dict.get("model", state_dict)
+	parameters = comfy.utils.calculate_parameters(state_dict)
+	unet_dtype = model_management.unet_dtype(model_params=parameters)
+	load_device = comfy.model_management.get_torch_device()
+	offload_device = comfy.model_management.unet_offload_device()
+	clip = None,
+	vae = None
+	model_patcher = None
+
+	# ignore fp8/etc and use directly for now
+	manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
+	root = os.path.join(model_path, "t2i")
+	if manual_cast_dtype:
+		print(f"DiT: falling back to {manual_cast_dtype}")
+		unet_dtype = manual_cast_dtype
+
+	#model_conf["unet_config"]["num_classes"] = state_dict["y_embedder.embedding_table.weight"].shape[0] - 1 # adj. for empty
+
+	if output_model:
+		model_conf = HunYuan_DiT(hydit_conf["G/2"])
+		model = HYDiT_Model(
+			model_conf,
+			model_type=comfy.model_base.ModelType.V_PREDICTION,
+			device=model_management.get_torch_device()
+		)
+
+		#print(model_conf.unet_config)
+		#assert(0)
+		all_args =get_args()
+
+		model.diffusion_model = ModifiedHunYuanDiT(model_conf.dit_conf, **model_conf.unet_config).half().to(load_device)
+
+		model.diffusion_model.load_state_dict(state_dict)
+		model.diffusion_model.eval()
+		model.diffusion_model.to(unet_dtype)
+
+		model_patcher = comfy.model_patcher.ModelPatcher(
+			model,
+			load_device = load_device,
+			offload_device = offload_device,
+			current_device = "cpu",
+		)
+	if output_clip:
+		clip = CLIP(root)
+
+	if output_vae:
+		if VAE_PATH:
+			vae_path = VAE_PATH
+		else:
+			vae_path = os.path.join(root, "sdxl-vae-fp16-fix", "diffusion_pytorch_model.safetensors")
+		#print(vae_path)
+		sd = comfy.utils.load_torch_file(vae_path)
+		vae = comfy.sd.VAE(sd=sd)
+
+	return (model_patcher, clip, vae)
+
+
+def load_checkpoint(model_path):
+
+	state_dict = comfy.utils.load_torch_file(model_path)
+
+
+	state_dict = state_dict.get("model", state_dict)
+	parameters = comfy.utils.calculate_parameters(state_dict)
+	unet_dtype = model_management.unet_dtype(model_params=parameters)
+	load_device = comfy.model_management.get_torch_device()
+	offload_device = comfy.model_management.unet_offload_device()
+	model_patcher = None
+
+	# ignore fp8/etc and use directly for now
+	manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
+	if manual_cast_dtype:
+		print(f"DiT: falling back to {manual_cast_dtype}")
+		unet_dtype = manual_cast_dtype
+
+
+	model_conf = HunYuan_DiT(hydit_conf["G/2"])
+	model = HYDiT_Model(
+		model_conf,
+		model_type=comfy.model_base.ModelType.V_PREDICTION,
+		device=model_management.get_torch_device()
+	)
+
+
+	model.diffusion_model = ModifiedHunYuanDiT(model_conf.dit_conf, **model_conf.unet_config).half().to(load_device)
+
+	model.diffusion_model.load_state_dict(state_dict)
+	model.diffusion_model.eval()
+	model.diffusion_model.to(unet_dtype)
+
+	model_patcher = comfy.model_patcher.ModelPatcher(
+		model,
+		load_device = load_device,
+		offload_device = offload_device,
+		current_device = "cpu",)
+	return (model_patcher,)
+
+
+
+def load_vae(vae_path):
+
+	sd = comfy.utils.load_torch_file(vae_path)
+	vae = comfy.sd.VAE(sd=sd)
+
+	return (vae,)

custom_nodes/comfyui-hydit_min/hydit/config_comfyui.py

@@ -0,0 +1,227 @@
+
+import argparse
+#import deepspeed
+
+from .constants import *
+from .diffusion.gaussian_diffusion import ModelVarType
+from .modules.models import HUNYUAN_DIT_CONFIG
+
+
+def model_var_type(value):
+    try:
+        return ModelVarType[value]
+    except KeyError:
+        raise ValueError(f"Invalid choice '{value}', valid choices are {[v.name for v in ModelVarType]}")
+
+
+def get_args(default_args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task-flag", type=str)
+
+    # General Setting
+    parser.add_argument("--batch-size", type=int, default=1, help="Per-GPU batch size")
+    parser.add_argument('--seed', type=int, default=42, help="A seed for all the prompts.")
+    parser.add_argument("--use-fp16", action="store_true", help="Use FP16 precision.")
+    parser.add_argument("--no-fp16", dest="use_fp16", action="store_false")
+    parser.set_defaults(use_fp16=True)
+    parser.add_argument("--extra-fp16", action="store_true", help="Use extra fp16 for vae and text_encoder.")
+
+    # HunYuan-DiT
+    parser.add_argument("--model", type=str, choices=list(HUNYUAN_DIT_CONFIG.keys()), default='DiT-g/2')
+    parser.add_argument("--image-size", type=int, nargs='+', default=[1024, 1024],
+                        help='Image size (h, w). If a single value is provided, the image will be treated to '
+                             '(value, value).')
+    parser.add_argument("--qk-norm", action="store_true",
+                        help="Query Key normalization. See http://arxiv.org/abs/2302.05442 for details.")
+    parser.set_defaults(qk_norm=True)
+    parser.add_argument("--norm", type=str, choices=["rms", "laryer"], default="layer", help="Normalization layer type")
+    parser.add_argument("--text-states-dim", type=int, default=1024, help="Hidden size of CLIP text encoder.")
+    parser.add_argument("--text-len", type=int, default=77, help="Token length of CLIP text encoder output.")
+    parser.add_argument("--text-states-dim-t5", type=int, default=2048, help="Hidden size of CLIP text encoder.")
+    parser.add_argument("--text-len-t5", type=int, default=256, help="Token length of T5 text encoder output.")
+
+    # LoRA config
+    parser.add_argument("--training-parts", type=str, default='all', choices=['all', 'lora'], help="Training parts")
+    parser.add_argument("--rank", type=int, default=64, help="Rank of LoRA")
+    parser.add_argument("--lora-ckpt", type=str, default=None, help="LoRA checkpoint")
+    parser.add_argument('--target-modules', type=str, nargs='+', default=['Wqkv', 'q_proj', 'kv_proj', 'out_proj'],
+                        help="Target modules for LoRA fine tune")
+    parser.add_argument("--output-merge-path", type=str, default=None, help="Output path for merged model")
+
+    # controlnet config
+    parser.add_argument("--control-type", type=str, default='canny', choices=['canny', 'depth', 'pose'],
+                        help="Controlnet condition type")
+    parser.add_argument("--control-weight", type=str, default=1.0,
+                        help="Controlnet weight, You can use a float to specify the weight for all layers, "
+                             "or use a list to separately specify the weight for each layer, for example, "
+                             "'[1.0 * (0.825 ** float(19 - i)) for i in range(19)]'")
+    parser.add_argument("--condition-image-path", type=str, default=None, help="Inference condition image path")
+
+    # Diffusion
+    parser.add_argument("--learn-sigma", action="store_true", help="Learn extra channels for sigma.")
+    parser.add_argument("--no-learn-sigma", dest="learn_sigma", action="store_false")
+    parser.set_defaults(learn_sigma=True)
+    parser.add_argument("--predict-type", type=str, choices=list(PREDICT_TYPE), default="v_prediction",
+                        help="Diffusion predict type")
+    parser.add_argument("--noise-schedule", type=str, choices=list(NOISE_SCHEDULES), default="scaled_linear",
+                        help="Noise schedule")
+    parser.add_argument("--beta-start", type=float, default=0.00085, help="Beta start value")
+    parser.add_argument("--beta-end", type=float, default=0.02, help="Beta end value")
+    parser.add_argument("--sigma-small", action="store_true")
+    parser.add_argument("--mse-loss-weight-type", type=str, default="constant",
+                        help="Min-SNR-gamma. Can be constant or min_snr_<gamma> where gamma is a integer. "
+                             "5 is recommended in the paper.")
+    parser.add_argument("--model-var-type", type=model_var_type, default=None, help="Specify the model variable type.")
+    parser.add_argument("--noise-offset", type=float, default=0.0, help="Add extra noise to the input image.")
+
+    # ========================================================================================================
+    # Inference
+    # ========================================================================================================
+
+    # Basic Setting
+    parser.add_argument("--prompt", type=str, default="一只小猫", help="The prompt for generating images.")
+    parser.add_argument("--model-root", type=str, default="ckpts",
+                        help="Root path of all the models, including t2i model and dialoggen model.")
+    parser.add_argument("--dit-weight", type=str, default=None,
+                        help="Path to the HunYuan-DiT model. If None, search the model in the args.model_root."
+                             "1. If it is a directory, search the model in the directory. Support two types of models: "
+                             "1) named `pytorch_model_*.pt`, where * is specified by the --load-key. "
+                             "2) named `*_model_states.pt`, where * can be `mp_rank_00`. *_model_states.pt contains "
+                             "both 'module' and 'ema' weights. Therefore, you still use --load-key to specify the "
+                             "weights to load. By default, load 'ema' weights. "
+                             "2. If it is a file, load the model directly. In this case, the --load-key is ignored."
+                        )
+
+    # Model setting
+    parser.add_argument("--load-key", type=str, choices=["ema", "module", "distill", 'merge'], default="ema",
+                        help="Load model key for HunYuanDiT checkpoint.")
+    parser.add_argument('--use-style-cond', action="store_true",
+                        help="Use style condition in hydit. Only for hydit version <= 1.1")
+    parser.add_argument('--size-cond', type=int, nargs='+', default=None,
+                        help="Size condition used in sampling. 2 values are required for height and width. "
+                             "If a single value is provided, the image will be treated to (value, value)."
+                             "Recommended values are [1024, 1024]. Only for hydit version <= 1.1")
+    parser.add_argument('--target-ratios', type=str, nargs='+', default=None,
+                        help="Target ratios for multi-resolution training.")
+    parser.add_argument("--cfg-scale", type=float, default=6.0, help="Guidance scale for classifier-free.")
+    parser.add_argument("--negative", type=str, default=None, help="Negative prompt.")
+
+    # Acceleration
+    parser.add_argument("--infer-mode", type=str, choices=["fa", "torch", "trt"], default="fa",
+                        help="Inference mode")
+    parser.add_argument("--onnx-workdir", type=str, default="onnx_model", help="Path to save ONNX model")
+
+    # Sampling
+    parser.add_argument("--sampler", type=str, choices=SAMPLER_FACTORY, default="ddpm", help="Diffusion sampler")
+    parser.add_argument("--infer-steps", type=int, default=100, help="Inference steps")
+
+    # Prompt enhancement
+    parser.add_argument("--enhance", action="store_true", help="Enhance prompt with mllm.")
+    parser.add_argument("--no-enhance", dest="enhance", action="store_false")
+    parser.add_argument("--load-4bit", help="load DialogGen model with 4bit quantization.", action="store_true")
+    parser.set_defaults(enhance=True)
+
+    # App
+    parser.add_argument("--lang", type=str, default="zh", choices=["zh", "en"], help="Language")
+
+
+    # ========================================================================================================
+    # Training
+    # ========================================================================================================
+
+    # Basic Setting
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--max-training-steps", type=int, default=10_000_000)
+    parser.add_argument("--gc-interval", type=int, default=40,
+                        help='To address the memory bottleneck encountered during the preprocessing of the dataset,'
+                             ' memory fragments are reclaimed here by invoking the gc.collect() function.')
+    parser.add_argument("--log-every", type=int, default=100)
+    parser.add_argument("--ckpt-every", type=int, default=100_000, help="Create a ckpt every a few steps.")
+    parser.add_argument("--ckpt-latest-every", type=int, default=10_000,
+                        help="Create a ckpt named `latest.pt` every a few steps.")
+    parser.add_argument("--ckpt-every-n-epoch", type=int, default=0,
+                        help="Create a ckpt every a few epochs. If 0, do not create ckpt based on epoch. Default is 0.")
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--global-seed", type=int, default=1234)
+    parser.add_argument("--warmup-min-lr", type=float, default=1e-6)
+    parser.add_argument("--warmup-num-steps", type=float, default=0)
+    parser.add_argument("--weight-decay", type=float, default=0, help="weight-decay in optimizer")
+    parser.add_argument("--rope-img", type=str, default=None, choices=['extend', 'base512', 'base1024'],
+                        help="Extend or interpolate the positional embedding of the image.")
+    parser.add_argument("--rope-real", action="store_true",
+                        help="Use real part and imaginary part separately for RoPE.")
+
+    # Classifier-free
+    parser.add_argument("--uncond-p", type=float, default=0.2,
+                        help="The probability of dropping training text used for CLIP feature extraction")
+    parser.add_argument("--uncond-p-t5", type=float, default=0.2,
+                        help="The probability of dropping training text used for mT5 feature extraction")
+
+    # Directory
+    parser.add_argument("--results-dir", type=str, default="results")
+    parser.add_argument("--resume", type=str, default=None, help="Resume experiment from a checkpoint")
+    parser.add_argument("--no-strict", dest="strict", action="store_false", help="Strict loading of checkpoint")
+    parser.set_defaults(strict=True)
+    parser.add_argument("--resume-deepspeed", action="store_true",
+                        help="Resume model and ema states from a checkpoint saved by Deepspeed version of DIT.")
+    parser.add_argument("--resume-split", action="store_true",
+                        help="Resume model and ema states from two checkpoint separated from DeepSpeed ckpt.")
+    parser.add_argument("--ema-to-module", action="store_true",
+                        help="If true, initialize the module with EMA weights.")
+    parser.add_argument("--module-to-ema", action="store_true",
+                        help="if true, initialize the ema with Module weights.")
+
+    # Dataset
+    parser.add_argument("--index-file", type=str, nargs='+', help="During training, provide a JSON file with data "
+                                                                  "indices.")
+    parser.add_argument("--random-flip", action="store_true", help="Random flip image")
+    parser.add_argument("--reset-loader", action="store_true",
+                        help="Reset the data loader. It is useful when resuming from a checkpoint but switch to a new "
+                             "dataset.")
+    parser.add_argument("--multireso", action="store_true", help="Use multi-resolution training.")
+    parser.add_argument("--reso-step", type=int, default=None, help="Step size for multi-resolution training.")
+
+    # Additional condition
+    parser.add_argument("--random-shrink-size-cond", action="store_true",
+                        help="Randomly shrink the original size condition.")
+    parser.add_argument("--merge-src-cond", action="store_true", help="Merge the source condition into a single value.")
+
+    # EMA Model
+    parser.add_argument("--use-ema", action="store_true", help="Use EMA model")
+    parser.add_argument("--ema-dtype", type=str, choices=['fp16', 'fp32', 'none'], default="none",
+                        help="EMA data type. If none, use the same data type as the model.")
+    parser.add_argument("--ema-decay", type=float, default=None,
+                        help="EMA decay rate. If None, use the default value of the model.")
+    parser.add_argument("--ema-warmup", action="store_true",
+                        help="EMA warmup. If True, perform ema_decay warmup from 0 to ema_decay.")
+    parser.add_argument("--ema-warmup-power", type=float, default=None,
+                        help="EMA power. If None, use the default value of the model.")
+    parser.add_argument("--ema-reset-decay", action="store_true",
+                        help="Reset EMA decay to 0 and restart increasing the EMA decay."
+                             "Only works when --ema-warmup is enabled.")
+    # Acceleration
+    parser.add_argument("--use-flash-attn", action="store_true", help="During training, "
+                                                                      "flash attention is used to accelerate training.")
+    parser.add_argument("--no-flash-attn", dest="use_flash_attn", action="store_false",
+                        help="During training, flash attention is not used to accelerate training.")
+    parser.add_argument("--use-zero-stage", type=int, default=1, help="Use AngelPTM zero stage. Support 2 and 3")
+    parser.add_argument("--grad-accu-steps", type=int, default=1, help="Gradient accumulation steps.")
+    parser.add_argument("--gradient-checkpointing", action="store_true", help="Use gradient checkpointing.")
+    parser.add_argument("--save-optimizer-state", action="store_true", help="Save optimizer state in the checkpoint.")
+
+    # ========================================================================================================
+    # Deepspeed config
+    # ========================================================================================================
+    #parser = deepspeed.add_config_arguments(parser)
+    parser.add_argument('--local_rank', type=int, default=None,
+                        help='local rank passed from distributed launcher.')
+    parser.add_argument('--deepspeed-optimizer', action='store_true',
+                        help='Switching to the optimizers in DeepSpeed')
+    parser.add_argument('--remote-device', type=str, default='none', choices=['none', 'cpu', 'nvme'],
+                        help='Remote device for ZeRO-3 initialized parameters.')
+    parser.add_argument('--zero-stage', type=int, default=1)
+
+    args = parser.parse_known_args()
+
+    return args

custom_nodes/comfyui-hydit_min/hydit/constants.py

@@ -0,0 +1,82 @@
+import torch
+
+# =======================================================
+
+NOISE_SCHEDULES = {
+    "linear",
+    "scaled_linear",
+    "squaredcos_cap_v2",
+}
+
+PREDICT_TYPE = {
+    "epsilon",
+    "sample",
+    "v_prediction",
+}
+
+# =======================================================
+
+NEGATIVE_PROMPT = '错误的眼睛，糟糕的人脸，毁容，糟糕的艺术，变形，多余的肢体，模糊的颜色，模糊，重复，病态，残缺，'
+
+# =======================================================
+TRT_MAX_BATCH_SIZE = 1
+TRT_MAX_WIDTH = 1280
+TRT_MAX_HEIGHT = 1280
+
+# =======================================================
+# Constants about models
+# =======================================================
+
+VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix"
+TOKENIZER = "ckpts/t2i/tokenizer"
+TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder'
+T5_ENCODER = {
+    'MT5': 'ckpts/t2i/mt5',
+    'attention_mask': True,
+    'layer_index': -1,
+    'attention_pool': True,
+    'torch_dtype': torch.float16,
+    'learnable_replace': True
+}
+
+SAMPLER_FACTORY = {
+    'ddpm': {
+        'scheduler': 'DDPMScheduler',
+        'name': 'DDPM',
+        'kwargs': {
+            'steps_offset': 1,
+            'clip_sample': False,
+            'clip_sample_range': 1.0,
+            'beta_schedule': 'scaled_linear',
+            'beta_start': 0.00085,
+            'beta_end': 0.02,
+            'prediction_type': 'v_prediction',
+        }
+    },
+    'ddim': {
+        'scheduler': 'DDIMScheduler',
+        'name': 'DDIM',
+        'kwargs': {
+            'steps_offset': 1,
+            'clip_sample': False,
+            'clip_sample_range': 1.0,
+            'beta_schedule': 'scaled_linear',
+            'beta_start': 0.00085,
+            'beta_end': 0.02,
+            'prediction_type': 'v_prediction',
+        }
+    },
+    'dpmms': {
+        'scheduler': 'DPMSolverMultistepScheduler',
+        'name': 'DPMMS',
+        'kwargs': {
+            'beta_schedule': 'scaled_linear',
+            'beta_start': 0.00085,
+            'beta_end': 0.02,
+            'prediction_type': 'v_prediction',
+            'trained_betas': None,
+            'solver_order': 2,
+            'algorithm_type': 'dpmsolver++',
+        }
+    },
+}

custom_nodes/comfyui-hydit_min/hydit/diffusion/__init__.py

@@ -0,0 +1,49 @@
+from . import gaussian_diffusion as gd
+from .respace import SpacedDiffusion, space_timesteps
+
+
+def create_diffusion(
+    *,
+    steps=1000,
+    learn_sigma=True,
+    sigma_small=False,
+    noise_schedule="linear",
+    use_kl=False,
+    predict_type='epsilon',
+    rescale_timesteps=False,
+    rescale_learned_sigmas=False,
+    timestep_respacing="",
+    mse_loss_weight_type='constant',
+    beta_start=0.0001,
+    beta_end=0.02,
+    noise_offset=0.0,
+):
+    betas = gd.get_named_beta_schedule(noise_schedule, steps, beta_start, beta_end)
+    if use_kl:
+        loss_type = gd.LossType.RESCALED_KL
+    elif rescale_learned_sigmas:
+        loss_type = gd.LossType.RESCALED_MSE
+    else:
+        loss_type = gd.LossType.MSE
+    if timestep_respacing is None or timestep_respacing == "":
+        timestep_respacing = [steps]
+    mean_type = gd.predict_type_dict[predict_type]
+
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(steps, timestep_respacing),
+        betas=betas,
+        model_mean_type=mean_type,
+        model_var_type=(
+            (
+                gd.ModelVarType.FIXED_LARGE
+                if not sigma_small
+                else gd.ModelVarType.FIXED_SMALL
+            )
+            if not learn_sigma
+            else gd.ModelVarType.LEARNED_RANGE
+        ),
+        loss_type=loss_type,
+        rescale_timesteps=rescale_timesteps,
+        mse_loss_weight_type=mse_loss_weight_type,
+        noise_offset=noise_offset,
+    )

custom_nodes/comfyui-hydit_min/hydit/diffusion/diffusion_utils.py

@@ -0,0 +1,69 @@
+import torch as th
+import numpy as np
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs
+

custom_nodes/comfyui-hydit_min/hydit/diffusion/respace.py

@@ -0,0 +1,144 @@
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+    If the stride is a string starting with "ddim", then the fixed striding
+    from the DDIM paper is used, and only one section is allowed.
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    Improved DDPM
+
+    A diffusion process which can skip steps in a base diffusion process.
+    :param use_timesteps: a collection (sequence or set) of timesteps from the
+                          original diffusion process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps, **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(
+        self, model, controlnet=None, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        if controlnet != None:
+            return super().training_losses(self._wrap_model(model), controlnet=self._wrap_model(controlnet), *args, **kwargs)
+        else:
+            return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def get_eps(self, model, *args, **kwargs):
+        return super().get_eps(self._wrap_model(model), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(
+            model, self.timestep_map, self.original_num_steps
+        )
+
+    def _scale_timesteps(self, t):
+        # Scaling is done by the wrapped model.
+        return t
+
+
+class _WrappedModel:
+    """
+    Improved DDPM
+
+    When using a subsequent timesteps (e.g., 250), we must wrap the model
+    for mapping the timesteps 1-250 with step 1 to 1-1000 with step 4
+    """
+    def __init__(self, model, timestep_map, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        # self.rescale_timesteps = rescale_timesteps
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        """
+        Here we must make a interpolation because `ts` maybe a float (e.g., 4.5)
+        in the PLMS/PNDM sampler.
+        """
+        ts = ts.float()
+        frac = ts.frac()
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts_1 = map_tensor[ts.floor().long()]
+        new_ts_2 = map_tensor[ts.ceil().long()]
+        new_ts = th.lerp(new_ts_1, new_ts_2, frac)
+        return self.model(x, new_ts, **kwargs)

custom_nodes/comfyui-hydit_min/hydit/ds_config.py

@@ -0,0 +1,140 @@
+# -*- coding: utf-8 -*-
+import os
+
+def deepspeed_config_from_args(args, global_batch_size):
+    if args.use_zero_stage == 2:
+        deepspeed_config = {
+            "train_batch_size": global_batch_size,
+            "train_micro_batch_size_per_gpu": args.batch_size,
+            "gradient_accumulation_steps": args.grad_accu_steps,
+            "steps_per_print": args.log_every,
+            "optimizer": {
+                "type": "AdamW",
+                "params": {
+                    "lr": args.lr,
+                    "betas": [
+                        0.9,
+                        0.999
+                    ],
+                    "eps": 1e-08,
+                    "weight_decay": args.weight_decay
+                }
+            },
+
+            "zero_optimization": {
+                "stage": 2,
+                "reduce_scatter": False,
+                "reduce_bucket_size": 1e9,
+            },
+
+            "gradient_clipping": 1.0,
+            "prescale_gradients": True,
+
+            "fp16": {
+                "enabled": args.use_fp16,
+                "loss_scale": 0,
+                "loss_scale_window": 500,
+                "hysteresis": 2,
+                "min_loss_scale": 1e-3,
+                "initial_scale_power": 15
+            },
+
+            "bf16": {
+                "enabled": False
+            },
+
+            "wall_clock_breakdown": False
+        }
+
+    elif args.use_zero_stage == 3:
+        deepspeed_config = {
+            "train_batch_size": args.global_batch_size,
+            # "train_micro_batch_size_per_gpu": args.batch_size,
+            "gradient_accumulation_steps": args.grad_accu_steps,
+            "steps_per_print": args.log_every,
+
+            "optimizer": {
+                "type": "AdamW",
+                "params": {
+                    "lr": args.lr,
+                    "betas": [
+                        0.9,
+                        0.999
+                    ],
+                    "eps": 1e-08,
+                    "weight_decay": args.weight_decay
+                }
+            },
+
+            "zero_optimization": {
+                "stage": 3,
+                "allgather_partitions": True,
+                "overlap_comm": True,
+                "reduce_scatter": True,
+                "contiguous_gradients": True,
+                "stage3_prefetch_bucket_size": 5e8,
+                "stage3_max_live_parameters" : 6e8,
+                "reduce_bucket_size": 1.2e9,
+                "sub_group_size": 1e9,
+                "sub_group_buffer_num": 10,
+                "pipeline_optimizer": True,
+                "max_contigous_event_size": 0,
+                "cache_sub_group_rate": 0.0,
+                "prefetch_cache_sub_group_rate": 1.0,
+                "max_contigous_params_size": -1,
+                "max_param_reduce_events": 0,
+                "stage3_param_persistence_threshold": 9e9,
+                "is_communication_time_profiling": False,
+                "save_large_model_multi_slice": True,
+                "use_fused_op_with_grad_norm_overflow": False,
+                "offload_optimizer": {
+                    "device": "cpu",
+                    "pin_memory": True
+                },
+                "offload_param": {
+                    "device": "cpu",
+                    "pin_memory": True
+                }
+            },
+
+            "gradient_clipping": 1.0,
+            "prescale_gradients": False,
+
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0,
+                "loss_scale_window": 500,
+                "hysteresis": 2,
+                "min_loss_scale": 1,
+                "initial_scale_power": 15
+            },
+
+            "bf16": {
+                "enabled": False
+            },
+
+            "wall_clock_breakdown": False,
+            "mem_chunk": {
+                "default_chunk_size": 536870911,
+                "use_fake_dist": False,
+                "client": {
+                    "mem_tracer": {
+                        "use_async_mem_monitor": True,
+                        "warmup_gpu_chunk_mem_ratio": 0.8,
+                        "overall_gpu_mem_ratio": 0.8,
+                        "overall_cpu_mem_ratio": 1.0,
+                        "margin_use_ratio": 0.8,
+                        "use_fake_dist": False
+                    },
+                    "opts": {
+                        "with_mem_cache": True,
+                        "with_async_move": True
+                    }
+                }
+            }
+        }
+    else:
+        raise ValueError
+    return deepspeed_config
+
+

custom_nodes/comfyui-hydit_min/hydit/inference_comfyui.py

@@ -0,0 +1,510 @@
+import random
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+# For reproducibility
+# torch.backends.cudnn.benchmark = False
+# torch.backends.cudnn.deterministic = True
+
+from diffusers import schedulers
+from diffusers.models import AutoencoderKL
+from loguru import logger
+from transformers import BertModel, BertTokenizer
+from transformers.modeling_utils import logger as tf_logger
+
+from .constants import SAMPLER_FACTORY, NEGATIVE_PROMPT, TRT_MAX_WIDTH, TRT_MAX_HEIGHT, TRT_MAX_BATCH_SIZE
+from .diffusion.pipeline import StableDiffusionPipeline
+from .modules.models import HunYuanDiT, HUNYUAN_DIT_CONFIG
+from .modules.posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
+from .modules.text_encoder import MT5Embedder
+from .utils.tools import set_seeds
+from peft import LoraConfig
+
+
+class Resolution:
+    def __init__(self, width, height):
+        self.width = width
+        self.height = height
+
+    def __str__(self):
+        return f'{self.height}x{self.width}'
+
+
+class ResolutionGroup:
+    def __init__(self):
+        self.data = [
+            Resolution(1024, 1024), # 1:1
+            Resolution(1280, 1280), # 1:1
+            Resolution(1024, 768),  # 4:3
+            Resolution(1152, 864),  # 4:3
+            Resolution(1280, 960),  # 4:3
+            Resolution(768, 1024),  # 3:4
+            Resolution(864, 1152),  # 3:4
+            Resolution(960, 1280),  # 3:4
+            Resolution(1280, 768),  # 16:9
+            Resolution(768, 1280),  # 9:16
+        ]
+        self.supported_sizes = set([(r.width, r.height) for r in self.data])
+
+    def is_valid(self, width, height):
+        return (width, height) in self.supported_sizes
+
+
+STANDARD_RATIO = np.array([
+    1.0,        # 1:1
+    4.0 / 3.0,  # 4:3
+    3.0 / 4.0,  # 3:4
+    16.0 / 9.0, # 16:9
+    9.0 / 16.0, # 9:16
+])
+STANDARD_SHAPE = [
+    [(1024, 1024), (1280, 1280)],   # 1:1
+    [(1280, 960)],                # 4:3
+    [(960, 1280)],                   # 3:4
+    [(1280, 768)],                              # 16:9
+    [(768, 1280)],                              # 9:16
+]
+STANDARD_AREA = [
+    np.array([w * h for w, h in shapes])
+    for shapes in STANDARD_SHAPE
+]
+
+
+def get_standard_shape(target_width, target_height):
+    """
+    Map image size to standard size.
+    """
+    target_ratio = target_width / target_height
+    closest_ratio_idx = np.argmin(np.abs(STANDARD_RATIO - target_ratio))
+    closest_area_idx = np.argmin(np.abs(STANDARD_AREA[closest_ratio_idx] - target_width * target_height))
+    width, height = STANDARD_SHAPE[closest_ratio_idx][closest_area_idx]
+    return width, height
+
+
+def _to_tuple(val):
+    if isinstance(val, (list, tuple)):
+        if len(val) == 1:
+            val = [val[0], val[0]]
+        elif len(val) == 2:
+            val = tuple(val)
+        else:
+            raise ValueError(f"Invalid value: {val}")
+    elif isinstance(val, (int, float)):
+        val = (val, val)
+    else:
+        raise ValueError(f"Invalid value: {val}")
+    return val
+
+
+def get_pipeline(args, vae, text_encoder, tokenizer, model, device, rank,
+                 embedder_t5, infer_mode, sampler=None):
+    """
+    Get scheduler and pipeline for sampling. The sampler and pipeline are both
+    based on diffusers and make some modifications.
+
+    Returns
+    -------
+    pipeline: StableDiffusionPipeline
+    sampler_name: str
+    """
+    sampler = sampler or args.sampler
+
+    # Load sampler from factory
+    kwargs = SAMPLER_FACTORY[sampler]['kwargs']
+    scheduler = SAMPLER_FACTORY[sampler]['scheduler']
+
+    # Update sampler according to the arguments
+    kwargs['beta_schedule'] = args.noise_schedule
+    kwargs['beta_start'] = args.beta_start
+    kwargs['beta_end'] = args.beta_end
+    kwargs['prediction_type'] = args.predict_type
+
+    # Build scheduler according to the sampler.
+    scheduler_class = getattr(schedulers, scheduler)
+    scheduler = scheduler_class(**kwargs)
+    #print(scheduler)
+    #assert(0)
+
+    # Set timesteps for inference steps.
+    scheduler.set_timesteps(args.infer_steps, device)
+
+    # Only enable progress bar for rank 0
+    progress_bar_config = {} if rank == 0 else {'disable': True}
+
+    pipeline = StableDiffusionPipeline(vae=vae,
+                                       text_encoder=text_encoder,
+                                       tokenizer=tokenizer,
+                                       unet=model,
+                                       scheduler=scheduler,
+                                       feature_extractor=None,
+                                       safety_checker=None,
+                                       requires_safety_checker=False,
+                                       progress_bar_config=progress_bar_config,
+                                       embedder_t5=embedder_t5,
+                                       infer_mode=infer_mode,
+                                       )
+
+    pipeline = pipeline.to(device)
+
+    return pipeline, sampler
+
+
+class End2End(object):
+    def __init__(self, args, models_root_path, MODEL_PATH = None, VAE_PATH = None, LOAR_PATH = None):
+        self.args = args
+
+        # Check arguments
+        t2i_root_path = Path(models_root_path) / "t2i"
+        self.root = t2i_root_path
+        logger.info(f"Got text-to-image model root path: {t2i_root_path}")
+
+        # Set device and disable gradient
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        torch.set_grad_enabled(False)
+        # Disable BertModel logging checkpoint info
+        tf_logger.setLevel('ERROR')
+
+        # ========================================================================
+        model_dir = self.root / "model"
+
+        # ========================================================================
+        logger.info(f"Loading CLIP Text Encoder...")
+        text_encoder_path = self.root / "clip_text_encoder"
+        self.clip_text_encoder = BertModel.from_pretrained(str(text_encoder_path), False, revision=None).to(self.device)
+        logger.info(f"Loading CLIP Text Encoder finished")
+        #print(self.clip_text_encoder)
+
+        # ========================================================================
+        logger.info(f"Loading CLIP Tokenizer...")
+        tokenizer_path = self.root / "tokenizer"
+        self.tokenizer = BertTokenizer.from_pretrained(str(tokenizer_path))
+        logger.info(f"Loading CLIP Tokenizer finished")
+        #print(self.tokenizer)
+        #assert(0)
+
+        # ========================================================================
+        logger.info(f"Loading T5 Text Encoder and T5 Tokenizer...")
+        t5_text_encoder_path = self.root / 'mt5'
+        embedder_t5 = MT5Embedder(t5_text_encoder_path, torch_dtype=torch.float16, max_length=256)
+        self.embedder_t5 = embedder_t5
+        self.embedder_t5.model.to(self.device)  # Only move encoder to device
+        logger.info(f"Loading t5_text_encoder and t5_tokenizer finished")
+
+        # ========================================================================
+        logger.info(f"Loading VAE...")
+        if VAE_PATH:
+            vae_path = VAE_PATH
+            self.vae = AutoencoderKL.from_single_file(str(vae_path)).to(self.device)
+        else:
+            vae_path = self.root / "sdxl-vae-fp16-fix"
+            self.vae = AutoencoderKL.from_pretrained(str(vae_path)).to(self.device)
+
+
+        logger.info(f"Loading VAE finished")
+
+        # ========================================================================
+        # Create model structure and load the checkpoint
+        logger.info(f"Building HunYuan-DiT model...")
+        #print(self.args.model)
+        #print(self.args)
+        model_config = HUNYUAN_DIT_CONFIG[self.args.model]
+        self.patch_size = model_config['patch_size']
+        self.head_size = model_config['hidden_size'] // model_config['num_heads']
+        self.resolutions, self.freqs_cis_img = self.standard_shapes()   # Used for TensorRT models
+        self.image_size = _to_tuple(self.args.image_size)
+        latent_size = (self.image_size[0] // 8, self.image_size[1] // 8)
+
+        self.infer_mode = self.args.infer_mode
+        if self.infer_mode in ['fa', 'torch']:
+            if MODEL_PATH:
+                model_path = Path(MODEL_PATH)
+            else:
+                model_path = model_dir / f"pytorch_model_{self.args.load_key}.pt"
+
+            if not model_path.exists():
+                raise ValueError(f"model_path not exists: {model_path}")
+            # Build model structure
+            self.model = HunYuanDiT(self.args,
+                                    input_size=latent_size,
+                                    **model_config,
+                                    log_fn=logger.info,
+                                    ).half().to(self.device)    # Force to use fp16
+            # Load model checkpoint
+            logger.info(f"Loading model checkpoint {model_path}...")
+            #self.load_torch_weights()
+            state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
+            self.model.load_state_dict(state_dict)
+            if LOAR_PATH:
+                self.model.load_adapter(LOAR_PATH)
+                self.model.merge_and_unload()
+
+            self.model.eval()
+        elif self.infer_mode == 'trt':
+            raise NotImplementedError("TensorRT model is not supported yet.")
+        else:
+            raise ValueError(f"Unknown infer_mode: {self.infer_mode}")
+
+        # ========================================================================
+        # Build inference pipeline. We use a customized StableDiffusionPipeline.
+        logger.info(f"Loading inference pipeline...")
+        #self.pipeline, self.sampler = self.load_sampler()
+        logger.info(f'Loading pipeline finished')
+
+        # ========================================================================
+        self.default_negative_prompt = NEGATIVE_PROMPT
+        logger.info("==================================================")
+        logger.info(f"                Model is ready.                  ")
+        logger.info("==================================================")
+
+    def load_torch_weights(self):
+        load_key = self.args.load_key
+        if self.args.dit_weight is not None:
+            dit_weight = Path(self.args.dit_weight)
+            if dit_weight.is_dir():
+                files = list(dit_weight.glob("*.pt"))
+                if len(files) == 0:
+                    raise ValueError(f"No model weights found in {dit_weight}")
+                if str(files[0]).startswith('pytorch_model_'):
+                    model_path = dit_weight / f"pytorch_model_{load_key}.pt"
+                    bare_model = True
+                elif any(str(f).endswith('_model_states.pt') for f in files):
+                    files = [f for f in files if str(f).endswith('_model_states.pt')]
+                    model_path = files[0]
+                    if len(files) > 1:
+                        logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}")
+                    bare_model = False
+                else:
+                    raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: "
+                                     f"{list(map(str, files))}. When given a directory as --dit-weight, only "
+                                     f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and "
+                                     f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
+                                     f"specific weight file, please provide the full path to the file.")
+            elif dit_weight.is_file():
+                model_path = dit_weight
+                bare_model = 'unknown'
+            else:
+                raise ValueError(f"Invalid model path: {dit_weight}")
+        else:
+            model_dir = self.root / "model"
+            model_path = model_dir / f"pytorch_model_{load_key}.pt"
+            bare_model = True
+
+        if not model_path.exists():
+            raise ValueError(f"model_path not exists: {model_path}")
+        logger.info(f"Loading torch model {model_path}...")
+        if model_path.suffix == '.safetensors':
+            raise NotImplementedError(f"Loading safetensors is not supported yet.")
+        else:
+            # Assume it's a single weight file in the *.pt format.
+            state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
+
+        if bare_model == 'unknown' and ('ema' in state_dict or 'module' in state_dict):
+            bare_model = False
+        if bare_model is False:
+            if load_key in state_dict:
+                state_dict = state_dict[load_key]
+            else:
+                raise KeyError(f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint "
+                               f"are: {list(state_dict.keys())}.")
+
+        if 'style_embedder.weight' in state_dict and not hasattr(self.model, 'style_embedder'):
+            raise ValueError(f"You might be attempting to load the weights of HunYuanDiT version <= 1.1. You need "
+                             f"to set `--use-style-cond --size-cond 1024 1024 --beta-end 0.03` to adapt to these weights."
+                             f"Alternatively, you can use weights of version >= 1.2, which no longer depend on "
+                             f"these two parameters.")
+        if 'style_embedder.weight' not in state_dict and hasattr(self.model, 'style_embedder'):
+            raise ValueError(f"You might be attempting to load the weights of HunYuanDiT version >= 1.2. You need "
+                             f"to remove `--use-style-cond` and `--size-cond 1024 1024` to adapt to these weights.")
+        # Don't set strict=False. Always explicitly check the state_dict.
+        self.model.load_state_dict(state_dict, strict=True)
+
+    def load_sampler(self, sampler=None):
+        pipeline, sampler = get_pipeline(self.args,
+                                         self.vae,
+                                         self.clip_text_encoder,
+                                         self.tokenizer,
+                                         self.model,
+                                         device=self.device,
+                                         rank=0,
+                                         embedder_t5=self.embedder_t5,
+                                         infer_mode=self.infer_mode,
+                                         sampler=sampler,
+                                         )
+        return pipeline, sampler
+
+    def calc_rope(self, height, width):
+        th = height // 8 // self.patch_size
+        tw = width // 8 // self.patch_size
+        base_size = 512 // 8 // self.patch_size
+
+        start, stop = get_fill_resize_and_crop((th, tw), base_size)
+        sub_args = [start, stop, (th, tw)]
+        #import pdb
+        #pdb.set_trace()
+        rope = get_2d_rotary_pos_embed(self.head_size, *sub_args)
+        return rope
+
+    def standard_shapes(self):
+        resolutions = ResolutionGroup()
+        freqs_cis_img = {}
+        for reso in resolutions.data:
+            freqs_cis_img[str(reso)] = self.calc_rope(reso.height, reso.width)
+        return resolutions, freqs_cis_img
+
+    def predict(self,
+                user_prompt,
+                image, 
+                height=1024,
+                width=1024,
+                seed=None,
+                enhanced_prompt=None,
+                negative_prompt=None,
+                infer_steps=100,
+                guidance_scale=6,
+                batch_size=1,
+                src_size_cond=None,
+                sampler=None,
+                control_weight = 1.0,
+                use_style_cond=False,
+                ):
+        # ========================================================================
+        # Arguments: seed
+        # ========================================================================
+        if seed is None:
+            seed = random.randint(0, 1_000_000)
+        if not isinstance(seed, int):
+            raise TypeError(f"`seed` must be an integer, but got {type(seed)}")
+        generator = set_seeds(seed, device=self.device)
+
+        # ========================================================================
+        # Arguments: target_width, target_height
+        # ========================================================================
+        if width <= 0 or height <= 0:
+            raise ValueError(f"`height` and `width` must be positive integers, got height={height}, width={width}")
+        logger.info(f"Input (height, width) = ({height}, {width})")
+        if self.infer_mode in ['fa', 'torch']:
+            # We must force height and width to align to 16 and to be an integer.
+            target_height = int((height // 16) * 16)
+            target_width = int((width // 16) * 16)
+            logger.info(f"Align to 16: (height, width) = ({target_height}, {target_width})")
+        elif self.infer_mode == 'trt':
+            target_width, target_height = get_standard_shape(width, height)
+            logger.info(f"Align to standard shape: (height, width) = ({target_height}, {target_width})")
+        else:
+            raise ValueError(f"Unknown infer_mode: {self.infer_mode}")
+
+        # ========================================================================
+        # Arguments: prompt, new_prompt, negative_prompt
+        # ========================================================================
+        if not isinstance(user_prompt, str):
+            raise TypeError(f"`user_prompt` must be a string, but got {type(user_prompt)}")
+        user_prompt = user_prompt.strip()
+        prompt = user_prompt
+
+        if enhanced_prompt is not None:
+            if not isinstance(enhanced_prompt, str):
+                raise TypeError(f"`enhanced_prompt` must be a string, but got {type(enhanced_prompt)}")
+            enhanced_prompt = enhanced_prompt.strip()
+            prompt = enhanced_prompt
+
+        # negative prompt
+        if negative_prompt is None or negative_prompt == '':
+            negative_prompt = self.default_negative_prompt
+        if not isinstance(negative_prompt, str):
+            raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}")
+
+        # ========================================================================
+        # Arguments: style. (A fixed argument. Don't Change it.)
+        # ========================================================================
+        if use_style_cond:
+            # Only for hydit <= 1.1
+            style = torch.as_tensor([0, 0] * batch_size, device=self.device)
+        else:
+            style = None
+
+        # ========================================================================
+        # Inner arguments: image_meta_size (Please refer to SDXL.)
+        # ========================================================================
+
+        if src_size_cond is None:
+            size_cond = None
+            image_meta_size = None
+        else:
+            # Only for hydit <= 1.1
+            if isinstance(src_size_cond, int):
+                src_size_cond = [src_size_cond, src_size_cond]
+            if not isinstance(src_size_cond, (list, tuple)):
+                raise TypeError(f"`src_size_cond` must be a list or tuple, but got {type(src_size_cond)}")
+            if len(src_size_cond) != 2:
+                raise ValueError(f"`src_size_cond` must be a tuple of 2 integers, but got {len(src_size_cond)}")
+            size_cond = list(src_size_cond) + [target_width, target_height, 0, 0]
+            image_meta_size = torch.as_tensor([size_cond] * 2 * batch_size, device=self.device)
+
+        # ========================================================================
+        start_time = time.time()
+        logger.debug(f"""
+                       prompt: {user_prompt}
+              enhanced prompt: {enhanced_prompt}
+                         seed: {seed}
+              (height, width): {(target_height, target_width)}
+              negative_prompt: {negative_prompt}
+                   batch_size: {batch_size}
+               guidance_scale: {guidance_scale}
+                  infer_steps: {infer_steps}
+              image_meta_size: {size_cond}
+        """)
+        reso = f'{target_height}x{target_width}'
+        if reso in self.freqs_cis_img:
+            freqs_cis_img = self.freqs_cis_img[reso]
+        else:
+            freqs_cis_img = self.calc_rope(target_height, target_width)
+
+        #if sampler is not None and sampler != self.sampler:
+        #    self.pipeline, self.sampler = self.load_sampler(sampler)
+
+        if isinstance(image, torch.Tensor):
+            samples = self.pipeline(
+                height=target_height,
+                width=target_width,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_images_per_prompt=batch_size,
+                guidance_scale=guidance_scale,
+                num_inference_steps=infer_steps,
+                image_meta_size=image_meta_size,
+                style=style,
+                return_dict=False,
+                generator=generator,
+                freqs_cis_img=freqs_cis_img,
+                use_fp16=self.args.use_fp16,
+                learn_sigma=self.args.learn_sigma,
+                image=image,
+                control_weight=control_weight,
+            )[0]
+        else:
+            samples = self.pipeline(
+                height=target_height,
+                width=target_width,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_images_per_prompt=batch_size,
+                guidance_scale=guidance_scale,
+                num_inference_steps=infer_steps,
+                image_meta_size=image_meta_size,
+                style=style,
+                return_dict=False,
+                generator=generator,
+                freqs_cis_img=freqs_cis_img,
+                use_fp16=self.args.use_fp16,
+                learn_sigma=self.args.learn_sigma,
+            )[0]            
+        gen_time = time.time() - start_time
+        logger.debug(f"Success, time: {gen_time}")
+
+        return {
+            'images': samples,
+            'seed': seed,
+        }

custom_nodes/comfyui-hydit_min/hydit/lr_scheduler.py

@@ -0,0 +1,761 @@
+"""
+Implementation of learning rate schedules.
+
+Taken and modified from PyTorch v1.0.1 source
+https://github.com/pytorch/pytorch/blob/v1.1.0/torch/optim/lr_scheduler.py
+"""
+
+import argparse
+from torch.optim import Optimizer
+import math
+
+LR_SCHEDULE = 'lr_schedule'
+LR_RANGE_TEST = 'LRRangeTest'
+ONE_CYCLE = 'OneCycle'
+WARMUP_LR = 'WarmupLR'
+WARMUP_DECAY_LR = 'WarmupDecayLR'
+VALID_LR_SCHEDULES = [LR_RANGE_TEST, ONE_CYCLE, WARMUP_LR, WARMUP_DECAY_LR]
+
+LR_RANGE_TEST_MIN_LR = 'lr_range_test_min_lr'
+LR_RANGE_TEST_STEP_RATE = 'lr_range_test_step_rate'
+LR_RANGE_TEST_STEP_SIZE = 'lr_range_test_step_size'
+LR_RANGE_TEST_STAIRCASE = 'lr_range_test_staircase'
+
+EDGE_VALUE = 'edge_value'
+MID_VALUE = 'mid_value'
+
+CYCLE_FIRST_STEP_SIZE = 'cycle_first_step_size'
+CYCLE_FIRST_STAIR_COUNT = 'cycle_first_stair_count'
+CYCLE_SECOND_STEP_SIZE = 'cycle_second_step_size'
+CYCLE_SECOND_STAIR_COUNT = 'cycle_second_stair_count'
+DECAY_STEP_SIZE = 'decay_step_size'
+
+CYCLE_MIN_LR = 'cycle_min_lr'
+CYCLE_MAX_LR = 'cycle_max_lr'
+DECAY_LR_RATE = 'decay_lr_rate'
+
+CYCLE_MIN_MOM = 'cycle_min_mom'
+CYCLE_MAX_MOM = 'cycle_max_mom'
+DECAY_MOM_RATE = 'decay_mom_rate'
+
+WARMUP_MIN_LR = 'warmup_min_lr'
+WARMUP_MAX_LR = 'warmup_max_lr'
+WARMUP_NUM_STEPS = 'warmup_num_steps'
+WARMUP_TYPE = 'warmup_type'
+WARMUP_LOG_RATE = 'log'
+WARMUP_LINEAR_RATE = 'linear'
+
+TOTAL_NUM_STEPS = 'total_num_steps'
+
+
+def add_tuning_arguments(parser):
+    group = parser.add_argument_group('Convergence Tuning', 'Convergence tuning configurations')
+
+    # LR scheduler
+    group.add_argument('--lr_schedule', type=str, default=None, help='LR schedule for training.')
+
+    # Learning rate range test
+    group.add_argument("--lr_range_test_min_lr", type=float, default=0.001, help='Starting lr value.')
+    group.add_argument("--lr_range_test_step_rate", type=float, default=1.0, help='scaling rate for LR range test.')
+    group.add_argument("--lr_range_test_step_size", type=int, default=1000, help='training steps per LR change.')
+    group.add_argument("--lr_range_test_staircase",
+                       type=bool,
+                       default=False,
+                       help='use staircase scaling for LR range test.')
+
+    # OneCycle schedule
+    group.add_argument("--cycle_first_step_size",
+                       type=int,
+                       default=1000,
+                       help='size of first step of 1Cycle schedule (training steps).')
+    group.add_argument("--cycle_first_stair_count",
+                       type=int,
+                       default=-1,
+                       help='first stair count for 1Cycle schedule.')
+    group.add_argument("--cycle_second_step_size",
+                       type=int,
+                       default=-1,
+                       help='size of second step of 1Cycle schedule (default first_step_size).')
+    group.add_argument("--cycle_second_stair_count",
+                       type=int,
+                       default=-1,
+                       help='second stair count for 1Cycle schedule.')
+    group.add_argument("--decay_step_size",
+                       type=int,
+                       default=1000,
+                       help='size of intervals for applying post cycle decay (training steps).')
+
+    # 1Cycle LR
+    group.add_argument("--cycle_min_lr", type=float, default=0.01, help='1Cycle LR lower bound.')
+    group.add_argument("--cycle_max_lr", type=float, default=0.1, help='1Cycle LR upper bound.')
+    group.add_argument("--decay_lr_rate", type=float, default=0.0, help='post cycle LR decay rate.')
+
+    # 1Cycle Momentum
+    group.add_argument('--cycle_momentum', default=False, action='store_true', help='Enable 1Cycle momentum schedule.')
+    group.add_argument("--cycle_min_mom", type=float, default=0.8, help='1Cycle momentum lower bound.')
+    group.add_argument("--cycle_max_mom", type=float, default=0.9, help='1Cycle momentum upper bound.')
+    group.add_argument("--decay_mom_rate", type=float, default=0.0, help='post cycle momentum decay rate.')
+
+    # Warmup LR
+    group.add_argument('--warmup_min_lr', type=float, default=0, help='WarmupLR minimum/initial LR value')
+    group.add_argument('--warmup_max_lr', type=float, default=0.001, help='WarmupLR maximum LR value.')
+    group.add_argument('--warmup_num_steps', type=int, default=1000, help='WarmupLR step count for LR warmup.')
+    group.add_argument('--warmup_type',
+                       type=str,
+                       default=WARMUP_LOG_RATE,
+                       help='WarmupLR increasing function during warmup')
+    return parser
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser = add_tuning_arguments(parser)
+
+    lr_sched_args, unknown_args = parser.parse_known_args()
+    return lr_sched_args, unknown_args
+
+
+def override_lr_range_test_params(args, params):
+    if hasattr(args, LR_RANGE_TEST_MIN_LR) and args.lr_range_test_min_lr is not None:
+        params[LR_RANGE_TEST_MIN_LR] = args.lr_range_test_min_lr
+
+    if hasattr(args, LR_RANGE_TEST_STEP_RATE) and args.lr_range_test_step_rate is not None:
+        params[LR_RANGE_TEST_STEP_RATE] = args.lr_range_test_step_rate
+
+    if hasattr(args, LR_RANGE_TEST_STEP_SIZE) and args.lr_range_test_step_size is not None:
+        params[LR_RANGE_TEST_STEP_SIZE] = args.lr_range_test_step_size
+
+    if hasattr(args, LR_RANGE_TEST_STAIRCASE) and args.lr_range_test_staircase is not None:
+        params[LR_RANGE_TEST_STAIRCASE] = args.lr_range_test_staircase
+
+
+def override_1cycle_params(args, params):
+    if hasattr(args, CYCLE_FIRST_STEP_SIZE) and args.cycle_first_step_size is not None:
+        params[CYCLE_FIRST_STEP_SIZE] = args.cycle_first_step_size
+
+    if hasattr(args, CYCLE_FIRST_STAIR_COUNT) and args.cycle_first_stair_count is not None:
+        params[CYCLE_FIRST_STAIR_COUNT] = args.cycle_first_stair_count
+
+    if hasattr(args, CYCLE_SECOND_STEP_SIZE) and args.cycle_second_step_size is not None:
+        params[CYCLE_SECOND_STEP_SIZE] = args.cycle_second_step_size
+
+    if hasattr(args, CYCLE_SECOND_STAIR_COUNT) and args.cycle_second_stair_count is not None:
+        params[CYCLE_SECOND_STAIR_COUNT] = args.cycle_second_stair_count
+
+    if hasattr(args, DECAY_STEP_SIZE) and args.decay_step_size is not None:
+        params[DECAY_STEP_SIZE] = args.decay_step_size
+
+    # 1Cycle LR params
+    if hasattr(args, CYCLE_MIN_LR) and args.cycle_min_lr is not None:
+        params[CYCLE_MIN_LR] = args.cycle_min_lr
+
+    if hasattr(args, CYCLE_MAX_LR) and args.cycle_max_lr is not None:
+        params[CYCLE_MAX_LR] = args.cycle_max_lr
+
+    if hasattr(args, DECAY_LR_RATE) and args.decay_lr_rate is not None:
+        params[DECAY_LR_RATE] = args.decay_lr_rate
+
+    # 1Cycle MOM params
+    if hasattr(args, CYCLE_MIN_MOM) and args.cycle_min_mom is not None:
+        params[CYCLE_MIN_MOM] = args.cycle_min_mom
+
+    if hasattr(args, CYCLE_MAX_MOM) and args.cycle_max_mom is not None:
+        params[CYCLE_MAX_MOM] = args.cycle_max_mom
+
+    if hasattr(args, DECAY_MOM_RATE) and args.decay_mom_rate is not None:
+        params[DECAY_MOM_RATE] = args.decay_mom_rate
+
+
+def override_warmupLR_params(args, params):
+    if hasattr(args, WARMUP_MIN_LR) and args.warmup_min_lr is not None:
+        params[WARMUP_MIN_LR] = args.warmup_min_lr
+
+    if hasattr(args, WARMUP_MAX_LR) and args.warmup_max_lr is not None:
+        params[WARMUP_MAX_LR] = args.warmup_max_lr
+
+    if hasattr(args, WARMUP_NUM_STEPS) and args.warmup_num_steps is not None:
+        params[WARMUP_NUM_STEPS] = args.warmup_num_steps
+
+    if hasattr(args, WARMUP_TYPE) and args.warmup_type is not None:
+        params[WARMUP_TYPE] = args.warmup_type
+
+
+def override_params(args, params):
+    # LR range test params
+    override_lr_range_test_params(args, params)
+
+    # 1Cycle params
+    override_1cycle_params(args, params)
+
+    # WarmupLR params
+    override_warmupLR_params(args, params)
+
+
+def get_config_from_args(args):
+    if not hasattr(args, LR_SCHEDULE) or args.lr_schedule is None:
+        return None, '--{} not specified on command line'.format(LR_SCHEDULE)
+
+    if not args.lr_schedule in VALID_LR_SCHEDULES:
+        return None, '{} is not supported LR schedule'.format(args.lr_schedule)
+
+    config = {}
+    config['type'] = args.lr_schedule
+    config['params'] = {}
+
+    if args.lr_schedule == LR_RANGE_TEST:
+        override_lr_range_test_params(args, config['params'])
+    elif args.lr_schedule == ONE_CYCLE:
+        override_1cycle_params(args, config['params'])
+    else:
+        override_warmupLR_params(args, config['params'])
+
+    return config, None
+
+
+def get_lr_from_config(config):
+    if not 'type' in config:
+        return None, 'LR schedule type not defined in config'
+
+    if not 'params' in config:
+        return None, 'LR schedule params not defined in config'
+
+    lr_schedule = config['type']
+    lr_params = config['params']
+
+    if not lr_schedule in VALID_LR_SCHEDULES:
+        return None, '{} is not a valid LR schedule'.format(lr_schedule)
+
+    if lr_schedule == LR_RANGE_TEST:
+        return lr_params[LR_RANGE_TEST_MIN_LR], ''
+    if lr_schedule == ONE_CYCLE:
+        return lr_params[CYCLE_MAX_LR], ''
+    # Warmup LR
+    return lr_params[WARMUP_MAX_LR], ''
+
+
+"""
+Only optimizers that are subclass of torch.optim.Optimizer are supported. So check the passed optimizer and wrapped
+optimizer to see if requirement is satisfied.
+TODO: Looking under the hood to examine the wrapped optimizer is a hack that requires a better long-term fix.
+"""
+
+
+def get_torch_optimizer(optimizer):
+    if isinstance(optimizer, Optimizer):
+        return optimizer
+
+    if hasattr(optimizer, 'optimizer') and isinstance(optimizer.optimizer, Optimizer):
+        return optimizer.optimizer
+
+    raise TypeError('{} is not a subclass of torch.optim.Optimizer'.format(type(optimizer).__name__))
+
+
+class LRRangeTest(object):
+    """Sets the learning rate of each parameter group according to
+    learning rate range test (LRRT) policy. The policy increases learning
+    rate starting from a base value with a constant frequency, as detailed in
+    the paper `A disciplined approach to neural network hyper-parameters: Part1`_.
+
+    LRRT policy is used for finding maximum LR that trains a model without divergence, and can be used to
+    configure the LR boundaries for Cyclic LR schedules.
+
+    LRRT changes the learning rate after every batch.
+    `step` should be called after a batch has been used for training.
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        lr_range_test_min_lr (float or list): Initial learning rate which is the
+            lower boundary in the range test for each parameter group.
+        lr_range_test_step_size (int): Interval of training steps to increase learning rate. Default: 2000
+        lr_range_test_step_rate (float): Scaling rate for range test. Default: 1.0
+        lr_range_test_staircase (bool): Scale in staircase fashion, rather than continuous. Default: False.
+        last_batch_iteration (int): The index of the last batch. This parameter is used when
+            resuming a training job. Since `step()` should be invoked after each
+            batch instead of after each epoch, this number represents the total
+            number of *batches* computed, not the total number of epochs computed.
+            When last_batch_iteration=-1, the schedule is started from the beginning.
+            Default: -1
+
+    Example:
+        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+        >>> scheduler = LRRangeTest(optimizer)
+        >>> data_loader = torch.utils.data.DataLoader(...)
+        >>> for epoch in range(10):
+        >>>     for batch in data_loader:
+        >>>         train_batch(...)
+        >>>         scheduler.step()
+
+        _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay:
+        https://arxiv.org/abs/1803.09820
+"""
+
+    def __init__(self,
+                 optimizer: Optimizer,
+                 lr_range_test_min_lr: float = 1e-3,
+                 lr_range_test_step_size: int = 2000,
+                 lr_range_test_step_rate: float = 1.0,
+                 lr_range_test_staircase: bool = False,
+                 last_batch_iteration: int = -1):
+
+        self.optimizer = get_torch_optimizer(optimizer)
+
+        if isinstance(lr_range_test_min_lr, list) or isinstance(lr_range_test_min_lr, tuple):
+            if len(lr_range_test_min_lr) != len(self.optimizer.param_groups):
+                raise ValueError("expected {} lr_range_test_min_lr, got {}".format(len(self.optimizer.param_groups),
+                                                                                   len(lr_range_test_min_lr)))
+            self.min_lr = list(lr_range_test_min_lr)
+        else:
+            self.min_lr = [lr_range_test_min_lr] * len(self.optimizer.param_groups)
+
+        self.step_size = lr_range_test_step_size
+        self.step_rate = lr_range_test_step_rate
+        self.last_batch_iteration = last_batch_iteration
+        self.staircase = lr_range_test_staircase
+        self.interval_fn = self._staircase_interval if lr_range_test_staircase else self._continuous_interval
+
+        if last_batch_iteration == -1:
+            self._update_optimizer(self.min_lr)
+
+    def _staircase_interval(self):
+        return math.floor(float(self.last_batch_iteration + 1) / self.step_size)
+
+    def _continuous_interval(self):
+        return float(self.last_batch_iteration + 1) / self.step_size
+
+    def _get_increase(self):
+        return (1 + self.step_rate * self.interval_fn())
+
+    def get_lr(self):
+        lr_increase = self._get_increase()
+        return [lr_range_test_min_lr * lr_increase for lr_range_test_min_lr in self.min_lr]
+
+    def get_last_lr(self):
+        """ Return last computed learning rate by current scheduler.
+        """
+        assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
+        return self._last_lr
+
+    def _update_optimizer(self, group_lrs):
+        for param_group, lr in zip(self.optimizer.param_groups, group_lrs):
+            param_group['lr'] = lr
+
+    def step(self, batch_iteration=None):
+        if batch_iteration is None:
+            batch_iteration = self.last_batch_iteration + 1
+        self.last_batch_iteration = batch_iteration
+        self._update_optimizer(self.get_lr())
+        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
+
+    def state_dict(self):
+        return {'last_batch_iteration': self.last_batch_iteration}
+
+    def load_state_dict(self, sd):
+        self.last_batch_iteration = sd['last_batch_iteration']
+
+
+
+class OneCycle(object):
+    """Sets the learning rate of each parameter group according to
+    1Cycle learning rate policy (1CLR). 1CLR is a variation of the
+    Cyclical Learning Rate (CLR) policy that involves one cycle followed by
+    decay. The policy simultaneously cycles the learning rate (and momentum)
+    between two boundaries with a constant frequency, as detailed in
+    the paper `A disciplined approach to neural network hyper-parameters`_.
+
+    1CLR policy changes the learning rate after every batch.
+    `step` should be called after a batch has been used for training.
+
+    This implementation was adapted from the github repo: `pytorch/pytorch`_
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        cycle_min_lr (float or list): Initial learning rate which is the
+            lower boundary in the cycle for each parameter group.
+        cycle_max_lr (float or list): Upper learning rate boundaries in the cycle
+            for each parameter group. Functionally,
+            it defines the cycle amplitude (cycle_max_lr - cycle_min_lr).
+            The lr at any cycle is the sum of cycle_min_lr
+            and some scaling of the amplitude; therefore
+            cycle_max_lr may not actually be reached depending on
+            scaling function.
+        decay_lr_rate(float): Decay rate for learning rate. Default: 0.
+        cycle_first_step_size (int): Number of training iterations in the
+            increasing half of a cycle. Default: 2000
+        cycle_second_step_size (int): Number of training iterations in the
+            decreasing half of a cycle. If cycle_second_step_size is None,
+            it is set to cycle_first_step_size. Default: None
+        cycle_first_stair_count(int): Number of stairs in first half of cycle phase. This means
+        lr/mom are changed in staircase fashion. Default 0, means staircase disabled.
+        cycle_second_stair_count(int): Number of stairs in second half of cycle phase. This means
+        lr/mom are changed in staircase fashion. Default 0, means staircase disabled.
+        decay_step_size (int): Intervals for applying decay in decay phase. Default: 0, means no decay.
+        cycle_momentum (bool): If ``True``, momentum is cycled inversely
+            to learning rate between 'cycle_min_mom' and 'cycle_max_mom'.
+            Default: True
+        cycle_min_mom (float or list): Initial momentum which is the
+            lower boundary in the cycle for each parameter group.
+            Default: 0.8
+        cycle_max_mom (float or list): Upper momentum boundaries in the cycle
+            for each parameter group. Functionally,
+            it defines the cycle amplitude (cycle_max_mom - cycle_min_mom).
+            The momentum at any cycle is the difference of cycle_max_mom
+            and some scaling of the amplitude; therefore
+            cycle_min_mom may not actually be reached depending on
+            scaling function. Default: 0.9
+        decay_mom_rate (float): Decay rate for momentum. Default: 0.
+        last_batch_iteration (int): The index of the last batch. This parameter is used when
+            resuming a training job. Since `step()` should be invoked after each
+            batch instead of after each epoch, this number represents the total
+            number of *batches* computed, not the total number of epochs computed.
+            When last_batch_iteration=-1, the schedule is started from the beginning.
+            Default: -1
+
+    Example:
+        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+        >>> scheduler = OneCycle(optimizer, 0.0001, 0.0010)
+        >>> data_loader = torch.utils.data.DataLoader(...)
+        >>> for epoch in range(10):
+        >>>     for batch in data_loader:
+        >>>         train_batch(...)
+        >>>         scheduler.step()
+
+
+    .. _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: https://arxiv.org/abs/1803.09820
+    """
+
+    def __init__(self,
+                 optimizer,
+                 cycle_min_lr,
+                 cycle_max_lr,
+                 decay_lr_rate=0.,
+                 cycle_first_step_size=2000,
+                 cycle_second_step_size=None,
+                 cycle_first_stair_count=0,
+                 cycle_second_stair_count=None,
+                 decay_step_size=0,
+                 cycle_momentum=True,
+                 cycle_min_mom=0.8,
+                 cycle_max_mom=0.9,
+                 decay_mom_rate=0.,
+                 last_batch_iteration=-1):
+
+        self.optimizer = get_torch_optimizer(optimizer)
+
+        # Initialize cycle shape
+        self._initialize_cycle(cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count,
+                               cycle_second_stair_count, decay_step_size)
+
+        # Initialize cycle lr
+        self._initialize_lr(self.optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration)
+
+        # Initialize cyclic momentum
+        self.cycle_momentum = cycle_momentum
+        if cycle_momentum:
+            self._initialize_momentum(self.optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate,
+                                      last_batch_iteration)
+
+        # Initialize batch iteration tracker
+        self.last_batch_iteration = last_batch_iteration
+
+    # Configure cycle shape
+
+    def _initialize_cycle(self, cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count,
+                          cycle_second_stair_count, decay_step_size):
+        cycle_first_step_size = float(cycle_first_step_size)
+        cycle_second_step_size = float(
+            cycle_second_step_size) if cycle_second_step_size is not None else cycle_first_step_size
+
+        self.total_size = cycle_first_step_size + cycle_second_step_size
+        self.step_ratio = cycle_first_step_size / self.total_size
+        self.first_stair_count = cycle_first_stair_count
+        self.second_stair_count = cycle_first_stair_count if cycle_second_stair_count is None else cycle_second_stair_count
+        self.decay_step_size = decay_step_size
+
+        if math.isclose(self.decay_step_size, 0):
+            self.skip_lr_decay = True
+            self.skip_mom_decay = True
+        else:
+            self.skip_lr_decay = False
+            self.skip_mom_decay = False
+
+    # Configure lr schedule
+    def _initialize_lr(self, optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration):
+        self.min_lrs = [cycle_min_lr] * len(optimizer.param_groups)
+        if last_batch_iteration == -1:
+            for lr, group in zip(self.min_lrs, optimizer.param_groups):
+                group['lr'] = lr
+
+        self.max_lrs = [cycle_max_lr] * len(optimizer.param_groups)
+        self.decay_lr_rate = decay_lr_rate
+
+        if math.isclose(self.decay_lr_rate, 0):
+            self.skip_lr_decay = True
+
+    # Configure momentum schedule
+    def _initialize_momentum(self, optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate, last_batch_iteration):
+        if 'betas' not in optimizer.defaults:
+            optimizer_name = type(optimizer).__name__
+            print(
+                f"cycle_momentum is disabled because optimizer {optimizer_name} does not support momentum, no betas attribute in defaults"
+            )
+            self.cycle_momentum = False
+            return
+
+        self.decay_mom_rate = decay_mom_rate
+        self.min_moms = [(cycle_min_mom, 0.99)] * len(optimizer.param_groups)
+        self.max_moms = [(cycle_max_mom, 0.99)] * len(optimizer.param_groups)
+
+        if last_batch_iteration == -1:
+            for momentum, group in zip(self.min_moms, optimizer.param_groups):
+                group['betas'] = momentum
+
+        if math.isclose(self.decay_mom_rate, 0):
+            self.skip_mom_decay = True
+
+    def _get_scale_factor(self):
+        batch_iteration = (self.last_batch_iteration + 1)
+        cycle = math.floor(1 + batch_iteration / self.total_size)
+        x = 1. + batch_iteration / self.total_size - cycle
+        if x <= self.step_ratio:
+            scale_factor = x / self.step_ratio
+        else:
+            scale_factor = (x - 1) / (self.step_ratio - 1)
+
+        return scale_factor
+
+    def _get_cycle_mom(self):
+        scale_factor = self._get_scale_factor()
+        momentums = []
+        for base_betas, max_betas in zip(self.min_moms, self.max_moms):
+            cycle_min_mom = base_betas[0]
+            cycle_max_mom = max_betas[0]
+            base_height = (cycle_max_mom - cycle_min_mom) * scale_factor
+            momentum = cycle_max_mom - base_height
+            momentums.append((momentum, base_betas[1]))
+        return momentums
+
+    def _get_cycle_lr(self):
+        scale_factor = self._get_scale_factor()
+        lrs = []
+        for cycle_min_lr, cycle_max_lr in zip(self.min_lrs, self.max_lrs):
+            base_height = (cycle_max_lr - cycle_min_lr) * scale_factor
+            lr = cycle_min_lr + base_height
+            lrs.append(lr)
+
+        return lrs
+
+    def _get_decay_mom(self, decay_batch_iteration):
+        if self.skip_mom_decay:
+            return self.max_moms
+
+        decay_interval = decay_batch_iteration / self.decay_step_size
+        mom_decay_factor = (1 + self.decay_mom_rate * decay_interval)
+        momentums = [(beta0 * mom_decay_factor, beta1) for beta0, beta1 in self.max_moms]
+
+        return momentums
+
+    def _get_decay_lr(self, decay_batch_iteration):
+        """Calculates the learning rate at batch index. This function is used
+        after the cycle completes and post cycle decaying of lr/mom is enabled.
+        This function treats `self.last_batch_iteration` as the last batch index.
+        """
+        if self.skip_lr_decay:
+            return self.min_lrs
+
+        decay_interval = decay_batch_iteration / self.decay_step_size
+        lr_decay_factor = (1 + self.decay_lr_rate * decay_interval)
+        lrs = [cycle_min_lr / lr_decay_factor for cycle_min_lr in self.min_lrs]
+
+        return lrs
+
+    def get_lr(self):
+        """Calculates the learning rate at batch index. This function treats
+        `self.last_batch_iteration` as the last batch index.
+        """
+        if self.last_batch_iteration < self.total_size:
+            return self._get_cycle_lr()
+        return self._get_decay_lr(self.last_batch_iteration - self.total_size + 1)
+
+    def get_mom(self):
+        """Calculates the momentum at batch index. This function treats
+        `self.last_batch_iteration` as the last batch index.
+        """
+        if not self.cycle_momentum:
+            return None
+
+        if self.last_batch_iteration < self.total_size:
+            return self._get_cycle_mom()
+        return self._get_decay_mom(self.last_batch_iteration - self.total_size + 1)
+
+    def get_last_lr(self):
+        """ Return last computed learning rate by current scheduler.
+        """
+        assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
+        return self._last_lr
+
+    def step(self, batch_iteration=None):
+        """ Updates the optimizer with the learning rate for the last batch index.
+        `self.last_batch_iteration` is treated as the last batch index.
+
+        If self.cycle_momentum is true, also updates optimizer momentum.
+        """
+        if batch_iteration is None:
+            batch_iteration = self.last_batch_iteration + 1
+
+        self.last_batch_iteration = batch_iteration
+        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
+            param_group['lr'] = lr
+        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
+
+        if self.cycle_momentum:
+            momentums = self.get_mom()
+            for param_group, momentum in zip(self.optimizer.param_groups, momentums):
+                param_group['betas'] = momentum
+
+    def state_dict(self):
+        return {'last_batch_iteration': self.last_batch_iteration}
+
+    def load_state_dict(self, sd):
+        self.last_batch_iteration = sd['last_batch_iteration']
+
+
+
+class WarmupLR(object):
+    """Increase the learning rate of each parameter group from min lr to max lr
+        over warmup_num_steps steps, and then fix at max lr.
+
+        Args:
+            optimizer (Optimizer): Wrapped optimizer.
+            warmup_min_lr (float or list): minimum learning rate. Default: 0
+            warmup_max_lr (float or list): maximum learning rate. Default: 0.001
+            warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000
+            warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log
+            last_batch_iteration (int): The index of the last batch. Default: -1.
+        Example:
+            >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+            >>> scheduler = WarmupLR(optimizer)
+            >>> data_loader = torch.utils.data.DataLoader(...)
+            >>> for epoch in range(10):
+            >>>     for batch in data_loader:
+            >>>         train_batch(...)
+            >>>         scheduler.step()
+
+    """
+
+    def __init__(self,
+                 optimizer: Optimizer,
+                 warmup_min_lr: float = 0.0,
+                 warmup_max_lr: float = 0.001,
+                 warmup_num_steps: int = 1000,
+                 warmup_type: str = WARMUP_LOG_RATE,
+                 last_batch_iteration: int = -1):
+
+        self.optimizer = get_torch_optimizer(optimizer)
+
+        self.min_lrs = self._format_param(self.optimizer, warmup_min_lr, "min_lr")
+        self.max_lrs = self._format_param(self.optimizer, warmup_max_lr, "max_lr")
+        self.delta_lrs = [big - small for big, small in zip(self.max_lrs, self.min_lrs)]
+        self.warmup_num_steps = max(2, warmup_num_steps)
+        # Currently only support linear and log function
+        if warmup_type not in {WARMUP_LOG_RATE, WARMUP_LINEAR_RATE}:
+            print(f"Using unknown warmup_type: {warmup_type}. The increasing function "
+                           f"is set to default (log)")
+            warmup_type = WARMUP_LOG_RATE
+        self.warmup_type = warmup_type
+        self.inverse_log_warm_up = 1.0 / math.log(self.warmup_num_steps)
+        self.last_batch_iteration = last_batch_iteration
+
+    def get_lr(self):
+        if self.last_batch_iteration < 0:
+            print("Attempting to get learning rate from scheduler before it has started")
+            return [0.0]
+        gamma = self._get_gamma()
+        return [min_lr + (delta_lr * gamma) for min_lr, delta_lr in zip(self.min_lrs, self.delta_lrs)]
+
+    def get_last_lr(self):
+        """ Return last computed learning rate by current scheduler.
+        """
+        assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
+        return self._last_lr
+
+    def step(self, last_batch_iteration=None):
+        if last_batch_iteration is None:
+            last_batch_iteration = self.last_batch_iteration + 1
+        self.last_batch_iteration = last_batch_iteration
+        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
+            param_group['lr'] = lr
+        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
+
+    def state_dict(self):
+        return {'last_batch_iteration': self.last_batch_iteration}
+
+    def load_state_dict(self, sd):
+        self.last_batch_iteration = sd['last_batch_iteration']
+
+    def _get_gamma(self):
+        if self.last_batch_iteration < self.warmup_num_steps:
+            if self.warmup_type == WARMUP_LOG_RATE:
+                return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1)
+            elif self.warmup_type == WARMUP_LINEAR_RATE:
+                return self.last_batch_iteration / self.warmup_num_steps
+        return 1.0
+
+    def _format_param(self, optimizer, param_value, param_name):
+        if isinstance(param_value, list) or isinstance(param_value, tuple):
+            if len(param_value) != len(optimizer.param_groups):
+                raise ValueError("expected {} value for {}, got {}".format(len(optimizer.param_groups), param_name,
+                                                                           FileNotFoundError(param_value)))
+            return list(param_value)
+        return [param_value] * len(optimizer.param_groups)
+
+
+
+class WarmupDecayLR(WarmupLR):
+    """Increase the learning rate of each parameter group from min lr to max lr
+        over warmup_num_steps steps, and then decay at linear rate over the remaining training steps.
+
+        Args:
+            optimizer (Optimizer): Wrapped optimizer.
+            total_num_steps (int): total number of training steps
+            warmup_min_lr (float or list): minimum learning rate. Default: 0
+            warmup_max_lr (float or list): maximum learning rate. Default: 0.001
+            warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000
+            warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log
+            last_batch_iteration (int): The index of the last batch. Default: -1.
+        Example:
+            >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+            >>> scheduler = WarmupDecayLR(optimizer, 1000000)
+            >>> data_loader = torch.utils.data.DataLoader(...)
+            >>> for epoch in range(10):
+            >>>     for batch in data_loader:
+            >>>         train_batch(...)
+            >>>         scheduler.step()
+
+    """
+
+    def __init__(self,
+                 optimizer: Optimizer,
+                 total_num_steps: int,
+                 warmup_min_lr: float = 0.0,
+                 warmup_max_lr: float = 0.001,
+                 warmup_num_steps: int = 1000,
+                 warmup_type: str = WARMUP_LOG_RATE,
+                 last_batch_iteration: int = -1):
+
+        self.total_num_steps = total_num_steps
+        super(WarmupDecayLR, self).__init__(optimizer, warmup_min_lr, warmup_max_lr, warmup_num_steps, warmup_type,
+                                            last_batch_iteration)
+        if self.total_num_steps < self.warmup_num_steps:
+            print('total_num_steps {} is less than warmup_num_steps {}'.format(
+                total_num_steps, warmup_num_steps))
+
+    def _get_gamma(self):
+        if self.last_batch_iteration < self.warmup_num_steps:
+            if self.warmup_type == WARMUP_LOG_RATE:
+                return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1)
+            elif self.warmup_type == WARMUP_LINEAR_RATE:
+                return self.last_batch_iteration / self.warmup_num_steps
+        return max(
+            0.0,
+            float(self.total_num_steps - self.last_batch_iteration) /
+            float(max(1.0, self.total_num_steps - self.warmup_num_steps)))

custom_nodes/comfyui-hydit_min/hydit/modules/attn_layers.py

@@ -0,0 +1,401 @@
+import torch
+import torch.nn as nn
+from typing import Tuple, Union, Optional
+from comfy.ldm.modules.attention import optimized_attention
+
+try:
+    import flash_attn
+    if hasattr(flash_attn, '__version__') and int(flash_attn.__version__[0]) == 2:
+        from flash_attn.flash_attn_interface import flash_attn_kvpacked_func
+        from flash_attn.modules.mha import FlashSelfAttention, FlashCrossAttention
+    else:
+        from flash_attn.flash_attn_interface import flash_attn_unpadded_kvpacked_func
+        from flash_attn.modules.mha import FlashSelfAttention, FlashCrossAttention
+except Exception as e:
+    print(f'flash_attn import failed: {e}')
+
+
+def reshape_for_broadcast(freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], x: torch.Tensor, head_first=False):
+    """
+    Reshape frequency tensor for broadcasting it with another tensor.
+
+    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
+    for the purpose of broadcasting the frequency tensor during element-wise operations.
+
+    Args:
+        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped.
+        x (torch.Tensor): Target tensor for broadcasting compatibility.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        torch.Tensor: Reshaped frequency tensor.
+
+    Raises:
+        AssertionError: If the frequency tensor doesn't match the expected shape.
+        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
+    """
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+
+    if isinstance(freqs_cis, tuple):
+        # freqs_cis: (cos, sin) in real space
+        if head_first:
+            assert freqs_cis[0].shape == (x.shape[-2], x.shape[-1]), f'freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}'
+            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        else:
+            assert freqs_cis[0].shape == (x.shape[1], x.shape[-1]), f'freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}'
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
+    else:
+        # freqs_cis: values in complex space
+        if head_first:
+            assert freqs_cis.shape == (x.shape[-2], x.shape[-1]), f'freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}'
+            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        else:
+            assert freqs_cis.shape == (x.shape[1], x.shape[-1]), f'freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}'
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis.view(*shape)
+
+
+def rotate_half(x):
+    x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+    return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+
+def apply_rotary_emb(
+        xq: torch.Tensor,
+        xk: Optional[torch.Tensor],
+        freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+        head_first: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor.
+
+    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
+    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
+    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
+    returned as real tensors.
+
+    Args:
+        xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D]
+        xk (torch.Tensor): Key tensor to apply rotary embeddings.   [B, S, H, D]
+        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Precomputed frequency tensor for complex exponentials.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+
+    """
+    xk_out = None
+    if isinstance(freqs_cis, tuple):
+        cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)    # [S, D]
+        cos, sin = cos.to(xq.device), sin.to(xq.device)
+        xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
+        if xk is not None:
+            xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
+    else:
+        xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))  # [B, S, H, D//2]
+        freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device)   # [S, D//2] --> [1, S, 1, D//2]
+        xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
+        if xk is not None:
+            xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))  # [B, S, H, D//2]
+            xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk)
+
+    return xq_out, xk_out
+
+
+class FlashSelfMHAModified(nn.Module):
+    """
+    Use QK Normalization.
+    """
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 qkv_bias=True,
+                 qk_norm=False,
+                 attn_drop=0.0,
+                 proj_drop=0.0,
+                 device=None,
+                 dtype=None,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        assert self.dim % num_heads == 0, "self.kdim must be divisible by num_heads"
+        self.head_dim = self.dim // num_heads
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+
+        self.Wqkv = nn.Linear(dim, 3 * dim, bias=qkv_bias, **factory_kwargs)
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.inner_attn = FlashSelfAttention(attention_dropout=attn_drop)
+        self.out_proj = nn.Linear(dim, dim, bias=qkv_bias, **factory_kwargs)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, freqs_cis_img=None):
+        """
+        Parameters
+        ----------
+        x: torch.Tensor
+            (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim)
+        freqs_cis_img: torch.Tensor
+            (batch, hidden_dim // 2), RoPE for image
+        """
+        b, s, d = x.shape
+
+        qkv = self.Wqkv(x)
+        qkv = qkv.view(b, s, 3, self.num_heads, self.head_dim)  # [b, s, 3, h, d]
+        q, k, v = qkv.unbind(dim=2) # [b, s, h, d]
+        q = self.q_norm(q).half()   # [b, s, h, d]
+        k = self.k_norm(k).half()
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, kk = apply_rotary_emb(q, k, freqs_cis_img)
+            assert qq.shape == q.shape and kk.shape == k.shape, f'qq: {qq.shape}, q: {q.shape}, kk: {kk.shape}, k: {k.shape}'
+            q, k = qq, kk
+
+        qkv = torch.stack([q, k, v], dim=2)     # [b, s, 3, h, d]
+        context = self.inner_attn(qkv)
+        out = self.out_proj(context.view(b, s, d))
+        out = self.proj_drop(out)
+
+        out_tuple = (out,)
+
+        return out_tuple
+
+
+class FlashCrossMHAModified(nn.Module):
+    """
+    Use QK Normalization.
+    """
+    def __init__(self,
+                 qdim,
+                 kdim,
+                 num_heads,
+                 qkv_bias=True,
+                 qk_norm=False,
+                 attn_drop=0.0,
+                 proj_drop=0.0,
+                 device=None,
+                 dtype=None,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.qdim = qdim
+        self.kdim = kdim
+        self.num_heads = num_heads
+        assert self.qdim % num_heads == 0, "self.qdim must be divisible by num_heads"
+        self.head_dim = self.qdim // num_heads
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+
+        self.scale = self.head_dim ** -0.5
+
+        self.q_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.kv_proj = nn.Linear(kdim, 2 * qdim, bias=qkv_bias, **factory_kwargs)
+
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+
+        self.inner_attn = FlashCrossAttention(attention_dropout=attn_drop)
+        self.out_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, y, freqs_cis_img=None):
+        """
+        Parameters
+        ----------
+        x: torch.Tensor
+            (batch, seqlen1, hidden_dim) (where hidden_dim = num_heads * head_dim)
+        y: torch.Tensor
+            (batch, seqlen2, hidden_dim2)
+        freqs_cis_img: torch.Tensor
+            (batch, hidden_dim // num_heads), RoPE for image
+        """
+        b, s1, _ = x.shape     # [b, s1, D]
+        _, s2, _ = y.shape     # [b, s2, 1024]
+
+        q = self.q_proj(x).view(b, s1, self.num_heads, self.head_dim)       # [b, s1, h, d]
+        kv = self.kv_proj(y).view(b, s2, 2, self.num_heads, self.head_dim)  # [b, s2, 2, h, d]
+        k, v = kv.unbind(dim=2)                 # [b, s2, h, d]
+        q = self.q_norm(q).half()               # [b, s1, h, d]
+        k = self.k_norm(k).half()               # [b, s2, h, d]
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, _ = apply_rotary_emb(q, None, freqs_cis_img)
+            assert qq.shape == q.shape, f'qq: {qq.shape}, q: {q.shape}'
+            q = qq                              # [b, s1, h, d]
+        kv = torch.stack([k, v], dim=2)         # [b, s1, 2, h, d]
+        context = self.inner_attn(q, kv)        # [b, s1, h, d]
+        context = context.view(b, s1, -1)       # [b, s1, D]
+
+        out = self.out_proj(context)
+        out = self.proj_drop(out)
+
+        out_tuple = (out,)
+
+        return out_tuple
+
+
+class CrossAttention(nn.Module):
+    """
+    Use QK Normalization.
+    """
+    def __init__(self,
+                 qdim,
+                 kdim,
+                 num_heads,
+                 qkv_bias=True,
+                 qk_norm=False,
+                 attn_drop=0.0,
+                 proj_drop=0.0,
+                 device=None,
+                 dtype=None,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.qdim = qdim
+        self.kdim = kdim
+        self.num_heads = num_heads
+        assert self.qdim % num_heads == 0, "self.qdim must be divisible by num_heads"
+        self.head_dim = self.qdim // num_heads
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+        self.scale = self.head_dim ** -0.5
+
+        self.q_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.kv_proj = nn.Linear(kdim, 2 * qdim, bias=qkv_bias, **factory_kwargs)
+
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.out_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, y, freqs_cis_img=None):
+        """
+        Parameters
+        ----------
+        x: torch.Tensor
+            (batch, seqlen1, hidden_dim) (where hidden_dim = num heads * head dim)
+        y: torch.Tensor
+            (batch, seqlen2, hidden_dim2)
+        freqs_cis_img: torch.Tensor
+            (batch, hidden_dim // 2), RoPE for image
+        """
+        b, s1, c = x.shape     # [b, s1, D]
+        _, s2, c = y.shape     # [b, s2, 1024]
+
+        q = self.q_proj(x).view(b, s1, self.num_heads, self.head_dim)   # [b, s1, h, d]
+        kv = self.kv_proj(y).view(b, s2, 2, self.num_heads, self.head_dim)    # [b, s2, 2, h, d]
+        k, v = kv.unbind(dim=2) # [b, s, h, d]
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, _ = apply_rotary_emb(q, None, freqs_cis_img)
+            assert qq.shape == q.shape, f'qq: {qq.shape}, q: {q.shape}'
+            q = qq
+
+
+        q = q.transpose(-2, -3).contiguous()        # q ->  B, L1, H, C - B, H, L1, C
+        k = k.transpose(-2, -3).contiguous()      # k ->  B, L2, H, C - B, H, C, L2
+        v = v.transpose(-2, -3).contiguous() 
+
+        """
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)                                # attn -> B, H, L1, L2
+        attn = attn.softmax(dim=-1)                 # attn -> B, H, L1, L2
+        attn = self.attn_drop(attn)
+        x = attn @ v             # v -> B, L2, H, C - B, H, L2, C    x-> B, H, L1, C
+
+        context = x.transpose(1, 2)                 # context -> B, H, L1, C - B, L1, H, C
+        context = context.contiguous().view(b, s1, -1)
+        """
+      
+
+        context = optimized_attention(q, k, v, self.num_heads, skip_reshape=True)
+  
+
+      
+        out = self.out_proj(context)  # context.reshape - B, L1, -1
+        out = self.proj_drop(out)
+
+        out_tuple = (out,)
+
+        return out_tuple
+
+
+class Attention(nn.Module):
+    """
+    We rename some layer names to align with flash attention
+    """
+    def __init__(self, dim, num_heads, qkv_bias=True, qk_norm=False, attn_drop=0., proj_drop=0.,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        assert self.dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.head_dim = self.dim // num_heads
+        # This assertion is aligned with flash attention
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+        self.scale = self.head_dim ** -0.5
+
+        # qkv --> Wqkv
+        self.Wqkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.out_proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, freqs_cis_img=None):
+        B, N, C = x.shape
+        qkv = self.Wqkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)   # [3, b, h, s, d]
+        q, k, v = qkv.unbind(0)     # [b, h, s, d]
+        q = self.q_norm(q)          # [b, h, s, d]
+        k = self.k_norm(k)          # [b, h, s, d]
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, kk = apply_rotary_emb(q, k, freqs_cis_img, head_first=True)
+            assert qq.shape == q.shape and kk.shape == k.shape, \
+                f'qq: {qq.shape}, q: {q.shape}, kk: {kk.shape}, k: {k.shape}'
+            q, k = qq, kk
+
+        x = optimized_attention(q, k, v, self.num_heads, skip_reshape=True)
+        """
+        import pdb
+        pdb.set_trace()
+        
+        q = q * self.scale 
+        attn = q @ k.transpose(-2, -1)              # [b, h, s, d] @ [b, h, d, s]
+        attn = attn.softmax(dim=-1)                 # [b, h, s, s]
+        attn = self.attn_drop(attn)
+        x = attn @ v                                # [b, h, s, d]
+
+        x = x.transpose(1, 2).reshape(B, N, C)      # [b, s, h, d]
+
+        x_comfyui = optimized_attention(q, k, v, self.num_heads, skip_reshape=True)
+        """
+
+
+        
+
+        x = self.out_proj(x)
+        x = self.proj_drop(x)
+        
+
+        out_tuple = (x,)
+
+        return out_tuple

custom_nodes/comfyui-hydit_min/hydit/modules/attn_layers.py.bak

@@ -0,0 +1,377 @@
+import torch
+import torch.nn as nn
+from typing import Tuple, Union, Optional
+
+try:
+    import flash_attn
+    if hasattr(flash_attn, '__version__') and int(flash_attn.__version__[0]) == 2:
+        from flash_attn.flash_attn_interface import flash_attn_kvpacked_func
+        from flash_attn.modules.mha import FlashSelfAttention, FlashCrossAttention
+    else:
+        from flash_attn.flash_attn_interface import flash_attn_unpadded_kvpacked_func
+        from flash_attn.modules.mha import FlashSelfAttention, FlashCrossAttention
+except Exception as e:
+    print(f'flash_attn import failed: {e}')
+
+
+def reshape_for_broadcast(freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], x: torch.Tensor, head_first=False):
+    """
+    Reshape frequency tensor for broadcasting it with another tensor.
+
+    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
+    for the purpose of broadcasting the frequency tensor during element-wise operations.
+
+    Args:
+        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped.
+        x (torch.Tensor): Target tensor for broadcasting compatibility.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        torch.Tensor: Reshaped frequency tensor.
+
+    Raises:
+        AssertionError: If the frequency tensor doesn't match the expected shape.
+        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
+    """
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+
+    if isinstance(freqs_cis, tuple):
+        # freqs_cis: (cos, sin) in real space
+        if head_first:
+            assert freqs_cis[0].shape == (x.shape[-2], x.shape[-1]), f'freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}'
+            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        else:
+            assert freqs_cis[0].shape == (x.shape[1], x.shape[-1]), f'freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}'
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
+    else:
+        # freqs_cis: values in complex space
+        if head_first:
+            assert freqs_cis.shape == (x.shape[-2], x.shape[-1]), f'freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}'
+            shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        else:
+            assert freqs_cis.shape == (x.shape[1], x.shape[-1]), f'freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}'
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis.view(*shape)
+
+
+def rotate_half(x):
+    x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+    return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+
+def apply_rotary_emb(
+        xq: torch.Tensor,
+        xk: Optional[torch.Tensor],
+        freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+        head_first: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor.
+
+    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
+    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
+    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
+    returned as real tensors.
+
+    Args:
+        xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D]
+        xk (torch.Tensor): Key tensor to apply rotary embeddings.   [B, S, H, D]
+        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Precomputed frequency tensor for complex exponentials.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+
+    """
+    xk_out = None
+    if isinstance(freqs_cis, tuple):
+        cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)    # [S, D]
+        cos, sin = cos.to(xq.device), sin.to(xq.device)
+        xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
+        if xk is not None:
+            xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
+    else:
+        xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))  # [B, S, H, D//2]
+        freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device)   # [S, D//2] --> [1, S, 1, D//2]
+        xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
+        if xk is not None:
+            xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))  # [B, S, H, D//2]
+            xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk)
+
+    return xq_out, xk_out
+
+
+class FlashSelfMHAModified(nn.Module):
+    """
+    Use QK Normalization.
+    """
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 qkv_bias=True,
+                 qk_norm=False,
+                 attn_drop=0.0,
+                 proj_drop=0.0,
+                 device=None,
+                 dtype=None,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        assert self.dim % num_heads == 0, "self.kdim must be divisible by num_heads"
+        self.head_dim = self.dim // num_heads
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+
+        self.Wqkv = nn.Linear(dim, 3 * dim, bias=qkv_bias, **factory_kwargs)
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.inner_attn = FlashSelfAttention(attention_dropout=attn_drop)
+        self.out_proj = nn.Linear(dim, dim, bias=qkv_bias, **factory_kwargs)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, freqs_cis_img=None):
+        """
+        Parameters
+        ----------
+        x: torch.Tensor
+            (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim)
+        freqs_cis_img: torch.Tensor
+            (batch, hidden_dim // 2), RoPE for image
+        """
+        b, s, d = x.shape
+
+        qkv = self.Wqkv(x)
+        qkv = qkv.view(b, s, 3, self.num_heads, self.head_dim)  # [b, s, 3, h, d]
+        q, k, v = qkv.unbind(dim=2) # [b, s, h, d]
+        q = self.q_norm(q).half()   # [b, s, h, d]
+        k = self.k_norm(k).half()
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, kk = apply_rotary_emb(q, k, freqs_cis_img)
+            assert qq.shape == q.shape and kk.shape == k.shape, f'qq: {qq.shape}, q: {q.shape}, kk: {kk.shape}, k: {k.shape}'
+            q, k = qq, kk
+
+        qkv = torch.stack([q, k, v], dim=2)     # [b, s, 3, h, d]
+        context = self.inner_attn(qkv)
+        out = self.out_proj(context.view(b, s, d))
+        out = self.proj_drop(out)
+
+        out_tuple = (out,)
+
+        return out_tuple
+
+
+class FlashCrossMHAModified(nn.Module):
+    """
+    Use QK Normalization.
+    """
+    def __init__(self,
+                 qdim,
+                 kdim,
+                 num_heads,
+                 qkv_bias=True,
+                 qk_norm=False,
+                 attn_drop=0.0,
+                 proj_drop=0.0,
+                 device=None,
+                 dtype=None,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.qdim = qdim
+        self.kdim = kdim
+        self.num_heads = num_heads
+        assert self.qdim % num_heads == 0, "self.qdim must be divisible by num_heads"
+        self.head_dim = self.qdim // num_heads
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+
+        self.scale = self.head_dim ** -0.5
+
+        self.q_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.kv_proj = nn.Linear(kdim, 2 * qdim, bias=qkv_bias, **factory_kwargs)
+
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+
+        self.inner_attn = FlashCrossAttention(attention_dropout=attn_drop)
+        self.out_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, y, freqs_cis_img=None):
+        """
+        Parameters
+        ----------
+        x: torch.Tensor
+            (batch, seqlen1, hidden_dim) (where hidden_dim = num_heads * head_dim)
+        y: torch.Tensor
+            (batch, seqlen2, hidden_dim2)
+        freqs_cis_img: torch.Tensor
+            (batch, hidden_dim // num_heads), RoPE for image
+        """
+        b, s1, _ = x.shape     # [b, s1, D]
+        _, s2, _ = y.shape     # [b, s2, 1024]
+
+        q = self.q_proj(x).view(b, s1, self.num_heads, self.head_dim)       # [b, s1, h, d]
+        kv = self.kv_proj(y).view(b, s2, 2, self.num_heads, self.head_dim)  # [b, s2, 2, h, d]
+        k, v = kv.unbind(dim=2)                 # [b, s2, h, d]
+        q = self.q_norm(q).half()               # [b, s1, h, d]
+        k = self.k_norm(k).half()               # [b, s2, h, d]
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, _ = apply_rotary_emb(q, None, freqs_cis_img)
+            assert qq.shape == q.shape, f'qq: {qq.shape}, q: {q.shape}'
+            q = qq                              # [b, s1, h, d]
+        kv = torch.stack([k, v], dim=2)         # [b, s1, 2, h, d]
+        context = self.inner_attn(q, kv)        # [b, s1, h, d]
+        context = context.view(b, s1, -1)       # [b, s1, D]
+
+        out = self.out_proj(context)
+        out = self.proj_drop(out)
+
+        out_tuple = (out,)
+
+        return out_tuple
+
+
+class CrossAttention(nn.Module):
+    """
+    Use QK Normalization.
+    """
+    def __init__(self,
+                 qdim,
+                 kdim,
+                 num_heads,
+                 qkv_bias=True,
+                 qk_norm=False,
+                 attn_drop=0.0,
+                 proj_drop=0.0,
+                 device=None,
+                 dtype=None,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.qdim = qdim
+        self.kdim = kdim
+        self.num_heads = num_heads
+        assert self.qdim % num_heads == 0, "self.qdim must be divisible by num_heads"
+        self.head_dim = self.qdim // num_heads
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+        self.scale = self.head_dim ** -0.5
+
+        self.q_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.kv_proj = nn.Linear(kdim, 2 * qdim, bias=qkv_bias, **factory_kwargs)
+
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.out_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, y, freqs_cis_img=None):
+        """
+        Parameters
+        ----------
+        x: torch.Tensor
+            (batch, seqlen1, hidden_dim) (where hidden_dim = num heads * head dim)
+        y: torch.Tensor
+            (batch, seqlen2, hidden_dim2)
+        freqs_cis_img: torch.Tensor
+            (batch, hidden_dim // 2), RoPE for image
+        """
+        b, s1, c = x.shape     # [b, s1, D]
+        _, s2, c = y.shape     # [b, s2, 1024]
+
+        q = self.q_proj(x).view(b, s1, self.num_heads, self.head_dim)   # [b, s1, h, d]
+        kv = self.kv_proj(y).view(b, s2, 2, self.num_heads, self.head_dim)    # [b, s2, 2, h, d]
+        k, v = kv.unbind(dim=2) # [b, s, h, d]
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, _ = apply_rotary_emb(q, None, freqs_cis_img)
+            assert qq.shape == q.shape, f'qq: {qq.shape}, q: {q.shape}'
+            q = qq
+
+        q = q * self.scale
+        q = q.transpose(-2, -3).contiguous()        # q ->  B, L1, H, C - B, H, L1, C
+        k = k.permute(0, 2, 3, 1).contiguous()      # k ->  B, L2, H, C - B, H, C, L2
+        attn = q @ k                                # attn -> B, H, L1, L2
+        attn = attn.softmax(dim=-1)                 # attn -> B, H, L1, L2
+        attn = self.attn_drop(attn)
+        x = attn @ v.transpose(-2, -3)              # v -> B, L2, H, C - B, H, L2, C    x-> B, H, L1, C
+        context = x.transpose(1, 2)                 # context -> B, H, L1, C - B, L1, H, C
+
+        context = context.contiguous().view(b, s1, -1)
+
+        out = self.out_proj(context)  # context.reshape - B, L1, -1
+        out = self.proj_drop(out)
+
+        out_tuple = (out,)
+
+        return out_tuple
+
+
+class Attention(nn.Module):
+    """
+    We rename some layer names to align with flash attention
+    """
+    def __init__(self, dim, num_heads, qkv_bias=True, qk_norm=False, attn_drop=0., proj_drop=0.,
+                 norm_layer=nn.LayerNorm,
+                 ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        assert self.dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.head_dim = self.dim // num_heads
+        # This assertion is aligned with flash attention
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+        self.scale = self.head_dim ** -0.5
+
+        # qkv --> Wqkv
+        self.Wqkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        # TODO: eps should be 1 / 65530 if using fp16
+        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.out_proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, freqs_cis_img=None):
+        B, N, C = x.shape
+        qkv = self.Wqkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)   # [3, b, h, s, d]
+        q, k, v = qkv.unbind(0)     # [b, h, s, d]
+        q = self.q_norm(q)          # [b, h, s, d]
+        k = self.k_norm(k)          # [b, h, s, d]
+
+        # Apply RoPE if needed
+        if freqs_cis_img is not None:
+            qq, kk = apply_rotary_emb(q, k, freqs_cis_img, head_first=True)
+            assert qq.shape == q.shape and kk.shape == k.shape, \
+                f'qq: {qq.shape}, q: {q.shape}, kk: {kk.shape}, k: {k.shape}'
+            q, k = qq, kk
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)              # [b, h, s, d] @ [b, h, d, s]
+        attn = attn.softmax(dim=-1)                 # [b, h, s, s]
+        attn = self.attn_drop(attn)
+        x = attn @ v                                # [b, h, s, d]
+
+        x = x.transpose(1, 2).reshape(B, N, C)      # [b, s, h, d]
+        x = self.out_proj(x)
+        x = self.proj_drop(x)
+
+        out_tuple = (x,)
+
+        return out_tuple

custom_nodes/comfyui-hydit_min/hydit/modules/ema.py

@@ -0,0 +1,127 @@
+from collections import OrderedDict
+from copy import deepcopy
+
+import torch
+from deepspeed.utils import instrument_w_nvtx
+from pathlib import Path
+
+def requires_grad(model, flag=True):
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+class EMA(object):
+    def __init__(self, args, model, device, logger):
+        if args.ema_dtype == 'fp32':
+            self.warmup = args.ema_warmup
+            self.update_after_step = 0
+            self.max_value = args.ema_decay if args.ema_decay is not None else 0.9999
+            self.inv_gamma = 1.0
+            self.power = args.ema_warmup_power if args.ema_warmup_power is not None else 2 / 3
+            self.min_value = 0.0
+        else:
+            self.warmup = args.ema_warmup
+            self.update_after_step = 0
+            self.max_value = args.ema_decay if args.ema_decay is not None else 0.992
+            self.inv_gamma = 1.0
+            self.power = args.ema_warmup_power if args.ema_warmup_power is not None else 0.446249
+            # 0.446249 == math.log(1 - 0.992) / math.log(50000)
+            self.min_value = 0.0
+
+        self.ema_reset_decay = args.ema_reset_decay
+        self.decay_steps = 0
+
+        if args.ema_dtype == 'none':
+            ema_dtype = 'fp16' if args.use_fp16 else 'fp32'
+        else:
+            ema_dtype = args.ema_dtype
+        
+        # 由于module.half()和module.float()会发生inplace类型修改，因此需要先copy后修改类型
+        self.ema_model = deepcopy(model)
+        if ema_dtype == 'fp16':
+            self.ema_model = self.ema_model.half().to(device)
+        elif ema_dtype == 'fp32':
+            self.ema_model = self.ema_model.float().to(device)
+        else:
+            raise ValueError(f"Unknown EMA dtype {ema_dtype}.")
+
+        requires_grad(self.ema_model, False)
+
+        logger.info(f"    Using EMA with date type {args.ema_dtype} "
+                    f"(decay={args.ema_decay}, warmup={args.ema_warmup}, warmup_power={args.ema_warmup_power}, "
+                    f"reset_decay={args.ema_reset_decay}).")
+
+    def get_decay(self):
+        """
+        @crowsonkb's notes on EMA Warmup:
+            If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan
+            to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),
+            gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999
+            at 215.4k steps).
+
+        @jarvizhang's notes on EMA max_value when enabling FP16:
+            If using FP16 for EMA, max_value=0.995 is better (Don't larger than 0.999, unless you know
+            what you are doing). This is because FP16 has less precision than FP32, so the EMA value can
+            be pushed out of the range of FP16.
+
+            gamma=1, power=0.446249 are good values for models (reaches decay factor 0.99 at 30K steps,
+            0.992 at 50K steps).
+        """
+        if self.warmup:
+            step = max(0, self.decay_steps - self.update_after_step - 1)
+            value = 1 - (1 + step / self.inv_gamma) ** -self.power
+
+            if step <= 0:
+                return 0.0
+
+            return max(self.min_value, min(value, self.max_value))
+        else:
+            return self.max_value
+
+    @torch.no_grad()
+    @instrument_w_nvtx
+    def update(self, model, step, decay=None):
+        """
+        Step the EMA model towards the current model.
+
+        Parameters
+        ----------
+        model: nn.Module
+            The current model
+        step: int
+            The current training step. This is used to determine the decay factor. If you want to control
+            the decay, you can pass in a custom step instead.
+            For example, if you want to restart the EMA decay, you can pass in step=0 at start and increase
+            step by step.
+        decay: float
+            The decay factor. If None, will be determined by the current step.
+        """
+        if decay is None:
+            if self.ema_reset_decay:
+                self.decay_steps += 1
+            else:
+                self.decay_steps = step
+            decay = self.get_decay()
+
+        ema_params = OrderedDict(self.ema_model.named_parameters())
+        model_params = OrderedDict(model.named_parameters())
+        for name, param in model_params.items():
+            # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+            ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+        return None
+
+    def state_dict(self, *args, **kwargs):
+        return self.ema_model.state_dict(*args, **kwargs)
+
+    def load_state_dict(self, *args, **kwargs):
+        return self.ema_model.load_state_dict(*args, **kwargs)
+
+    def train(self):
+        self.ema_model.train()
+
+    def eval(self):
+        self.ema_model.eval()

custom_nodes/comfyui-hydit_min/hydit/modules/embedders.py

@@ -0,0 +1,111 @@
+import math
+import torch
+import torch.nn as nn
+from einops import repeat
+
+from timm.models.layers import to_2tuple
+
+
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding
+
+    Image to Patch Embedding using Conv2d
+
+    A convolution based approach to patchifying a 2D image w/ embedding projection.
+
+    Based on the impl in https://github.com/google-research/vision_transformer
+
+    Hacked together by / Copyright 2020 Ross Wightman
+
+    Remove the _assert function in forward function to be compatible with multi-resolution images.
+    """
+    def __init__(
+            self,
+            img_size=224,
+            patch_size=16,
+            in_chans=3,
+            embed_dim=768,
+            norm_layer=None,
+            flatten=True,
+            bias=True,
+    ):
+        super().__init__()
+        if isinstance(img_size, int):
+            img_size = to_2tuple(img_size)
+        elif isinstance(img_size, (tuple, list)) and len(img_size) == 2:
+            img_size = tuple(img_size)
+        else:
+            raise ValueError(f"img_size must be int or tuple/list of length 2. Got {img_size}")
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def update_image_size(self, img_size):
+        self.img_size = img_size
+        self.grid_size = (img_size[0] // self.patch_size[0], img_size[1] // self.patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+    def forward(self, x):
+        # B, C, H, W = x.shape
+        # _assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).")
+        # _assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+def timestep_embedding(t, dim, max_period=10000, repeat_only=False):
+    """
+    Create sinusoidal timestep embeddings.
+    :param t: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an (N, D) Tensor of positional embeddings.
+    """
+    # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    if not repeat_only:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32)
+            / half
+        ).to(device=t.device)   # size: [dim/2], 一个指数衰减的曲线
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+    else:
+        embedding = repeat(t, "b -> b d", d=dim)
+    return embedding
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256, out_size=None):
+        super().__init__()
+        if out_size is None:
+            out_size = hidden_size
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, out_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    def forward(self, t):
+        t_freq = timestep_embedding(t, self.frequency_embedding_size).type(self.mlp[0].weight.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb

custom_nodes/comfyui-hydit_min/hydit/modules/fp16_layers.py

@@ -0,0 +1,86 @@
+import torch
+from torch.autograd import Variable
+from torch.nn.parameter import Parameter
+
+
+_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
+_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
+
+
+def conversion_helper(val, conversion):
+    """Apply conversion to val. Recursively apply conversion if `val`
+    #is a nested tuple/list structure."""
+    if isinstance(val, dict):
+        res_dict = {}
+        for k, v in val.items():
+            if  k!= 'cos_cis_img' and k != 'sin_cis_img':
+                res_dict[k] = conversion_helper(v, conversion)
+            else:
+                res_dict[k] = v
+        return res_dict
+    if not isinstance(val, (tuple, list)):
+        return conversion(val)
+    rtn = [conversion_helper(v, conversion) for v in val]
+    if isinstance(val, tuple):
+        rtn = tuple(rtn)
+    return rtn
+
+
+def fp32_to_float16(val, float16_convertor):
+    """Convert fp32 `val` to fp16/bf16"""
+    def half_conversion(val):
+        val_typecheck = val
+        if isinstance(val_typecheck, (Parameter, Variable)):
+            val_typecheck = val.data
+        if isinstance(val_typecheck, _FLOAT_TYPES):
+            val = float16_convertor(val)
+        return val
+    return conversion_helper(val, half_conversion)
+
+
+def float16_to_fp32(val):
+    """Convert fp16/bf16 `val` to fp32"""
+    def float_conversion(val):
+        val_typecheck = val
+        if isinstance(val_typecheck, (Parameter, Variable)):
+            val_typecheck = val.data
+        if isinstance(val_typecheck, (_HALF_TYPES,)):
+            val = val.float()
+        return val
+    return conversion_helper(val, float_conversion)
+
+
+class Float16Module(torch.nn.Module):
+
+    def __init__(self, module, args):
+        super(Float16Module, self).__init__()
+
+        self.add_module('module', module.half())
+
+        def float16_convertor(val):
+            return val.half()
+
+        self.float16_convertor = float16_convertor
+
+        self.config = self.module.config
+        self.dtype = torch.float16
+
+    def forward(self, *inputs, **kwargs):
+        inputs = fp32_to_float16(inputs, self.float16_convertor)
+        kwargs = fp32_to_float16(kwargs, self.float16_convertor)
+        outputs = self.module(*inputs, **kwargs)
+        outputs = float16_to_fp32(outputs)
+        return outputs
+
+    def state_dict(self, destination=None, prefix='', keep_vars=False):
+        return self.module.state_dict(destination, prefix, keep_vars)
+
+
+    def state_dict_for_save_checkpoint(self, destination=None, prefix='',
+                                       keep_vars=False):
+        return self.module.state_dict_for_save_checkpoint(destination, prefix,
+                                                          keep_vars)
+
+    def load_state_dict(self, state_dict, strict=True):
+        self.module.load_state_dict(state_dict, strict=strict)
+

custom_nodes/comfyui-hydit_min/hydit/modules/models.py

@@ -0,0 +1,556 @@
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+#from diffusers.configuration_utils import ConfigMixin, register_to_config
+#from diffusers.models import ModelMixin
+#from peft.utils import (
+#    ModulesToSaveWrapper,
+#    _get_submodules,
+#)
+from timm.models.vision_transformer import Mlp
+from torch.utils import checkpoint
+from tqdm import tqdm
+#from transformers.integrations import PeftAdapterMixin
+
+from .attn_layers import Attention, CrossAttention
+from .embedders import TimestepEmbedder, PatchEmbed, timestep_embedding
+from .norm_layers import RMSNorm
+from .poolers import AttentionPool
+#from .posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class FP32_Layernorm(nn.LayerNorm):
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        origin_dtype = inputs.dtype
+        return F.layer_norm(inputs.float(), self.normalized_shape, self.weight.float(), self.bias.float(),
+                            self.eps).to(origin_dtype)
+
+
+class FP32_SiLU(nn.SiLU):
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.silu(inputs.float(), inplace=False).to(inputs.dtype)
+
+
+class HunYuanDiTBlock(nn.Module):
+    """
+    A HunYuanDiT block with `add` conditioning.
+    """
+    def __init__(self,
+                 hidden_size,
+                 c_emb_size,
+                 num_heads,
+                 mlp_ratio=4.0,
+                 text_states_dim=1024,
+                 use_flash_attn=False,
+                 qk_norm=False,
+                 norm_type="layer",
+                 skip=False,
+                 ):
+        super().__init__()
+        self.use_flash_attn = use_flash_attn
+        use_ele_affine = True
+
+        if norm_type == "layer":
+            norm_layer = FP32_Layernorm
+        elif norm_type == "rms":
+            norm_layer = RMSNorm
+        else:
+            raise ValueError(f"Unknown norm_type: {norm_type}")
+
+        # ========================= Self-Attention =========================
+        self.norm1 = norm_layer(hidden_size, elementwise_affine=use_ele_affine, eps=1e-6)
+        if use_flash_attn:
+            self.attn1 = FlashSelfMHAModified(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=qk_norm)
+        else:
+            self.attn1 = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=qk_norm)
+
+        # ========================= FFN =========================
+        self.norm2 = norm_layer(hidden_size, elementwise_affine=use_ele_affine, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+
+        # ========================= Add =========================
+        # Simply use add like SDXL.
+        self.default_modulation = nn.Sequential(
+            FP32_SiLU(),
+            nn.Linear(c_emb_size, hidden_size, bias=True)
+        )
+
+        # ========================= Cross-Attention =========================
+        if use_flash_attn:
+            self.attn2 = FlashCrossMHAModified(hidden_size, text_states_dim, num_heads=num_heads, qkv_bias=True,
+                                               qk_norm=qk_norm)
+        else:
+            self.attn2 = CrossAttention(hidden_size, text_states_dim, num_heads=num_heads, qkv_bias=True,
+                                        qk_norm=qk_norm)
+        self.norm3 = norm_layer(hidden_size, elementwise_affine=True, eps=1e-6)
+
+        # ========================= Skip Connection =========================
+        if skip:
+            self.skip_norm = norm_layer(2 * hidden_size, elementwise_affine=True, eps=1e-6)
+            self.skip_linear = nn.Linear(2 * hidden_size, hidden_size)
+        else:
+            self.skip_linear = None
+
+        self.gradient_checkpointing = False
+
+    def _forward(self, x, c=None, text_states=None, freq_cis_img=None, skip=None):
+        # Long Skip Connection
+        if self.skip_linear is not None:
+            cat = torch.cat([x, skip], dim=-1)
+            cat = self.skip_norm(cat)
+            x = self.skip_linear(cat)
+
+        # Self-Attention
+        shift_msa = self.default_modulation(c).unsqueeze(dim=1)
+        attn_inputs = (
+            self.norm1(x) + shift_msa, freq_cis_img,
+        )
+        x = x + self.attn1(*attn_inputs)[0]
+
+        # Cross-Attention
+        cross_inputs = (
+            self.norm3(x), text_states, freq_cis_img
+        )
+        x = x + self.attn2(*cross_inputs)[0]
+
+        # FFN Layer
+        mlp_inputs = self.norm2(x)
+        x = x + self.mlp(mlp_inputs)
+
+        return x
+
+    def forward(self, x, c=None, text_states=None, freq_cis_img=None, skip=None):
+        if self.gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(self._forward, x, c, text_states, freq_cis_img, skip)
+        return self._forward(x, c, text_states, freq_cis_img, skip)
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of HunYuanDiT.
+    """
+    def __init__(self, final_hidden_size, c_emb_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(final_hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(final_hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            FP32_SiLU(),
+            nn.Linear(c_emb_size, 2 * final_hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class HunYuanDiT(nn.Module):
+    """
+    HunYuanDiT: Diffusion model with a Transformer backbone.
+
+    Inherit ModelMixin and ConfigMixin to be compatible with the sampler StableDiffusionPipeline of diffusers.
+
+    Inherit PeftAdapterMixin to be compatible with the PEFT training pipeline.
+
+    Parameters
+    ----------
+    args: argparse.Namespace
+        The arguments parsed by argparse.
+    input_size: tuple
+        The size of the input image.
+    patch_size: int
+        The size of the patch.
+    in_channels: int
+        The number of input channels.
+    hidden_size: int
+        The hidden size of the transformer backbone.
+    depth: int
+        The number of transformer blocks.
+    num_heads: int
+        The number of attention heads.
+    mlp_ratio: float
+        The ratio of the hidden size of the MLP in the transformer block.
+    log_fn: callable
+        The logging function.
+    """
+    #@register_to_config
+    def __init__(self,
+                 args: Any,
+                 input_size: tuple = (32, 32),
+                 patch_size: int = 2,
+                 in_channels: int = 4,
+                 hidden_size: int = 1152,
+                 depth: int = 28,
+                 num_heads: int = 16,
+                 mlp_ratio: float = 4.0,
+                 log_fn: callable = print,
+                    **kwargs,
+    ):
+        super().__init__()
+        self.args = args
+        self.log_fn = log_fn
+        self.depth = depth
+        self.learn_sigma = args.learn_sigma
+        self.in_channels = in_channels
+        self.out_channels = in_channels * 2 if args.learn_sigma else in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.text_states_dim = args.text_states_dim
+        self.text_states_dim_t5 = args.text_states_dim_t5
+        self.text_len = args.text_len
+        self.text_len_t5 = args.text_len_t5
+        self.norm = args.norm
+        self.dtype = torch.float16
+        #import pdb
+        #pdb.set_trace()
+
+        use_flash_attn = args.infer_mode == 'fa' or args.use_flash_attn
+        if use_flash_attn:
+            log_fn(f"    Enable Flash Attention.")
+        qk_norm = args.qk_norm  # See http://arxiv.org/abs/2302.05442 for details.
+
+        self.mlp_t5 = nn.Sequential(
+            nn.Linear(self.text_states_dim_t5, self.text_states_dim_t5 * 4, bias=True),
+            FP32_SiLU(),
+            nn.Linear(self.text_states_dim_t5 * 4, self.text_states_dim, bias=True),
+        )
+        # learnable replace
+        self.text_embedding_padding = nn.Parameter(
+            torch.randn(self.text_len + self.text_len_t5, self.text_states_dim, dtype=torch.float32))
+
+        # Attention pooling
+        pooler_out_dim = 1024
+        self.pooler = AttentionPool(self.text_len_t5, self.text_states_dim_t5, num_heads=8, output_dim=pooler_out_dim)
+
+        # Dimension of the extra input vectors
+        self.extra_in_dim = pooler_out_dim
+
+        if args.size_cond:
+            # Image size and crop size conditions
+            self.extra_in_dim += 6 * 256
+
+        if args.use_style_cond:
+            # Here we use a default learned embedder layer for future extension.
+            self.style_embedder = nn.Embedding(1, hidden_size)
+            self.extra_in_dim += hidden_size
+
+        # Text embedding for `add`
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.extra_embedder = nn.Sequential(
+            nn.Linear(self.extra_in_dim, hidden_size * 4),
+            FP32_SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size, bias=True),
+        )
+
+        # Image embedding
+        num_patches = self.x_embedder.num_patches
+        log_fn(f"    Number of tokens: {num_patches}")
+
+        # HUnYuanDiT Blocks
+        self.blocks = nn.ModuleList([
+            HunYuanDiTBlock(hidden_size=hidden_size,
+                            c_emb_size=hidden_size,
+                            num_heads=num_heads,
+                            mlp_ratio=mlp_ratio,
+                            text_states_dim=self.text_states_dim,
+                            use_flash_attn=use_flash_attn,
+                            qk_norm=qk_norm,
+                            norm_type=self.norm,
+                            skip=layer > depth // 2,
+                            )
+            for layer in range(depth)
+        ])
+
+        self.final_layer = FinalLayer(hidden_size, hidden_size, patch_size, self.out_channels)
+        self.unpatchify_channels = self.out_channels
+
+        self.initialize_weights()
+
+    def check_condition_validation(self, image_meta_size, style):
+        if self.args.size_cond is None and image_meta_size is not None:
+            raise ValueError(f"When `size_cond` is None, `image_meta_size` should be None, but got "
+                             f"{type(image_meta_size)}. ")
+        if self.args.size_cond is not None and image_meta_size is None:
+            raise ValueError(f"When `size_cond` is not None, `image_meta_size` should not be None. ")
+        if not self.args.use_style_cond and style is not None:
+            raise ValueError(f"When `use_style_cond` is False, `style` should be None, but got {type(style)}. ")
+        if self.args.use_style_cond and style is None:
+            raise ValueError(f"When `use_style_cond` is True, `style` should be not None.")
+
+    def enable_gradient_checkpointing(self):
+        for block in self.blocks:
+            block.gradient_checkpointing = True
+
+    def disable_gradient_checkpointing(self):
+        for block in self.blocks:
+            block.gradient_checkpointing = False
+
+    def forward(self,
+                x,
+                t,
+                encoder_hidden_states=None,
+                text_embedding_mask=None,
+                encoder_hidden_states_t5=None,
+                text_embedding_mask_t5=None,
+                image_meta_size=None,
+                style=None,
+                cos_cis_img=None,
+                sin_cis_img=None,
+                return_dict=True,
+                controls=None,
+                ):
+        """
+        Forward pass of the encoder.
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            (B, D, H, W)
+        t: torch.Tensor
+            (B)
+        encoder_hidden_states: torch.Tensor
+            CLIP text embedding, (B, L_clip, D)
+        text_embedding_mask: torch.Tensor
+            CLIP text embedding mask, (B, L_clip)
+        encoder_hidden_states_t5: torch.Tensor
+            T5 text embedding, (B, L_t5, D)
+        text_embedding_mask_t5: torch.Tensor
+            T5 text embedding mask, (B, L_t5)
+        image_meta_size: torch.Tensor
+            (B, 6)
+        style: torch.Tensor
+            (B)
+        cos_cis_img: torch.Tensor
+        sin_cis_img: torch.Tensor
+        return_dict: bool
+            Whether to return a dictionary.
+        """
+        #import pdb
+        #pdb.set_trace()
+        text_states = encoder_hidden_states                     # 2,77,1024
+        text_states_t5 = encoder_hidden_states_t5               # 2,256,2048
+        text_states_mask = text_embedding_mask.bool()           # 2,77
+        text_states_t5_mask = text_embedding_mask_t5.bool()     # 2,256
+        b_t5, l_t5, c_t5 = text_states_t5.shape
+        text_states_t5 = self.mlp_t5(text_states_t5.view(-1, c_t5))
+        text_states = torch.cat([text_states, text_states_t5.view(b_t5, l_t5, -1)], dim=1)  # 2,205，1024
+        clip_t5_mask = torch.cat([text_states_mask, text_states_t5_mask], dim=-1)
+
+        clip_t5_mask = clip_t5_mask
+        text_states = torch.where(clip_t5_mask.unsqueeze(2), text_states, self.text_embedding_padding.to(text_states))
+
+        _, _, oh, ow = x.shape
+        th, tw = oh // self.patch_size, ow // self.patch_size
+
+        # ========================= Build time and image embedding =========================
+        t = self.t_embedder(t)
+        x = self.x_embedder(x)
+
+        # Get image RoPE embedding according to `reso`lution.
+        freqs_cis_img = (cos_cis_img, sin_cis_img)
+
+        # ========================= Concatenate all extra vectors =========================
+        # Build text tokens with pooling
+        extra_vec = self.pooler(encoder_hidden_states_t5)
+
+        self.check_condition_validation(image_meta_size, style)
+        # Build image meta size tokens if applicable
+        if image_meta_size is not None:
+            image_meta_size = timestep_embedding(image_meta_size.view(-1), 256)   # [B * 6, 256]
+            if self.args.use_fp16:
+                image_meta_size = image_meta_size.half()
+            image_meta_size = image_meta_size.view(-1, 6 * 256)
+            extra_vec = torch.cat([extra_vec, image_meta_size], dim=1)  # [B, D + 6 * 256]
+
+        # Build style tokens
+        if style is not None:
+            style_embedding = self.style_embedder(style)
+            extra_vec = torch.cat([extra_vec, style_embedding], dim=1)
+
+        # Concatenate all extra vectors
+        c = t + self.extra_embedder(extra_vec)  # [B, D]
+
+        # ========================= Forward pass through HunYuanDiT blocks =========================
+        skips = []
+        for layer, block in enumerate(self.blocks):
+            if layer > self.depth // 2:
+                if controls is not None:
+                    skip = skips.pop() + controls.pop()
+                else:
+                    skip = skips.pop()
+                x = block(x, c, text_states, freqs_cis_img, skip)   # (N, L, D)
+            else:
+                x = block(x, c, text_states, freqs_cis_img)         # (N, L, D)
+
+            if layer < (self.depth // 2 - 1):
+                skips.append(x)
+        if controls is not None and len(controls) != 0:
+            raise ValueError("The number of controls is not equal to the number of skip connections.")
+
+        # ========================= Final layer =========================
+        x = self.final_layer(x, c)                              # (N, L, patch_size ** 2 * out_channels)
+        x = self.unpatchify(x, th, tw)                          # (N, out_channels, H, W)
+
+        if return_dict:
+            return {'x': x}
+        return x
+    
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.extra_embedder[0].weight, std=0.02)
+        nn.init.normal_(self.extra_embedder[2].weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in HunYuanDiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.default_modulation[-1].weight, 0)
+            nn.init.constant_(block.default_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x, h, w):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.unpatchify_channels
+        p = self.x_embedder.patch_size[0]
+        # h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
+        return imgs
+
+    def _replace_module(self, parent, child_name, new_module, child) -> None:
+        setattr(parent, child_name, new_module)
+        # It's not necessary to set requires_grad here, as that is handled by
+        # _mark_only_adapters_as_trainable
+
+        # child layer wraps the original module, unpack it
+        if hasattr(child, "base_layer"):
+            child = child.get_base_layer()
+        elif hasattr(child, "quant_linear_module"):
+            # TODO maybe not necessary to have special treatment?
+            child = child.quant_linear_module
+
+        if not hasattr(new_module, "base_layer"):
+            new_module.weight = child.weight
+            if hasattr(child, "bias"):
+                new_module.bias = child.bias
+
+        if getattr(child, "state", None) is not None:
+            if hasattr(new_module, "base_layer"):
+                new_module.base_layer.state = child.state
+            else:
+                new_module.state = child.state
+            new_module.to(child.weight.device)
+
+        # dispatch to correct device
+        for name, module in new_module.named_modules():
+            # if any(prefix in name for prefix in PREFIXES):
+            #     module.to(child.weight.device)
+            if "ranknum" in name:
+                module.to(child.weight.device)
+
+    def merge_and_unload(self,
+                         merge=True,
+                        progressbar: bool = False,
+                        safe_merge: bool = False,
+                        adapter_names = None,):
+        if merge:
+            if getattr(self, "quantization_method", None) == "gptq":
+                raise ValueError("Cannot merge layers when the model is gptq quantized")
+
+        def merge_recursively(module):
+            # helper function to recursively merge the base_layer of the target
+            path = []
+            layer = module
+            while hasattr(layer, "base_layer"):
+                path.append(layer)
+                layer = layer.base_layer
+            for layer_before, layer_after in zip(path[:-1], path[1:]):
+                layer_after.merge(safe_merge=safe_merge, adapter_names=adapter_names)
+                layer_before.base_layer = layer_after.base_layer
+            module.merge(safe_merge=safe_merge, adapter_names=adapter_names)
+
+        key_list = [key for key, _ in self.named_modules()]
+        desc = "Unloading " + ("and merging " if merge else "") + "model"
+
+        for key in tqdm(key_list, disable=not progressbar, desc=desc):
+            try:
+                parent, target, target_name = _get_submodules(self, key)
+            except AttributeError:
+                continue
+
+            if hasattr(target, "base_layer"):
+                if merge:
+                    merge_recursively(target)
+                self._replace_module(parent, target_name, target.get_base_layer(), target)
+            elif isinstance(target, ModulesToSaveWrapper):
+                # save any additional trainable modules part of `modules_to_save`
+                new_module = target.modules_to_save[target.active_adapter]
+                if hasattr(new_module, "base_layer"):
+                    # check if the module is itself a tuner layer
+                    if merge:
+                        new_module.merge(safe_merge=safe_merge, adapter_names=adapter_names)
+                    new_module = new_module.get_base_layer()
+                setattr(parent, target_name, new_module)
+
+
+
+#################################################################################
+#                            HunYuanDiT Configs                                 #
+#################################################################################
+
+HUNYUAN_DIT_CONFIG = {
+    'DiT-g/2': {'depth': 40, 'hidden_size': 1408, 'patch_size': 2, 'num_heads': 16, 'mlp_ratio': 4.3637},
+    'DiT-XL/2': {'depth': 28, 'hidden_size': 1152, 'patch_size': 2, 'num_heads': 16},
+}
+
+
+def DiT_g_2(args, **kwargs):
+    return HunYuanDiT(args, depth=40, hidden_size=1408, patch_size=2, num_heads=16, mlp_ratio=4.3637, **kwargs)
+
+
+def DiT_XL_2(args, **kwargs):
+    return HunYuanDiT(args, depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
+
+
+HUNYUAN_DIT_MODELS = {
+    'DiT-g/2':  DiT_g_2,
+    'DiT-XL/2': DiT_XL_2,
+}

custom_nodes/comfyui-hydit_min/hydit/modules/norm_layers.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, elementwise_affine=True, eps: float = 1e-6):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        if hasattr(self, "weight"):
+            output = output * self.weight
+        return output
+
+
+class GroupNorm32(nn.GroupNorm):
+    def __init__(self, num_groups, num_channels, eps=1e-5, dtype=None):
+        super().__init__(num_groups=num_groups, num_channels=num_channels, eps=eps, dtype=dtype)
+
+    def forward(self, x):
+        y = super().forward(x).to(x.dtype)
+        return y
+
+def normalization(channels, dtype=None):
+    """
+    Make a standard normalization layer.
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    return GroupNorm32(num_channels=channels, num_groups=32, dtype=dtype)

custom_nodes/comfyui-hydit_min/hydit/modules/poolers.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class AttentionPool(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.permute(1, 0, 2)  # NLC -> LNC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (L+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (L+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)

custom_nodes/comfyui-hydit_min/hydit/modules/posemb_layers.py

@@ -0,0 +1,224 @@
+import torch
+import numpy as np
+from typing import Union
+
+
+def _to_tuple(x):
+    if isinstance(x, int):
+        return x, x
+    else:
+        return x
+
+
+def get_fill_resize_and_crop(src, tgt):
+    th, tw = _to_tuple(tgt)
+    h, w = _to_tuple(src)
+
+    tr = th / tw        # base resolution
+    r = h / w           # target resolution
+
+    # resize
+    if r > tr:
+        resize_height = th
+        resize_width = int(round(th / h * w))
+    else:
+        resize_width = tw
+        resize_height = int(round(tw / w * h))    # resize the target resolution down based on the base resolution
+
+    crop_top = int(round((th - resize_height) / 2.0))
+    crop_left = int(round((tw - resize_width) / 2.0))
+
+    return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)
+
+
+def get_meshgrid(start, *args):
+    if len(args) == 0:
+        # start is grid_size
+        num = _to_tuple(start)
+        start = (0, 0)
+        stop = num
+    elif len(args) == 1:
+        # start is start, args[0] is stop, step is 1
+        start = _to_tuple(start)
+        stop = _to_tuple(args[0])
+        num = (stop[0] - start[0], stop[1] - start[1])
+    elif len(args) == 2:
+        # start is start, args[0] is stop, args[1] is num
+        start = _to_tuple(start)
+        stop = _to_tuple(args[0])
+        num = _to_tuple(args[1])
+    else:
+        raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
+
+    grid_h = np.linspace(start[0], stop[0], num[0], endpoint=False, dtype=np.float32)
+    grid_w = np.linspace(start[1], stop[1], num[1], endpoint=False, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)   # [2, W, H]
+    return grid
+
+#################################################################################
+#                   Sine/Cosine Positional Embedding Functions                  #
+#################################################################################
+# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+def get_2d_sincos_pos_embed(embed_dim, start, *args, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid = get_meshgrid(start, *args)   # [2, H, w]
+    # grid_h = np.arange(grid_size, dtype=np.float32)
+    # grid_w = np.arange(grid_size, dtype=np.float32)
+    # grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    # grid = np.stack(grid, axis=0)   # [2, W, H]
+
+    grid = grid.reshape([2, 1, *grid.shape[1:]])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)    # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (W,H)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)   # (M, D/2)
+    emb_cos = np.cos(out)   # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+#################################################################################
+#                   Rotary Positional Embedding Functions                       #
+#################################################################################
+# https://github.com/facebookresearch/llama/blob/main/llama/model.py#L443
+
+def get_2d_rotary_pos_embed(embed_dim, start, *args, use_real=True):
+    """
+    This is a 2d version of precompute_freqs_cis, which is a RoPE for image tokens with 2d structure.
+
+    Parameters
+    ----------
+    embed_dim: int
+        embedding dimension size
+    start: int or tuple of int
+        If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, step is 1;
+        If len(args) == 2, start is start, args[0] is stop, args[1] is num.
+    use_real: bool
+        If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+
+    Returns
+    -------
+    pos_embed: torch.Tensor
+        [HW, D/2]
+    """
+    grid = get_meshgrid(start, *args)   # [2, H, w]
+    grid = grid.reshape([2, 1, *grid.shape[1:]])   # Returns a sampling matrix with the same resolution as the target resolution
+    pos_embed = get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=use_real)
+    return pos_embed
+
+
+def get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=False):
+    assert embed_dim % 4 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_rotary_pos_embed(embed_dim // 2, grid[0].reshape(-1), use_real=use_real)  # (H*W, D/4)
+    emb_w = get_1d_rotary_pos_embed(embed_dim // 2, grid[1].reshape(-1), use_real=use_real)  # (H*W, D/4)
+
+    if use_real:
+        cos = torch.cat([emb_h[0], emb_w[0]], dim=1)    # (H*W, D/2)
+        sin = torch.cat([emb_h[1], emb_w[1]], dim=1)    # (H*W, D/2)
+        return cos, sin
+    else:
+        emb = torch.cat([emb_h, emb_w], dim=1)    # (H*W, D/2)
+        return emb
+
+
+def get_1d_rotary_pos_embed(dim: int, pos: Union[np.ndarray, int], theta: float = 10000.0, use_real=False):
+    """
+    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+    The returned tensor contains complex values in complex64 data type.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        pos (np.ndarray, int): Position indices for the frequency tensor. [S] or scalar
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (bool, optional): If True, return real part and imaginary part separately.
+                                   Otherwise, return complex numbers.
+
+    Returns:
+        torch.Tensor: Precomputed frequency tensor with complex exponentials. [S, D/2]
+
+    """
+    if isinstance(pos, int):
+        pos = np.arange(pos)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))  # [D/2]
+    t = torch.from_numpy(pos).to(freqs.device)  # type: ignore  # [S]
+    freqs = torch.outer(t, freqs).float()  # type: ignore   # [S, D/2]
+    if use_real:
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1)  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1)  # [S, D]
+        return freqs_cos, freqs_sin
+    else:
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64     # [S, D/2]
+        return freqs_cis
+
+
+
+def calc_sizes(rope_img, patch_size, th, tw):
+    if rope_img == 'extend':
+        # Expansion mode
+        sub_args = [(th, tw)]
+    elif rope_img.startswith('base'):
+        # Based on the specified dimensions, other dimensions are obtained through interpolation.
+        base_size = int(rope_img[4:]) // 8 // patch_size
+        start, stop = get_fill_resize_and_crop((th, tw), base_size)
+        sub_args = [start, stop, (th, tw)]
+    else:
+        raise ValueError(f"Unknown rope_img: {rope_img}")
+    return sub_args
+
+
+def init_image_posemb(rope_img,
+                      resolutions,
+                      patch_size,
+                      hidden_size,
+                      num_heads,
+                      log_fn,
+                      rope_real=True,
+                      ):
+    freqs_cis_img = {}
+    for reso in resolutions:
+        th, tw = reso.height // 8 // patch_size, reso.width // 8 // patch_size
+        sub_args = calc_sizes(rope_img, patch_size, th, tw)
+        freqs_cis_img[str(reso)] = get_2d_rotary_pos_embed(hidden_size // num_heads, *sub_args, use_real=rope_real)
+        log_fn(f"    Using image RoPE ({rope_img}) ({'real' if rope_real else 'complex'}): {sub_args} | ({reso}) "
+               f"{freqs_cis_img[str(reso)][0].shape if rope_real else freqs_cis_img[str(reso)].shape}")
+    return freqs_cis_img

custom_nodes/comfyui-hydit_min/hydit/modules/text_encoder.py

@@ -0,0 +1,128 @@
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, T5EncoderModel, T5ForConditionalGeneration,T5Config,modeling_utils
+import os
+import comfy.utils
+
+class MT5Embedder(nn.Module):
+    available_models = ["t5-v1_1-xxl"]
+
+    def __init__(
+        self,
+        model_dir="t5-v1_1-xxl",
+        model_kwargs=None,
+        torch_dtype=None,
+        use_tokenizer_only=False,
+        conditional_generation=False,
+        max_length=128,
+        ksampler = False
+    ):
+        super().__init__()
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.torch_dtype = torch_dtype or torch.bfloat16
+        self.max_length = max_length
+        if model_kwargs is None:
+            model_kwargs = {
+                # "low_cpu_mem_usage": True,
+                "torch_dtype": self.torch_dtype,
+            }
+        model_kwargs["device_map"] = {"shared": self.device, "encoder": self.device}
+        if not ksampler:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+            if use_tokenizer_only:
+                return
+            """
+            if conditional_generation:
+                self.model = None
+                self.generation_model = T5ForConditionalGeneration.from_pretrained(
+                    model_dir
+                )
+                return
+            """
+            self.model = T5EncoderModel.from_pretrained(model_dir, **model_kwargs).eval().to(self.torch_dtype)
+        else:
+            tokenizer_path = os.path.join( # local
+				os.path.dirname(os.path.realpath(__file__)), "..", "..",
+				"tokenizer_mt5",
+			)
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+            #self.tokenizer.eval().to(self.device) 
+            if use_tokenizer_only:
+                return
+            """
+            if conditional_generation:
+                self.model = None
+                self.generation_model = T5ForConditionalGeneration.from_pretrained(
+                    model_dir
+                )
+                return
+            """
+            textmodel_json_config = os.path.join(
+				os.path.dirname(os.path.realpath(__file__)), "..", "..",
+				f"config_mt5.json"
+			)
+            config = T5Config.from_json_file(textmodel_json_config)
+            with modeling_utils.no_init_weights():
+                self.model = T5EncoderModel(config)
+            sd = comfy.utils.load_torch_file(model_dir)
+            self.model.load_state_dict(sd, strict=False)
+            self.model.eval().to(self.device)    
+   
+    def get_tokens_and_mask(self, texts):
+        text_tokens_and_mask = self.tokenizer(
+            texts,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        tokens = text_tokens_and_mask["input_ids"][0]
+        mask = text_tokens_and_mask["attention_mask"][0]
+        # tokens = torch.tensor(tokens).clone().detach()
+        # mask = torch.tensor(mask, dtype=torch.bool).clone().detach()
+        return tokens, mask
+
+    def get_text_embeddings(self, texts, attention_mask=True, layer_index=-1):
+        text_tokens_and_mask = self.tokenizer(
+            texts,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+
+        with torch.no_grad():
+            outputs = self.model(
+                input_ids=text_tokens_and_mask["input_ids"].to(self.device),
+                attention_mask=text_tokens_and_mask["attention_mask"].to(self.device)
+                if attention_mask
+                else None,
+                output_hidden_states=True,
+            )
+            text_encoder_embs = outputs["hidden_states"][layer_index].detach()
+
+        return text_encoder_embs, text_tokens_and_mask["attention_mask"].to(self.device)
+
+    @torch.no_grad()
+    def __call__(self, tokens, attention_mask, layer_index=-1):
+        with torch.cuda.amp.autocast():
+            outputs = self.model(
+                input_ids=tokens,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+
+        z = outputs.hidden_states[layer_index].detach()
+        return z
+
+    def general(self, text: str):
+        # input_ids = input_ids = torch.tensor([list(text.encode("utf-8"))]) + num_special_tokens
+        input_ids = self.tokenizer(text, max_length=128).input_ids
+        print(input_ids)
+        outputs = self.generation_model(input_ids)
+        return outputs
+    

custom_nodes/comfyui-hydit_min/hydit/utils/tools.py

@@ -0,0 +1,190 @@
+import random
+import logging
+from pathlib import Path
+import shutil
+
+import numpy as np
+from PIL import Image
+import torch
+import torch.distributed as dist
+from tqdm.auto import tqdm
+import math
+import torch.nn.functional as F
+import os
+
+def get_trainable_params(model):
+    params = model.parameters()
+    params = [p for p in params if p.requires_grad]
+    return params
+
+
+def set_seeds(seed_list, device=None):
+    if isinstance(seed_list, (tuple, list)):
+        seed = sum(seed_list)
+    else:
+        seed = seed_list
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    return torch.Generator(device).manual_seed(seed)
+
+def get_start_epoch(resume_path, ckpt, steps_per_epoch):
+    if 'epoch' in ckpt:
+        start_epoch = ckpt['epoch']
+    else:
+        start_epoch = 0
+    if 'steps' in ckpt:
+        train_steps = ckpt['steps']
+    else:
+        try:
+            train_steps = int(Path(resume_path).stem)
+        except:
+            train_steps = start_epoch * steps_per_epoch
+
+    start_epoch_step = train_steps % steps_per_epoch + 1
+    return start_epoch, start_epoch_step, train_steps
+
+def assert_shape(*args):
+    if len(args) < 2:
+        return
+    cond = True
+    fail_str = f"{args[0] if isinstance(args[0], (list, tuple)) else args[0].shape}"
+    for i in range(1, len(args)):
+        shape1 = args[i] if isinstance(args[i], (list, tuple)) else args[i].shape
+        shape2 = args[i - 1] if isinstance(args[i - 1], (list, tuple)) else args[i - 1].shape
+        cond = cond and (shape1 == shape2)
+        fail_str += f" vs {args[i] if isinstance(args[i], (list, tuple)) else args[i].shape}"
+    assert cond, fail_str
+
+
+def create_logger(logging_dir=None, logging_file=None, ddp=True):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    if not ddp or (ddp and dist.get_rank() == 0):  # real logger
+        if logging_file is not None:
+            file_handler = [logging.FileHandler(logging_file)]
+        elif logging_dir is not None:
+            file_handler = [logging.FileHandler(f"{logging_dir}/log.txt")]
+        else:
+            file_handler = []
+        logging.basicConfig(
+            level=logging.INFO,
+            format='[\033[34m%(asctime)s\033[0m] %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S',
+            handlers=[logging.StreamHandler()] + file_handler
+        )
+        logger = logging.getLogger(__name__)
+    else:
+        logger = logging.getLogger(__name__)
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+def create_exp_folder(args, rank):
+    if rank == 0:
+        os.makedirs(args.results_dir, exist_ok=True)
+    existed_experiments = list(Path(args.results_dir).glob("*dit*"))
+    if len(existed_experiments) == 0:
+        experiment_index = 1
+    else:
+        existed_experiments.sort()
+        print('existed_experiments', existed_experiments)
+        experiment_index = max([int(x.stem.split('-')[0]) for x in existed_experiments]) + 1
+    dist.barrier()
+    model_string_name = args.task_flag if args.task_flag else args.model.replace("/", "-")
+    experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}"       # Create an experiment folder
+    checkpoint_dir = f"{experiment_dir}/checkpoints"                                        # Stores saved model checkpoints
+    if rank == 0:
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        logger = create_logger(experiment_dir)
+        logger.info(f"Experiment directory created at {experiment_dir}")
+    else:
+        logger = create_logger()
+        experiment_dir = ""
+
+    return experiment_dir, checkpoint_dir, logger
+
+
+def model_resume(args, model, ema, logger, len_loader):
+    """
+    Load pretrained weights.
+    """
+    start_epoch = 0
+    start_epoch_step = 0
+    train_steps = 0
+    resume_path = args.resume
+    if not Path(resume_path).exists():
+        raise FileNotFoundError(f"    Cannot find checkpoint from {resume_path}")
+
+    logger.info(f"    Resume deepspeed={args.resume_deepspeed}, "
+                f"Resume split={args.resume_split}, "
+                f"Resume from checkpoint {resume_path}")
+    # Resume model and ema states (not include optimizer states) from a checkpoint saved by Deepspeed version of DIT.
+    if args.resume_deepspeed:
+        assert 'mp_rank_00_model_states.pt' in os.listdir(resume_path), f'    Cannot find dp chkpt from {resume_path}'
+        resume_ckpt = torch.load(os.path.join(resume_path, 'mp_rank_00_model_states.pt'),
+                                 map_location=lambda storage, loc: storage)
+        # Resume main model
+        if args.ema_to_module:
+            logger.info("    Resume main model from the ema states.")
+            model.load_state_dict(resume_ckpt['ema'], strict=args.strict)
+        else:
+            logger.info("    Resume main model from the main states.")
+            model.load_state_dict(resume_ckpt['module'], strict=args.strict)
+        # Resume EMA model
+        if args.use_ema:
+            if args.module_to_ema:
+                logger.info("    Resume EMA model from the main states.")
+                ema.load_state_dict(resume_ckpt['module'], strict=args.strict)
+            else:
+                logger.info("    Resume EMA model from the EMA states.")
+                ema.load_state_dict(resume_ckpt['ema'], strict=args.strict)
+        if not args.reset_loader:
+            start_epoch, start_epoch_step, train_steps = get_start_epoch(args.resume, resume_ckpt, len_loader)
+    # Resume model and ema states (not include optimizer states) from two checkpoints separated from DeepSpeed ckpt.
+    elif args.resume_split:
+        # Resume main model
+        if args.ema_to_module:
+            assert 'pytorch_model_ema.pt' in os.listdir(
+                resume_path), f'    Cannot find pytorch_model_ema.pt from {resume_path}'
+            logger.info(f"    Resume main model from ema states.")
+            resume_ckpt_ema = torch.load(os.path.join(resume_path, 'pytorch_model_ema.pt'),
+                                         map_location=lambda storage, loc: storage)
+            model.load_state_dict(resume_ckpt_ema, strict=args.strict)
+        else:
+            assert 'pytorch_model_module.pt' in os.listdir(
+                resume_path), f'    Cannot find pytorch_model_module.pt from {resume_path}'
+            logger.info(f"    Resume main model from main states.")
+            resume_ckpt_module = torch.load(os.path.join(resume_path, 'pytorch_model_module.pt'),
+                                            map_location=lambda storage, loc: storage)
+            model.load_state_dict(resume_ckpt_module, strict=args.strict)
+        # Resume ema model
+        if args.use_ema:
+            if args.module_to_ema:
+                if "resume_ckpt_module" in locals():
+                    logger.info(f"    Resume ema model from main states.")
+                    ema.load_state_dict(resume_ckpt_module, strict=args.strict)
+                else:
+                    assert 'pytorch_model_module.pt' in os.listdir(
+                        resume_path), f'    Cannot find pytorch_model_module.pt from {resume_path}'
+                    logger.info(f"    Resume ema model from module states.")
+                    resume_ckpt_module = torch.load(os.path.join(resume_path, 'pytorch_model_module.pt'),
+                                                    map_location=lambda storage, loc: storage)
+                    ema.load_state_dict(resume_ckpt_module, strict=args.strict)
+            else:
+                if "resume_ckpt_ema" in locals():
+                    logger.info(f"    Resume ema model from EMA states.")
+                    ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
+                else:
+                    assert 'pytorch_model_ema.pt' in os.listdir(
+                        resume_path), f'    Cannot find pytorch_model_ema.pt from {resume_path}'
+                    logger.info(f"    Resume ema model from EMA states.")
+                    resume_ckpt_ema = torch.load(os.path.join(resume_path, 'pytorch_model_ema.pt'),
+                                                 map_location=lambda storage, loc: storage)
+                    ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
+    else:
+        raise ValueError("    “If `resume` is True, then either `resume_split` or `resume_deepspeed` must be true.”")
+
+    return model, ema, start_epoch, start_epoch_step, train_steps

custom_nodes/comfyui-hydit_min/nodes.py

@@ -0,0 +1,94 @@
+import os
+import torch
+import folder_paths
+from .constant import T5_PATH
+from .dit import load_checkpoint, load_vae
+from .clip import CLIP
+
+
+class DiffusersCheckpointLoader:
+    def __init__(self):
+        self.tmp_dir = folder_paths.get_temp_directory()
+        self.dtype = torch.float32
+        
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { 
+                             "model_name": (folder_paths.get_filename_list("checkpoints"),),  }}
+
+    RETURN_TYPES = ("MODEL", )
+
+    FUNCTION = "load_checkpoint"
+
+    CATEGORY = "HunYuan"
+
+    def load_checkpoint(self, model_name):
+        MODEL_PATH = folder_paths.get_full_path("checkpoints", model_name)
+        out = load_checkpoint(MODEL_PATH)
+        return out
+    
+
+class DiffusersVAELoader:
+    def __init__(self):
+        self.tmp_dir = folder_paths.get_temp_directory()
+        self.dtype = torch.float32
+        
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { 
+                             "model_name": (folder_paths.get_filename_list("vae"),),  }}
+
+    RETURN_TYPES = ("VAE", )
+
+    FUNCTION = "load_vae"
+
+    CATEGORY = "HunYuan"
+
+    def load_vae(self, model_name):
+        MODEL_PATH = folder_paths.get_full_path("vae", model_name)
+        out = load_vae(MODEL_PATH)
+        return out
+    
+
+
+class DiffusersCLIPLoader:
+    def __init__(self):
+        self.tmp_dir = folder_paths.get_temp_directory()
+        self.dtype = torch.float32
+        
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { 
+                             "text_encoder_path": (folder_paths.get_filename_list("clip"),),  
+                             "t5_text_encoder_path": (os.listdir(T5_PATH),), }}
+
+    RETURN_TYPES = ("CLIP", )
+
+    FUNCTION = "load_clip"
+
+    CATEGORY = "HunYuan"
+
+    def load_clip(self, text_encoder_path, t5_text_encoder_path):
+        CLIP_PATH = folder_paths.get_full_path("clip", text_encoder_path)
+        t5_file = os.path.join(T5_PATH, t5_text_encoder_path)
+        root = None
+        out = CLIP(root, CLIP_PATH, t5_file)
+        return (out,)
+    
+
+
+
+NODE_CLASS_MAPPINGS = {
+    "DiffusersCheckpointLoader": DiffusersCheckpointLoader,
+    "DiffusersVAELoader": DiffusersVAELoader,
+    "DiffusersCLIPLoader": DiffusersCLIPLoader,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DiffusersCheckpointLoader": "HunYuan Checkpoint Loader",
+    "DiffusersVAELoader": "HunYuan VAE Loader", 
+    "DiffusersCLIPLoader": "HunYuan CLIP Loader",  
+}

custom_nodes/comfyui-hydit_min/requirements.txt

@@ -0,0 +1,15 @@
+--extra-index-url https://pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/cu117
+timm==1.0.3
+diffusers==0.21.2
+peft==0.10.0
+protobuf==5.27.0
+accelerate==0.30.1
+loguru==0.7.2
+sentencepiece==0.2.0
+cuda-python==12.5.0
+polygraphy==0.49.9
+pandas==2.2.2
+omegaconf==2.3.0
+xformers==0.0.20
+pytorch-lightning==2.2.5

custom_nodes/comfyui-hydit_min/supported_dit_models.py

@@ -0,0 +1,115 @@
+import comfy.supported_models_base
+import comfy.latent_formats
+import comfy.model_patcher
+import comfy.model_base
+import comfy.utils
+import torch
+from collections import namedtuple
+from .hydit.modules.models import HunYuanDiT as HYDiT
+import pdb
+from .hydit.modules.posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
+from .hydit.modules.models import HUNYUAN_DIT_CONFIG
+
+def calc_rope(height, width):
+    th = height // 8 // HUNYUAN_DIT_CONFIG['DiT-g/2']['patch_size']
+    tw = width // 8 // HUNYUAN_DIT_CONFIG['DiT-g/2']['patch_size']
+    base_size = 512 // 8 // HUNYUAN_DIT_CONFIG['DiT-g/2']['patch_size']
+    start, stop = get_fill_resize_and_crop((th, tw), base_size)
+    sub_args = [start, stop, (th, tw)]
+    head_size = HUNYUAN_DIT_CONFIG['DiT-g/2']['hidden_size'] // HUNYUAN_DIT_CONFIG['DiT-g/2']['num_heads']
+    rope = get_2d_rotary_pos_embed(head_size, *sub_args)
+    return rope
+
+def batch_embeddings(embeds, batch_size):
+	bs_embed, seq_len, _ = embeds.shape
+	embeds = embeds.repeat(1, batch_size, 1)
+	embeds = embeds.view(bs_embed * batch_size, seq_len, -1)
+	return embeds
+    
+class HunYuan_DiT(comfy.supported_models_base.BASE):
+    Conf = namedtuple('DiT', ['learn_sigma', 'text_states_dim', 'text_states_dim_t5', 'text_len', 'text_len_t5', 'norm', 'infer_mode', 
+                              'use_fp16', 'use_flash_attn', 'qk_norm', 'size_cond', 'use_style_cond'])
+    conf = {
+        'learn_sigma': True,
+        'text_states_dim': 1024,
+        'text_states_dim_t5': 2048,
+        'text_len': 77,
+        'text_len_t5': 256,
+        'norm': 'layer',
+        'infer_mode': 'torch',
+        'use_fp16': True,
+        'use_flash_attn': False,
+        'qk_norm': True,
+        'size_cond': None, 
+        'use_style_cond': False,
+
+    }
+
+    unet_config = {}
+    unet_extra_config = {
+        "num_heads": 16
+    }
+    latent_format = comfy.latent_formats.SDXL
+
+    dit_conf = Conf(**conf)
+
+    def __init__(self, model_conf):
+        self.unet_config = model_conf.get("unet_config", {})
+        #print(model_conf)
+        print(self.unet_config)
+        self.sampling_settings = model_conf.get("sampling_settings", {})
+        self.latent_format = self.latent_format()
+        self.unet_config["disable_unet_model_creation"] = True
+        #self.unet_config["disable_unet_model_creation"] = self.unet_config.get("disable_unet_model_creation", True)
+
+    def model_type(self, state_dict, prefix=""):
+        return comfy.model_base.ModelType.V_PREDICTION
+    
+class HYDiT_Model(comfy.model_base.BaseModel):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        addit_embeds = kwargs['cross_attn'].addit_embeds
+        for name in addit_embeds:
+            out[name] = comfy.conds.CONDRegular(addit_embeds[name])
+
+        return out
+    
+class ModifiedHunYuanDiT(HYDiT):
+    def forward_core(self, *args, **kwargs):
+        return super().forward(*args, **kwargs)
+
+    def forward(self, x, timesteps, context, t5_embeds=None, attention_mask=None, t5_attention_mask=None, image_meta_size=None, **kwargs):
+        batch_size, _, height, width = x.shape
+        ratio = 8
+        #assert(0)
+        #pdb.set_trace()
+     
+        
+        style = None
+        size_cond = None
+        image_meta_size = None
+        rope = calc_rope(int(ratio*height), int(ratio*width))
+        #import pdb
+        #pdb.set_trace()
+
+
+        noise_pred = self.forward_core(
+            x = x.to(self.dtype),
+            t = timesteps.to(self.dtype),
+            encoder_hidden_states = context.to(self.dtype),
+            text_embedding_mask   = attention_mask.to(self.dtype),
+            encoder_hidden_states_t5 = t5_embeds.to(self.dtype),
+            text_embedding_mask_t5   = t5_attention_mask.to(self.dtype),
+            image_meta_size = image_meta_size,
+            style = style,
+            cos_cis_img = rope[0],
+            sin_cis_img = rope[1],
+            return_dict=False
+        )
+        noise_pred = noise_pred.to(torch.float)
+        eps, _ = noise_pred[:, :self.in_channels], noise_pred[:, self.in_channels:]
+        return eps
+

custom_nodes/comfyui-hydit_min/tokenizer_clip/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

custom_nodes/comfyui-hydit_min/tokenizer_clip/tokenizer_config.json

@@ -0,0 +1,16 @@
+{
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "name_or_path": "hfl/chinese-roberta-wwm-ext",
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "special_tokens_map_file": "/home/chenweifeng/.cache/huggingface/hub/models--hfl--chinese-roberta-wwm-ext/snapshots/5c58d0b8ec1d9014354d691c538661bf00bfdb44/special_tokens_map.json",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]",
+  "model_max_length": 77
+}

custom_nodes/comfyui-hydit_min/tokenizer_mt5/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "/home/patrick/t5/mt5-xl",
+  "architectures": [
+    "MT5ForConditionalGeneration"
+  ],
+  "d_ff": 5120,
+  "d_kv": 64,
+  "d_model": 2048,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "mt5",
+  "num_decoder_layers": 24,
+  "num_heads": 32,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "tokenizer_class": "T5Tokenizer",
+  "transformers_version": "4.10.0.dev0",
+  "use_cache": true,
+  "vocab_size": 250112
+}

custom_nodes/comfyui-hydit_min/tokenizer_mt5/special_tokens_map.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>"}

custom_nodes/comfyui-hydit_min/tokenizer_mt5/tokenizer_config.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>", "extra_ids": 0, "additional_special_tokens": null, "special_tokens_map_file": "/home/patrick/.cache/torch/transformers/685ac0ca8568ec593a48b61b0a3c272beee9bc194a3c7241d15dcadb5f875e53.f76030f3ec1b96a8199b2593390c610e76ca8028ef3d24680000619ffb646276", "tokenizer_file": null, "name_or_path": "google/mt5-small"}

custom_nodes/comfyui-hydit_min/utils.py

@@ -0,0 +1,215 @@
+import io
+import torch
+import requests
+import numpy as np
+from PIL import Image
+from omegaconf import OmegaConf
+from torchvision.transforms import ToTensor
+from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
+    assign_to_checkpoint,
+    conv_attn_to_linear,
+    create_vae_diffusers_config,
+    renew_vae_attention_paths,
+    renew_vae_resnet_paths,
+)
+from diffusers import (
+    AutoencoderKL,
+    DDIMScheduler,
+    DDPMScheduler,
+    DEISMultistepScheduler,
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    KDPM2DiscreteScheduler,
+    UniPCMultistepScheduler,
+)
+
+SCHEDULERS = {
+    'DDIM' : DDIMScheduler,
+    'DDPM' : DDPMScheduler,
+    'DEISMultistep' : DEISMultistepScheduler,
+    'DPMSolverMultistep' : DPMSolverMultistepScheduler,
+    'DPMSolverSinglestep' : DPMSolverSinglestepScheduler,
+    'EulerAncestralDiscrete' : EulerAncestralDiscreteScheduler,
+    'EulerDiscrete' : EulerDiscreteScheduler,
+    'HeunDiscrete' : HeunDiscreteScheduler,
+    'KDPM2AncestralDiscrete' : KDPM2AncestralDiscreteScheduler,
+    'KDPM2Discrete' : KDPM2DiscreteScheduler,
+    'UniPCMultistep' : UniPCMultistepScheduler
+}
+
+SCHEDULERS_hunyuan = ["ddpm", "ddim", "dpmms"]
+
+def token_auto_concat_embeds(pipe, positive, negative):
+    max_length = pipe.tokenizer.model_max_length
+    positive_length = pipe.tokenizer(positive, return_tensors="pt").input_ids.shape[-1]
+    negative_length = pipe.tokenizer(negative, return_tensors="pt").input_ids.shape[-1]
+    
+    print(f'Token length is model maximum: {max_length}, positive length: {positive_length}, negative length: {negative_length}.')
+    if max_length < positive_length or max_length < negative_length:
+        print('Concatenated embedding.')
+        if positive_length > negative_length:
+            positive_ids = pipe.tokenizer(positive, return_tensors="pt").input_ids.to("cuda")
+            negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=positive_ids.shape[-1], return_tensors="pt").input_ids.to("cuda")
+        else:
+            negative_ids = pipe.tokenizer(negative, return_tensors="pt").input_ids.to("cuda")  
+            positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=negative_ids.shape[-1],  return_tensors="pt").input_ids.to("cuda")
+    else:
+        positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=max_length,  return_tensors="pt").input_ids.to("cuda")
+        negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda")
+    
+    positive_concat_embeds = []
+    negative_concat_embeds = []
+    for i in range(0, positive_ids.shape[-1], max_length):
+        positive_concat_embeds.append(pipe.text_encoder(positive_ids[:, i: i + max_length])[0])
+        negative_concat_embeds.append(pipe.text_encoder(negative_ids[:, i: i + max_length])[0])
+    
+    positive_prompt_embeds = torch.cat(positive_concat_embeds, dim=1)
+    negative_prompt_embeds = torch.cat(negative_concat_embeds, dim=1)
+    return positive_prompt_embeds, negative_prompt_embeds
+
+# Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py
+def custom_convert_ldm_vae_checkpoint(checkpoint, config):
+    vae_state_dict = checkpoint
+
+    new_checkpoint = {}
+
+    new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
+    new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
+    new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+    new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
+    new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
+    new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+
+    new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
+    new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
+    new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+    new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
+    new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
+    new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+
+    new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+    new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+    new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
+    new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
+
+    # Retrieves the keys for the encoder down blocks only
+    num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
+    down_blocks = {
+        layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)
+    }
+
+    # Retrieves the keys for the decoder up blocks only
+    num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
+    up_blocks = {
+        layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)
+    }
+
+    for i in range(num_down_blocks):
+        resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
+
+        if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
+                f"encoder.down.{i}.downsample.conv.weight"
+            )
+            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
+                f"encoder.down.{i}.downsample.conv.bias"
+            )
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    conv_attn_to_linear(new_checkpoint)
+
+    for i in range(num_up_blocks):
+        block_id = num_up_blocks - 1 - i
+        resnets = [
+            key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+        ]
+
+        if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+                f"decoder.up.{block_id}.upsample.conv.weight"
+            ]
+            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+                f"decoder.up.{block_id}.upsample.conv.bias"
+            ]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    conv_attn_to_linear(new_checkpoint)
+    return new_checkpoint
+
+# Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py
+def vae_pt_to_vae_diffuser(
+    checkpoint_path: str,
+    output_path: str,
+):
+    # Only support V1
+    r = requests.get(
+        " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"
+    )
+    io_obj = io.BytesIO(r.content)
+
+    original_config = OmegaConf.load(io_obj)
+    image_size = 512
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if checkpoint_path.endswith("safetensors"):
+        from safetensors import safe_open
+
+        checkpoint = {}
+        with safe_open(checkpoint_path, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                checkpoint[key] = f.get_tensor(key)
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location=device)["state_dict"]
+
+    # Convert the VAE model.
+    vae_config = create_vae_diffusers_config(original_config, image_size=image_size)
+    converted_vae_checkpoint = custom_convert_ldm_vae_checkpoint(checkpoint, vae_config)
+
+    vae = AutoencoderKL(**vae_config)
+    vae.load_state_dict(converted_vae_checkpoint)
+    vae.save_pretrained(output_path)
+
+
+def convert_images_to_tensors(images: list[Image.Image]):
+    return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images])
+
+def convert_tensors_to_images(images: torch.tensor):
+    return [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in images]
+
+def resize_images(images: list[Image.Image], size: tuple[int, int]):
+    return [image.resize(size) for image in images]

custom_nodes/comfyui-hydit_min/workflow/workflow_ksampler.json

@@ -0,0 +1,424 @@
+{
+  "last_node_id": 33,
+  "last_link_id": 9,
+  "nodes": [
+    {
+      "id": 33,
+      "type": "DiffusersCLIPLoader",
+      "pos": [
+        -77,
+        221
+      ],
+      "size": {
+        "0": 315,
+        "1": 82
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "links": [
+            8,
+            9
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "DiffusersCLIPLoader"
+      },
+      "widgets_values": [
+        "pytorch_model.bin",
+        "pytorch_model.bin"
+      ]
+    },
+    {
+      "id": 25,
+      "type": "CLIPTextEncode",
+      "pos": [
+        260,
+        108
+      ],
+      "size": {
+        "0": 400,
+        "1": 200
+      },
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": 8
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "links": [
+            1
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CLIPTextEncode"
+      },
+      "widgets_values": [
+        "描绘的风格是写实，画面主要描述一双泥泞的靴子在雨天里，靴子颜色是棕色，泥沙溅在 Boot 的表面，背景是湿漉漉的地板，泥泞的环境，天气是阴沉的雨天，镜头是近景"
+      ]
+    },
+    {
+      "id": 32,
+      "type": "DiffusersVAELoader",
+      "pos": [
+        359,
+        498
+      ],
+      "size": {
+        "0": 303.9019775390625,
+        "1": 58
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "links": [
+            7
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "DiffusersVAELoader"
+      },
+      "widgets_values": [
+        "diffusion_pytorch_model.bin"
+      ]
+    },
+    {
+      "id": 26,
+      "type": "CLIPTextEncode",
+      "pos": [
+        -80,
+        356
+      ],
+      "size": {
+        "0": 400,
+        "1": 200
+      },
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": 9
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "links": [
+            2
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CLIPTextEncode"
+      },
+      "widgets_values": [
+        "错误的眼睛，糟糕的人脸，毁容，糟糕的艺术，变形，多余的肢体，模糊的颜色，模糊，重复，病态，残缺，"
+      ]
+    },
+    {
+      "id": 28,
+      "type": "VAEDecode",
+      "pos": [
+        731,
+        159
+      ],
+      "size": {
+        "0": 210,
+        "1": 46
+      },
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "samples",
+          "type": "LATENT",
+          "link": 4
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            5
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAEDecode"
+      }
+    },
+    {
+      "id": 24,
+      "type": "KSampler",
+      "pos": [
+        681,
+        290
+      ],
+      "size": {
+        "0": 315,
+        "1": 262
+      },
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "MODEL",
+          "link": 6,
+          "slot_index": 0
+        },
+        {
+          "name": "positive",
+          "type": "CONDITIONING",
+          "link": 1
+        },
+        {
+          "name": "negative",
+          "type": "CONDITIONING",
+          "link": 2
+        },
+        {
+          "name": "latent_image",
+          "type": "LATENT",
+          "link": 3
+        }
+      ],
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            4
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "KSampler"
+      },
+      "widgets_values": [
+        123,
+        "fixed",
+        30,
+        6,
+        "ddim",
+        "ddim_uniform",
+        1
+      ]
+    },
+    {
+      "id": 31,
+      "type": "DiffusersCheckpointLoader",
+      "pos": [
+        -79,
+        107
+      ],
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "links": [
+            6
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "DiffusersCheckpointLoader"
+      },
+      "widgets_values": [
+        "pytorch_model_distill.pt"
+      ]
+    },
+    {
+      "id": 27,
+      "type": "EmptyLatentImage",
+      "pos": [
+        351,
+        350
+      ],
+      "size": {
+        "0": 315,
+        "1": 106
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            3
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "EmptyLatentImage"
+      },
+      "widgets_values": [
+        1024,
+        1280,
+        2
+      ]
+    },
+    {
+      "id": 29,
+      "type": "PreviewImage",
+      "pos": [
+        1031,
+        175
+      ],
+      "size": [
+        357.9370410924005,
+        264.0938661488955
+      ],
+      "flags": {},
+      "order": 8,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 5
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      }
+    }
+  ],
+  "links": [
+    [
+      1,
+      25,
+      0,
+      24,
+      1,
+      "CONDITIONING"
+    ],
+    [
+      2,
+      26,
+      0,
+      24,
+      2,
+      "CONDITIONING"
+    ],
+    [
+      3,
+      27,
+      0,
+      24,
+      3,
+      "LATENT"
+    ],
+    [
+      4,
+      24,
+      0,
+      28,
+      0,
+      "LATENT"
+    ],
+    [
+      5,
+      28,
+      0,
+      29,
+      0,
+      "IMAGE"
+    ],
+    [
+      6,
+      31,
+      0,
+      24,
+      0,
+      "MODEL"
+    ],
+    [
+      7,
+      32,
+      0,
+      28,
+      1,
+      "VAE"
+    ],
+    [
+      8,
+      33,
+      0,
+      25,
+      0,
+      "CLIP"
+    ],
+    [
+      9,
+      33,
+      0,
+      26,
+      0,
+      "CLIP"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.8954302432552415,
+      "offset": [
+        370.7466007037628,
+        90.3483300454219
+      ]
+    }
+  },
+  "version": 0.4
+}