HunYuanDiT Initial proof of concept

#35

HunYuanDiT/tenc.py

@@ -0,0 +1,168 @@
+# This is for loading the CLIP (bert?) + mT5 encoder for HunYuanDiT
+import os
+import torch
+from transformers import AutoTokenizer, modeling_utils
+from transformers import T5Config, T5EncoderModel, BertConfig, BertModel
+
+import comfy.model_patcher
+import comfy.utils
+
+class mT5Model(torch.nn.Module):
+	def __init__(self, textmodel_json_config=None, device="cpu", max_length=256, freeze=True, dtype=None):
+		super().__init__()
+		self.device = device
+		self.dtype = dtype
+		self.max_length = max_length
+		if textmodel_json_config is None:
+			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.transformer = T5EncoderModel(config)
+		self.to(dtype)
+		if freeze:
+			self.freeze()
+
+	def freeze(self):
+		self.transformer = self.transformer.eval()
+		for param in self.parameters():
+			param.requires_grad = False
+
+	def load_sd(self, sd):
+		return self.transformer.load_state_dict(sd, strict=False)
+
+	def to(self, *args, **kwargs):
+		self.transformer.to(*args, **kwargs)
+
+class hyCLIPModel(torch.nn.Module):
+	def __init__(self, textmodel_json_config=None, device="cpu", max_length=77, freeze=True, dtype=None):
+		super().__init__()
+		self.device = device
+		self.dtype = dtype
+		self.max_length = max_length
+		if textmodel_json_config is None:
+			textmodel_json_config = os.path.join(
+				os.path.dirname(os.path.realpath(__file__)),
+				f"config_clip.json"
+			)
+		config = BertConfig.from_json_file(textmodel_json_config)
+		with modeling_utils.no_init_weights():
+			self.transformer = BertModel(config)
+		self.to(dtype)
+		if freeze:
+			self.freeze()
+
+	def freeze(self):
+		self.transformer = self.transformer.eval()
+		for param in self.parameters():
+			param.requires_grad = False
+
+	def load_sd(self, sd):
+		return self.transformer.load_state_dict(sd, strict=False)
+
+	def to(self, *args, **kwargs):
+		self.transformer.to(*args, **kwargs)
+
+class EXM_HyDiT_Tenc_Temp:
+	def __init__(self, no_init=False, device="cpu", dtype=None, model_class="mT5", *kwargs):
+		if no_init:
+			return
+
+		if device == "auto":
+			size = 0
+			self.load_device = model_management.text_encoder_device()
+			self.offload_device = model_management.text_encoder_offload_device()
+			self.init_device = "cpu"
+		elif device == "cpu":
+			size = 0
+			self.load_device = "cpu"
+			self.offload_device = "cpu"
+			self.init_device="cpu"
+		elif device.startswith("cuda"):
+			print("Direct CUDA device override!\nVRAM will not be freed by default.")
+			size = 0
+			self.load_device = device
+			self.offload_device = device
+			self.init_device = device
+		else:
+			size = 0
+			self.load_device = model_management.get_torch_device()
+			self.offload_device = "cpu"
+			self.init_device="cpu"
+
+		self.dtype = dtype
+		self.device = device
+		if model_class == "mT5":
+			self.cond_stage_model = mT5Model(
+				device         = device,
+				dtype          = dtype,
+			)
+			tokenizer_args = {"subfolder": "t2i/mt5"}
+		else:
+			self.cond_stage_model = hyCLIPModel(
+				device         = device,
+				dtype          = dtype,
+			)
+			tokenizer_args = {"subfolder": "t2i/tokenizer",}
+		self.tokenizer = AutoTokenizer.from_pretrained(
+			"Tencent-Hunyuan/HunyuanDiT",
+			**tokenizer_args
+		)
+		self.patcher = comfy.model_patcher.ModelPatcher(
+			self.cond_stage_model,
+			load_device    = self.load_device,
+			offload_device = self.offload_device,
+			current_device = self.load_device,
+			size           = size,
+		)
+
+	def clone(self):
+		n = EXM_HyDiT_Tenc_Temp(no_init=True)
+		n.patcher = self.patcher.clone()
+		n.cond_stage_model = self.cond_stage_model
+		n.tokenizer = self.tokenizer
+		return n
+
+	def load_sd(self, sd):
+		return self.cond_stage_model.load_sd(sd)
+
+	def get_sd(self):
+		return self.cond_stage_model.state_dict()
+
+	def load_model(self):
+		if self.load_device != "cpu":
+			model_management.load_model_gpu(self.patcher)
+		return self.patcher
+
+	def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
+		return self.patcher.add_patches(patches, strength_patch, strength_model)
+
+	def get_key_patches(self):
+		return self.patcher.get_key_patches()
+
+def load_clip(model_path, **kwargs):
+	model = EXM_HyDiT_Tenc_Temp(model_class="clip", **kwargs)
+	sd = comfy.utils.load_torch_file(model_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 = model.load_sd(state_dict)
+	if len(m) > 0 or len(e) > 0:
+		print(f"HYDiT: clip missing {len(m)} keys ({len(e)} extra)")
+	return model
+
+def load_t5(model_path, **kwargs):
+	model = EXM_HyDiT_Tenc_Temp(model_class="mT5", **kwargs)
+	sd = comfy.utils.load_torch_file(model_path)
+	m, e = model.load_sd(sd)
+	if len(m) > 0 or len(e) > 0:
+		print(f"HYDiT: mT5 missing {len(m)} keys ({len(e)} extra)")
+	return model