nodes/MinusZoneAI-ComfyUI-Kolors-MZ
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

'update'

Browse Source
This commit is contained in:
wailovet
2024-07-10 23:56:48 +08:00
parent 52e3b37ee1
commit 211e2ebfa9
1 changed files with 12 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mz_kolors_core.py
+12 -26
View File
@@ -97,33 +97,19 @@ def MZ_ChatGLM3Loader_call(args):
except:
pass
text_encoder_resp = Utils.cache_get("chatglm3_text_encoder")
if text_encoder_resp is None or text_encoder_resp['chatglm3_checkpoint'] != chatglm3_checkpoint:
if text_encoder_resp is not None and text_encoder_resp['chatglm3_checkpoint'] != chatglm3_checkpoint:
del text_encoder_resp
gc.collect()
with (init_empty_weights() if is_accelerate_available else nullcontext()):
text_encoder = ChatGLMModel(text_encoder_config)
if '4bit' in chatglm3_checkpoint:
text_encoder.quantize(4)
elif '8bit' in chatglm3_checkpoint:
text_encoder.quantize(8)
text_encoder_sd = load_torch_file(chatglm3_checkpoint_path)
if is_accelerate_available:
for key in text_encoder_sd:
set_module_tensor_to_device(
text_encoder, key, device=offload_device, value=text_encoder_sd[key])
else:
text_encoder.load_state_dict()
Utils.cache_set("chatglm3_text_encoder", {
"text_encoder": text_encoder,
"chatglm3_checkpoint": chatglm3_checkpoint
})
with (init_empty_weights() if is_accelerate_available else nullcontext()):
text_encoder = ChatGLMModel(text_encoder_config).eval()
if '4bit' in chatglm3_checkpoint:
text_encoder.quantize(4)
elif '8bit' in chatglm3_checkpoint:
text_encoder.quantize(8)
text_encoder_sd = load_torch_file(chatglm3_checkpoint_path)
if is_accelerate_available:
for key in text_encoder_sd:
set_module_tensor_to_device(
text_encoder, key, device=offload_device, value=text_encoder_sd[key])
else:
text_encoder = text_encoder_resp['text_encoder']
text_encoder.load_state_dict()
tokenizer_path = os.path.join(
os.path.dirname(__file__), 'configs', "tokenizer")