diff --git a/HunYuanDiT/loader.py b/HunYuanDiT/loader.py index 35940de..942b267 100644 --- a/HunYuanDiT/loader.py +++ b/HunYuanDiT/loader.py @@ -73,5 +73,6 @@ def load_hydit(model_path, model_conf): load_device = load_device, offload_device = offload_device, current_device = "cpu", + size = 6 * (1024**3), ) return model_patcher diff --git a/HunYuanDiT/nodes.py b/HunYuanDiT/nodes.py index 3b17c29..8bc39f5 100644 --- a/HunYuanDiT/nodes.py +++ b/HunYuanDiT/nodes.py @@ -97,6 +97,7 @@ class HYDiTTextEncode: def encode(self, text, text_t5, CLIP, T5): # T5 + T5.load_model() t5_pre = T5.tokenizer( text, max_length = T5.cond_stage_model.max_length, @@ -117,6 +118,7 @@ class HYDiTTextEncode: t5_embs = t5_outs["hidden_states"][-1].float().cpu() # "clip" + CLIP.load_model() clip_pre = CLIP.tokenizer( text, max_length = CLIP.cond_stage_model.max_length, diff --git a/HunYuanDiT/tenc.py b/HunYuanDiT/tenc.py index 311543b..b432c2d 100644 --- a/HunYuanDiT/tenc.py +++ b/HunYuanDiT/tenc.py @@ -35,7 +35,7 @@ class mT5Model(torch.nn.Module): return self.transformer.load_state_dict(sd, strict=False) def to(self, *args, **kwargs): - self.transformer.to(*args, **kwargs) + return 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): @@ -64,47 +64,50 @@ class hyCLIPModel(torch.nn.Module): return self.transformer.load_state_dict(sd, strict=False) def to(self, *args, **kwargs): - self.transformer.to(*args, **kwargs) + return 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 + size = 8 if model_class == "mT5" else 2 + if dtype == torch.float32: + size *= 2 + size *= (1024**3) + 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 + size = 0 # doesn't matter 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 + size = 0 # not used 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 + self.device = self.load_device if model_class == "mT5": self.cond_stage_model = mT5Model( - device = device, - dtype = dtype, + device = self.load_device, + dtype = self.dtype, ) tokenizer_args = {"subfolder": "t2i/mt5"} else: self.cond_stage_model = hyCLIPModel( - device = device, - dtype = dtype, + device = self.load_device, + dtype = self.dtype, ) tokenizer_args = {"subfolder": "t2i/tokenizer",} self.tokenizer = AutoTokenizer.from_pretrained(