nodes/kijai-ComfyUI-FluxTrainer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add lora path to load network_weights in the model select node

Browse Source
This commit is contained in:
kijai
2024-09-08 16:20:40 +03:00
parent 95de0cbcc9
commit 284bb6f112
1 changed files with 14 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+14 -7
View File
@@ -37,11 +37,14 @@ class FluxTrainModelSelect:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"transformer": (folder_paths.get_filename_list("unet"), ),
"vae": (folder_paths.get_filename_list("vae"), ),
"clip_l": (folder_paths.get_filename_list("clip"), ),
"t5": (folder_paths.get_filename_list("clip"), ),
},
"transformer": (folder_paths.get_filename_list("unet"), ),
"vae": (folder_paths.get_filename_list("vae"), ),
"clip_l": (folder_paths.get_filename_list("clip"), ),
"t5": (folder_paths.get_filename_list("clip"), ),
},
"optional": {
"lora_path": ("STRING",{"multiline": True, "forceInput": True, "default": "", "tooltip": "pre-trained LoRA path to load (network_weights)"}),
}
}
RETURN_TYPES = ("TRAIN_FLUX_MODELS",)
@@ -49,7 +52,7 @@ class FluxTrainModelSelect:
FUNCTION = "loadmodel"
CATEGORY = "FluxTrainer"
def loadmodel(self, transformer, vae, clip_l, t5):
def loadmodel(self, transformer, vae, clip_l, t5, lora_path=None):
transformer_path = folder_paths.get_full_path("unet", transformer)
vae_path = folder_paths.get_full_path("vae", vae)
@@ -60,7 +63,8 @@ class FluxTrainModelSelect:
"transformer": transformer_path,
"vae": vae_path,
"clip_l": clip_path,
"t5": t5_path
"t5": t5_path,
"lora_path": lora_path if lora_path is not None else None
}
return (flux_models,)
@@ -446,6 +450,9 @@ class InitFluxLoRATraining:
else:
config_dict["gradient_checkpointing"] = True
if flux_models["lora_path"]:
config_dict["network_weights"] = flux_models["lora_path"]
config_dict.update(kwargs)
config_dict.update(optimizer_settings)