124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
import torch
|
|
import os
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
|
|
class IteratorListFiles:
|
|
"""Lists files in a directory, returning paths as a list of strings.
|
|
Get the list of files to process."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"extension": ("STRING", {"default": "png, jpg, jpeg, webp, bmp, txt, json, csv, md"}),
|
|
"limit": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("file_paths",)
|
|
FUNCTION = "list_files"
|
|
CATEGORY = "Iterator"
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
def list_files(self, directory, extension="png, jpg, jpeg, webp, bmp, txt, json, csv, md", limit=0):
|
|
if not os.path.isdir(directory):
|
|
print(f"Directory not found: {directory}")
|
|
return ([],)
|
|
|
|
extensions = [e.strip().lower() for e in extension.split(",")]
|
|
|
|
files = []
|
|
try:
|
|
for f in os.listdir(directory):
|
|
ext = os.path.splitext(f)[1].lower().replace(".", "")
|
|
if ext in extensions:
|
|
files.append(os.path.join(directory, f))
|
|
except Exception as e:
|
|
print(f"Error listing files: {e}")
|
|
return ([],)
|
|
|
|
files.sort()
|
|
|
|
if limit > 0:
|
|
files = files[:limit]
|
|
|
|
return (files,)
|
|
|
|
class IteratorLoadImage:
|
|
"""Loads a single image from a file path.
|
|
Load the content (Image)."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"file_path": ("STRING", {"default": "", "forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "filename")
|
|
FUNCTION = "load_image"
|
|
CATEGORY = "Iterator"
|
|
|
|
def load_image(self, file_path):
|
|
if not os.path.isfile(file_path):
|
|
raise FileNotFoundError(f"File not found: {file_path}")
|
|
|
|
i = Image.open(file_path)
|
|
i = ImageOps.exif_transpose(i)
|
|
if i.mode == 'I':
|
|
i = i.point(lambda x: x * (1 / 255))
|
|
image = i.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
filename = os.path.basename(file_path)
|
|
return (image, filename)
|
|
|
|
class IteratorLoadString:
|
|
"""Loads text content from a single file.
|
|
Load the content (String)."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"file_path": ("STRING", {"default": "", "forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("content", "filename")
|
|
FUNCTION = "load_string"
|
|
CATEGORY = "Iterator"
|
|
|
|
def load_string(self, file_path):
|
|
if not os.path.isfile(file_path):
|
|
raise FileNotFoundError(f"File not found: {file_path}")
|
|
|
|
# Try multiple encodings
|
|
encodings = ['utf-8-sig', 'utf-16', 'shift_jis', 'cp932', 'euc_jp', 'cp1252', 'latin-1']
|
|
content = None
|
|
|
|
for enc in encodings:
|
|
try:
|
|
with open(file_path, 'r', encoding=enc) as f:
|
|
content = f.read()
|
|
break
|
|
except (UnicodeDecodeError, UnicodeError):
|
|
continue
|
|
|
|
if content is None:
|
|
# Final fallback: read as utf-8 with replacement for invalid characters
|
|
with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
|
|
content = f.read()
|
|
|
|
filename = os.path.basename(file_path)
|
|
return (content, filename)
|