Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

NNCF force eval and fix embeddings

Browse Source
This commit is contained in:
Disty0
2024-03-10 23:51:59 +03:00
parent 70ede51166
commit 327bea1eeb
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_models_compile.py
+6 -1
View File
@@ -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