diff --git a/mz_kolors_core.py b/mz_kolors_core.py index 1865690..1cc5f33 100644 --- a/mz_kolors_core.py +++ b/mz_kolors_core.py @@ -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")