ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 19:21:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix issue with saving chroma full finetune.

Browse Source
This commit is contained in:
Jaret Burkett
2025-05-28 07:42:30 -06:00
parent 79499fa795
commit 79bb9be92b
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -411,6 +411,8 @@ class ChromaModel(BaseModel):
return self.text_encoder[1].encoder.block[0].layer[0].SelfAttention.q.weight.requires_grad return self.text_encoder[1].encoder.block[0].layer[0].SelfAttention.q.weight.requires_grad
def save_model(self, output_path, meta, save_dtype): def save_model(self, output_path, meta, save_dtype):
if not output_path.endswith(".safetensors"):
output_path = output_path + ".safetensors"
# only save the unet # only save the unet
transformer: Chroma = unwrap_model(self.model) transformer: Chroma = unwrap_model(self.model)
state_dict = transformer.state_dict() state_dict = transformer.state_dict()