From 327bea1eebde1f3d2b5e7d3b0fb0f018e4c8faa1 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 10 Mar 2024 23:51:59 +0300 Subject: [PATCH] NNCF force eval and fix embeddings --- modules/sd_models_compile.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 7e1bf1da7..bc114fd18 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -1,3 +1,4 @@ +import copy import time import logging import torch @@ -78,7 +79,11 @@ def nncf_compress_weights(sd_model): def nncf_compress_model(model): return_device = model.device - model = nncf.compress_weights(model.to(devices.device)).to(return_device) + if hasattr(model, "get_input_embeddings"): + backup_embeddings = copy.deepcopy(model.get_input_embeddings()) + model = nncf.compress_weights(model.eval().to(devices.device)).to(return_device) + if hasattr(model, "set_input_embeddings"): + model.set_input_embeddings(backup_embeddings) devices.torch_gc(force=True) return model