github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Adjusted some names and allow to ignore text encoder on extract

Browse Source
This commit is contained in:
Jaret Burkett
2023-07-16 12:52:39 -06:00
parent a61914ddd9
commit 9df7af1720
5 changed files with 39 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
jobs/process/BaseExtractProcess.py
View File

@@ -30,6 +30,8 @@ class BaseExtractProcess(BaseProcess):
self.config = config
self.dtype = self.get_conf('dtype', self.job.dtype)
self.torch_dtype = get_torch_dtype(self.dtype)
self.extract_unet = self.get_conf('extract_unet', self.job.extract_unet)
self.extract_text_encoder = self.get_conf('extract_text_encoder', self.job.extract_text_encoder)
def run(self):
# here instead of init because child init needs to go first

8
jobs/process/ExtractLoconProcess.py
View File

@@ -45,15 +45,17 @@ class ExtractLoconProcess(BaseExtractProcess):
print(f"Running process: {self.mode}, lin: {self.linear_param}, conv: {self.conv_param}")
state_dict, extract_diff_meta = extract_diff(
self.job.base_model,
self.job.extract_model,
self.job.model_base,
self.job.model_extract,
self.mode,
self.linear_param,
self.conv_param,
self.job.device,
self.use_sparse_bias,
self.sparsity,
not self.disable_cp
not self.disable_cp,
extract_unet=self.extract_unet,
extract_text_encoder=self.extract_text_encoder
)
self.add_meta(extract_diff_meta)

6
jobs/process/ExtractLoraProcess.py
View File

@@ -44,8 +44,8 @@ class ExtractLoraProcess(BaseExtractProcess):
print(f"Running process: {self.mode}, dim: {self.dim}")
state_dict, extract_diff_meta = extract_diff(
self.job.base_model,
self.job.extract_model,
self.job.model_base,
self.job.model_extract,
self.mode,
self.dim,
0,
@@ -54,6 +54,8 @@ class ExtractLoraProcess(BaseExtractProcess):
self.sparsity,
small_conv=False,
linear_only=True,
extract_unet=self.extract_unet,
extract_text_encoder=self.extract_text_encoder
)
self.add_meta(extract_diff_meta)