|
import os |
|
import re |
|
from typing import Union |
|
import torch |
|
import folder_paths |
|
|
|
class SaveStateDict: |
|
|
|
@classmethod |
|
def INPUT_TYPES(cls): |
|
return { |
|
'required': { |
|
'weights': ('DICT',), |
|
'filename': ('STRING', { 'multiline': False, 'default': 'merged_model.safetensors' }), |
|
'overwrite': (['False', 'True'],), |
|
} |
|
} |
|
|
|
OUTPUT_NODE = True |
|
|
|
RETURN_TYPES = () |
|
|
|
FUNCTION = 'execute' |
|
|
|
CATEGORY = 'model' |
|
|
|
def execute(self, weights: dict, filename: str, overwrite: Union[str,bool]): |
|
if isinstance(overwrite, str): |
|
overwrite = overwrite.lower() == 'true' |
|
|
|
subdir = os.path.dirname(os.path.normpath(filename)) |
|
basename = os.path.basename(os.path.normpath(filename)) |
|
|
|
output_dir = folder_paths.get_output_directory() |
|
full_output_dir = os.path.join(output_dir, subdir) |
|
|
|
if os.path.commonpath((output_dir, os.path.realpath(full_output_dir))) != output_dir: |
|
print('Saving image outside the output folder is not allowed.') |
|
return {} |
|
|
|
full_path = os.path.join(full_output_dir, basename) |
|
|
|
if os.path.exists(full_path): |
|
print(f'{full_path} already exists.') |
|
if overwrite: |
|
print(f'overwriting: {full_path}') |
|
else: |
|
base, ext = os.path.splitext(basename) |
|
|
|
def replace(m: re.Match): |
|
x = m.group(0) |
|
if len(x) == 0: |
|
return '0' |
|
else: |
|
n = int(x) |
|
return str(n + 1) |
|
|
|
base_renamed = re.sub(r'\d+$|$', replace, base) |
|
full_path_renamed = os.path.join(full_output_dir, base_renamed + ext) |
|
print(f'rename {full_path} -> {full_path_renamed}') |
|
full_path = full_path_renamed |
|
|
|
os.makedirs(full_output_dir, exist_ok=True) |
|
|
|
print(f'Saving the state_dict to {full_path}') |
|
|
|
ext = os.path.splitext(full_path)[1] |
|
saver = None |
|
|
|
if ext == '.pt' or ext == '.ckpt': |
|
saver = self.save_torch |
|
elif ext == '.safetensor' or ext == '.safetensors': |
|
saver = self.save_safetensors |
|
else: |
|
full_path += '.safetensors' |
|
saver = self.save_safetensors |
|
|
|
saver(weights, full_path) |
|
|
|
return {} |
|
|
|
def save_torch(self, model: dict, path: str): |
|
torch.save(model, path) |
|
|
|
def save_safetensors(self, model: dict, path: str): |
|
import safetensors.torch |
|
safetensors.torch.save_file(model, path) |
|
|
