diff --git a/styles_csv_loader.py b/styles_csv_loader.py index 8927e31..d7bf2cf 100644 --- a/styles_csv_loader.py +++ b/styles_csv_loader.py @@ -1,14 +1,13 @@ import os import re +import folder_paths class StylesCSVLoader: """ Loads csv file with styles. For migration purposes from automatic11111 webui. """ - def __init__(self): - if not hasattr(self, "styles_csv"): - self.styles_csv = self.load_styles_csv("styles.csv") + @staticmethod def load_styles_csv(styles_path): """Loads csv file with styles. It has only one column. Ignore the first row (header). @@ -18,16 +17,26 @@ class StylesCSVLoader: Returns: list: List of styles. Each style is a dict with keys: style_name and value: [positive_prompt, negative_prompt] """ + styles = {"Error loading styles.csv, check the console": ["",""]} if not os.path.exists(styles_path): - return {"Paste styles.csv in ComfyUI root and press \"Refresh\"...": ["",""]} # return empty style if file does not exist - with open(styles_path, "r", encoding="utf-8") as f: - styles = [[x.replace('"', '').replace('\n','') for x in re.split(',(?=(?:[^"]*"[^"]*")*[^"]*$)', line)] for line in f.readlines()[1:]] - styles = {x[0]: [x[1],x[2]] for x in styles} + 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 styles + try: + with open(styles_path, "r", encoding="utf-8") as f: + styles = [[x.replace('"', '').replace('\n','') for x in re.split(',(?=(?:[^"]*"[^"]*")*[^"]*$)', line)] for line in f.readlines()[1:]] + styles = {x[0]: [x[1],x[2]] for x in 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 styles @classmethod def INPUT_TYPES(s): - s.styles_csv = s.load_styles_csv("styles.csv") + s.styles_csv = s.load_styles_csv(os.path.join(folder_paths.base_path, "styles.csv")) return { "required": { "styles": (list(s.styles_csv.keys()),),