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

Allow user to have the full flux checkpoint local

Browse Source
This commit is contained in:
Jaret Burkett
2024-08-12 09:57:16 -06:00
parent af108bb964
commit 599fafe01f
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/stable_diffusion_model.py
View File

@@ -465,6 +465,11 @@ class StableDiffusion:
if os.path.exists(transformer_path): if os.path.exists(transformer_path):
subfolder = None subfolder = None
transformer_path = os.path.join(transformer_path, 'transformer') transformer_path = os.path.join(transformer_path, 'transformer')
# check if the path is a full checkpoint.
te_folder_path = os.path.join(model_path, 'text_encoder')
# if we have the te, this folder is a full checkpoint, use it as the base
if os.path.exists(te_folder_path):
base_model_path = model_path
transformer = FluxTransformer2DModel.from_pretrained( transformer = FluxTransformer2DModel.from_pretrained(
transformer_path, transformer_path,