|
import os |
|
import folder_paths |
|
|
|
class SaveText: |
|
|
|
def __init__(self): |
|
self.output_dir = folder_paths.get_output_directory() |
|
self.type = "output" |
|
|
|
@classmethod |
|
def INPUT_TYPES(cls): |
|
return { |
|
'required': { |
|
'filename_prefix': ('STRING', { 'default': 'ComfyUI' }), |
|
'ext': ('STRING', { 'default': 'txt' }), |
|
'text': ('STRING', { 'multiline': True, 'default': '' }), |
|
} |
|
} |
|
|
|
OUTPUT_NODE = True |
|
|
|
RETURN_TYPES = () |
|
|
|
FUNCTION = 'execute' |
|
|
|
CATEGORY = 'utils' |
|
|
|
def execute(self, filename_prefix: str, ext: str, text: str): |
|
def map_filename(filename): |
|
prefix_len = len(os.path.basename(filename_prefix)) |
|
prefix = filename[:prefix_len + 1] |
|
try: |
|
digits = int(filename[prefix_len + 1:].split('_')[0]) |
|
except: |
|
digits = 0 |
|
return (digits, prefix) |
|
|
|
subfolder = os.path.dirname(os.path.normpath(filename_prefix)) |
|
filename = os.path.basename(os.path.normpath(filename_prefix)) |
|
|
|
full_output_folder = os.path.join(self.output_dir, subfolder) |
|
|
|
if os.path.commonpath((self.output_dir, os.path.abspath(full_output_folder))) != self.output_dir: |
|
print("Saving image outside the output folder is not allowed.") |
|
return {} |
|
|
|
try: |
|
counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1 |
|
except ValueError: |
|
counter = 1 |
|
except FileNotFoundError: |
|
os.makedirs(full_output_folder, exist_ok=True) |
|
counter = 1 |
|
|
|
if ext is None or len(ext) == 0 or ext == '.': |
|
ext = '.txt' |
|
if not ext.startswith('.'): |
|
ext = '.' + ext |
|
|
|
file = f"{filename}_{counter:05}_{ext}" |
|
with open(os.path.join(full_output_folder, file), 'w') as io: |
|
io.write(text) |
|
counter += 1 |
|
|
|
return {} |
|
|