263 lines
8.0 KiB
Python
263 lines
8.0 KiB
Python
import yaml # type: ignore
|
|
import os
|
|
import folder_paths # type: ignore
|
|
import re
|
|
import time
|
|
import numpy as np # type: ignore
|
|
import torch # type: ignore
|
|
from PIL import Image # type: ignore
|
|
import json
|
|
import base64
|
|
import io
|
|
|
|
|
|
class FileHandler:
|
|
def __init__(self, filename):
|
|
self.filename = filename
|
|
self.lines = None
|
|
|
|
def read_line(self, line_index):
|
|
"""Read a specific line from the file by its index (0-based)."""
|
|
with open(self.filename, "r", encoding="utf-8") as file:
|
|
lines = file.readlines()
|
|
lines = [line for line in lines if line.strip()] # Ignore empty lines
|
|
if self.lines == None:
|
|
self.lines = lines
|
|
if 0 <= line_index < len(lines):
|
|
return lines[line_index].strip()
|
|
else:
|
|
raise IndexError("Line index out of range.")
|
|
|
|
def len_lines(self):
|
|
with open(self.filename, "r", encoding="utf-8") as file:
|
|
lines = file.readlines()
|
|
lines = [line for line in lines if line.strip()] # Ignore empty lines
|
|
return len(lines)
|
|
|
|
def append_line(self, line):
|
|
"""Append a line to the end of the file."""
|
|
with open(self.filename, "a") as file:
|
|
file.write(line + "\n")
|
|
|
|
def load_lines(self):
|
|
"""Load all lines from the file into a list."""
|
|
with open(self.filename, "r", encoding="utf-8") as file:
|
|
return [line.strip() for line in file.readlines()]
|
|
|
|
def escape_quotes(self, text):
|
|
return text.replace('"', '\\"').replace("'", "\\'")
|
|
|
|
def unescape_quotes(self, text):
|
|
return text.replace('\\"', '"').replace("\\'", "'")
|
|
|
|
|
|
def get_safe_path(base_dir, filename):
|
|
"""Sanitize and resolve filename, preventing path traversal attacks."""
|
|
safe_name = os.path.basename(filename)
|
|
if not safe_name or safe_name in (".", ".."):
|
|
raise ValueError("Invalid filename")
|
|
base_resolved = os.path.realpath(base_dir)
|
|
resolved = os.path.realpath(os.path.join(base_resolved, safe_name))
|
|
if os.path.commonpath([base_resolved, resolved]) != base_resolved:
|
|
raise ValueError("Path traversal detected")
|
|
return resolved
|
|
|
|
|
|
def base64_to_pil(base64_string):
|
|
header, encoded = base64_string.split(",", 1) # Remove the data URL header
|
|
image_data = base64.b64decode(encoded)
|
|
return Image.open(io.BytesIO(image_data))
|
|
|
|
|
|
class MyCustomError(Exception):
|
|
def __init__(self, message="Something went wrong"):
|
|
super().__init__(message)
|
|
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
|
|
any_type = AnyType("*")
|
|
|
|
|
|
class FlexibleOptionalInputType(dict):
|
|
def __init__(self, type):
|
|
self.type = type
|
|
|
|
def __getitem__(self, key):
|
|
return (self.type,)
|
|
|
|
def __contains__(self, key):
|
|
return True
|
|
|
|
|
|
def time_it(func, *args, **kwargs):
|
|
start_time = time.time() # Record start time
|
|
result = func(*args, **kwargs) # Call the passed function with arguments
|
|
end_time = time.time() # Record end time
|
|
print(f"Execution time: {end_time - start_time:.6f} seconds")
|
|
return result
|
|
|
|
|
|
def clean_text(text):
|
|
# Remove duplicate commas
|
|
text = re.sub(r",+", ",", text)
|
|
|
|
# Replace any occurrence of " ," with ","
|
|
text = re.sub(r"\s+,", ",", text)
|
|
|
|
# Ensure there's a space after commas followed by a word without space
|
|
text = re.sub(r",(\S)", r", \1", text)
|
|
|
|
# Replace multiple spaces with a single space
|
|
text = re.sub(r"\s+", " ", text)
|
|
|
|
# Replace any occurrence of ".,", ",." with "."
|
|
text = re.sub(r"\.,|,\.", ".", text)
|
|
|
|
# Strip leading and trailing spaces
|
|
return text.strip().replace(" .", ".")
|
|
|
|
|
|
def load_yaml_data(_file_path):
|
|
try:
|
|
# Open the file with UTF-8 encoding
|
|
with open(_file_path, "r", encoding="utf-8") as yaml_file:
|
|
# Load YAML content as Python objects
|
|
_yaml_data = yaml.safe_load(yaml_file)
|
|
|
|
# Ensure the data is returned as a list
|
|
if isinstance(_yaml_data, list):
|
|
return _yaml_data
|
|
else:
|
|
raise ValueError("YAML content is not a list of objects.")
|
|
|
|
except FileNotFoundError:
|
|
print(f"Error: The file '{_file_path}' was not found.")
|
|
except yaml.YAMLError as e:
|
|
print(f"Error: Invalid YAML format. {e}")
|
|
except Exception as e:
|
|
print(f"An error occurred: {e}")
|
|
|
|
return None
|
|
|
|
|
|
def get_yaml_names(_folder_path):
|
|
names = []
|
|
|
|
for file_name in os.listdir(_folder_path):
|
|
if file_name.endswith(".yaml") or file_name.endswith(".yml"):
|
|
names.append(file_name)
|
|
|
|
return names
|
|
|
|
|
|
possible_names = ["comfyui-itools", "ComfyUI-iTools"]
|
|
|
|
|
|
def check_detect_project_dir():
|
|
paths = folder_paths.folder_names_and_paths["custom_nodes"][0]
|
|
for path in paths:
|
|
for name in possible_names:
|
|
proj_dir = os.path.join(path, name)
|
|
if os.path.exists(proj_dir):
|
|
return proj_dir
|
|
raise FileNotFoundError("No valid project directory found on the device.")
|
|
|
|
|
|
project_dir = check_detect_project_dir()
|
|
|
|
|
|
def get_user_extra_style_choice():
|
|
try:
|
|
ud_dir = os.path.join(folder_paths.base_path, "user", "default")
|
|
settings_file = os.path.join(ud_dir, "comfy.settings.json")
|
|
with open(settings_file, "r") as file:
|
|
settings = json.load(file)
|
|
return settings.get("iTools.Nodes.More Styles", False)
|
|
except (OSError, json.JSONDecodeError, AttributeError):
|
|
return False
|
|
|
|
|
|
allow_extra_styles = get_user_extra_style_choice()
|
|
|
|
styles = get_yaml_names(os.path.join(project_dir, "styles"))
|
|
|
|
if allow_extra_styles:
|
|
more_styles = get_yaml_names(os.path.join(project_dir, "styles", "more examples"))
|
|
styles.extend(x for x in more_styles if x not in styles)
|
|
|
|
|
|
def read_styles(_yaml_data):
|
|
if not isinstance(_yaml_data, list):
|
|
print("Error: input data must be a list")
|
|
return None
|
|
|
|
names = []
|
|
|
|
for item in _yaml_data:
|
|
if isinstance(item, dict):
|
|
if "name" in item:
|
|
names.append(item["name"])
|
|
|
|
return names
|
|
|
|
|
|
def tensor2pil(image):
|
|
return Image.fromarray(
|
|
np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
)
|
|
|
|
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
|
|
def pil2mask(image):
|
|
# Convert grayscale image to numpy array
|
|
numpy_array = torch.tensor(np.array(image), dtype=torch.float32)
|
|
# Normalize to binary mask: 0 for transparent (or dark), 1 for opaque (or bright)
|
|
mask = (numpy_array > 0).float() # Converts values to 1 if > 0, otherwise 0
|
|
# Adding a batch dimension
|
|
if mask.dim() == 2:
|
|
mask = mask.unsqueeze(0) # Shape: [1, H, W]
|
|
return mask
|
|
|
|
|
|
# get allow beta nodes
|
|
def get_user_dev_mode():
|
|
try:
|
|
ud_dir = os.path.join(folder_paths.base_path, "user", "default")
|
|
settings_file = os.path.join(ud_dir, "comfy.settings.json")
|
|
with open(settings_file, "r") as file:
|
|
settings = json.load(file)
|
|
return settings.get("iTools.Nodes.Dev Mode", True)
|
|
except (OSError, json.JSONDecodeError, AttributeError):
|
|
return True
|
|
|
|
|
|
# get allow dev nodes
|
|
def get_user_dev_mode2():
|
|
try:
|
|
ud_dir = os.path.join(folder_paths.base_path, "user", "default")
|
|
settings_file = os.path.join(ud_dir, "comfy.settings.json")
|
|
with open(settings_file, "r") as file:
|
|
settings = json.load(file)
|
|
return settings.get("iTools.Nodes.Dev Mode2", False)
|
|
except (OSError, json.JSONDecodeError, AttributeError):
|
|
return False
|
|
|
|
|
|
# get node display name preferences
|
|
def get_user_node_display_name_preferences():
|
|
try:
|
|
ud_dir = os.path.join(folder_paths.base_path, "user", "default")
|
|
settings_file = os.path.join(ud_dir, "comfy.settings.json")
|
|
with open(settings_file, "r") as file:
|
|
settings = json.load(file)
|
|
return settings.get("iTools.Nodes.Node Display Name Preferences", True)
|
|
except (OSError, json.JSONDecodeError, AttributeError):
|
|
return True
|