David Vincent Fischer
695d2616e7
fixed an issue where the loader would not recognize csv files with trailing comma or omitted columns.
103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
import os
|
|
import re
|
|
import folder_paths
|
|
|
|
class StylesCSVLoader:
|
|
"""
|
|
Loads csv file with styles. For migration purposes from automatic1111 webui.
|
|
"""
|
|
|
|
@staticmethod
|
|
def load_styles_csv(styles_path: str):
|
|
"""Loads csv file with styles. Each row has three columns:
|
|
- style_name: string (labeled as "name" in the CSV)
|
|
- positive_prompt: string (labeled as "prompt" in the CSV)
|
|
- negative_prompt: string (may be empty or missing)
|
|
Returns:
|
|
dict: Dictionary of styles, where each key is style_name and value is [positive_prompt, negative_prompt]
|
|
"""
|
|
error_styles = {
|
|
"Error loading styles.csv, check the console": ["", ""]}
|
|
|
|
if not os.path.exists(styles_path):
|
|
print(f"""Error. No styles.csv found. Put your styles.csv in the root directory of ComfyUI. Then press "Refresh".
|
|
Your current root directory is: {folder_paths.base_path}
|
|
""")
|
|
return error_styles
|
|
|
|
try:
|
|
styles = {}
|
|
with open(styles_path, "r", encoding="utf-8") as f:
|
|
reader = csv.reader(f)
|
|
next(reader) # Skip header row
|
|
|
|
# Start counting from 2 (after header)
|
|
for i, row in enumerate(reader, start=2):
|
|
try:
|
|
# Initialize variables to avoid "used before assigned" errors
|
|
style_name = ""
|
|
positive_prompt = ""
|
|
negative_prompt = ""
|
|
|
|
# Handle rows with fewer than 3 columns
|
|
if len(row) == 0:
|
|
# Skip empty rows
|
|
continue
|
|
elif len(row) == 1:
|
|
print(
|
|
f"Error parsing line {i}: {','.join(row)} -> Expected at least 2 columns, got {len(row)}")
|
|
continue
|
|
elif len(row) == 2:
|
|
# Missing negative prompt - treat as empty string
|
|
style_name, positive_prompt = row
|
|
# negative_prompt already initialized as empty string
|
|
elif len(row) >= 3:
|
|
# Normal case or extra columns (ignore extras)
|
|
style_name, positive_prompt, negative_prompt = row[0], row[1], row[2]
|
|
|
|
styles[style_name] = [positive_prompt, negative_prompt]
|
|
except Exception as e:
|
|
# Use a safe version of the row for error reporting
|
|
row_str = ','.join(row) if isinstance(
|
|
row, list) else str(row)
|
|
print(f"Error parsing line {i}: {row_str} -> {e}")
|
|
|
|
return styles
|
|
|
|
except Exception as e:
|
|
print(f"""Error loading styles.csv. Make sure it is in the root directory of ComfyUI. Then press "Refresh".
|
|
Your current root directory is: {folder_paths.base_path}
|
|
Error: {e}
|
|
""")
|
|
return error_styles
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
cls.styles_csv = cls.load_styles_csv(
|
|
os.path.join(folder_paths.base_path, "styles.csv"))
|
|
return {
|
|
"required": {
|
|
"styles": (list(cls.styles_csv.keys()),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("positive prompt", "negative prompt")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "loaders"
|
|
|
|
def execute(self, styles):
|
|
return (
|
|
self.styles_csv[styles][0],
|
|
self.styles_csv[styles][1]
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Load Styles CSV": StylesCSVLoader
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"StylesCSVLoader": "Load Styles CSV Node"
|
|
}
|