298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-07-01 21:07:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

SaveLora now saves in the same filename format as all the other nodes. (#8538)

Browse Source
This commit is contained in:
comfyanonymous 2025-06-15 00:44:59 -07:00 committed by GitHub
parent d6a2137fc3
commit ae75a084df
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
comfy_extras/nodes_train.py
View File

@ -552,6 +552,9 @@ class LoraModelLoader:
class SaveLoRA: class SaveLoRA:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
@ -565,7 +568,7 @@ class SaveLoRA:
"prefix": ( "prefix": (
"STRING", "STRING",
{ {
"default": "trained_lora", "default": "loras/ComfyUI_trained_lora",
"tooltip": "The prefix to use for the saved LoRA file.", "tooltip": "The prefix to use for the saved LoRA file.",
}, },
), ),
@ -588,12 +591,13 @@ class SaveLoRA:
OUTPUT_NODE = True OUTPUT_NODE = True
def save(self, lora, prefix, steps=None): def save(self, lora, prefix, steps=None):
date = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(prefix, self.output_dir)
if steps is None: if steps is None:
output_file = f"models/loras/{prefix}_{date}_lora.safetensors" output_checkpoint = f"{filename}_{counter:05}_.safetensors"
else: else:
output_file = f"models/loras/{prefix}_{steps}_steps_{date}_lora.safetensors" output_checkpoint = f"{filename}_{steps}_steps_{counter:05}_.safetensors"
safetensors.torch.save_file(lora, output_file) output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
safetensors.torch.save_file(lora, output_checkpoint)
return {} return {}