From 599fafe01f741e356491024ad60b9447f9befe9e Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 12 Aug 2024 09:57:16 -0600 Subject: [PATCH] Allow user to have the full flux checkpoint local --- toolkit/stable_diffusion_model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 35c0a37c..0896080f 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -465,6 +465,11 @@ class StableDiffusion: if os.path.exists(transformer_path): subfolder = None 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_path,