diff --git a/styles_csv_loader.py b/styles_csv_loader.py index a8283d0..88fa1be 100644 --- a/styles_csv_loader.py +++ b/styles_csv_loader.py @@ -4,36 +4,54 @@ import folder_paths class StylesCSVLoader: """ - Loads csv file with styles. For migration purposes from automatic11111 webui. + Loads csv file with styles. For migration purposes from automatic1111 webui. """ - + @staticmethod def load_styles_csv(styles_path: str): - """Loads csv file with styles. It has only one column. - Ignore the first row (header). - positive_prompt are strings separated by comma. Each string is a prompt. - negative_prompt are strings separated by comma. Each string is a prompt. + """Loads csv file with styles. Each row has three columns: + - style_name: string + - positive_prompt: string (comma-separated prompts, optionally in quotes) + - negative_prompt: string (comma-separated prompts, optionally in quotes) Returns: - list: List of styles. Each style is a dict with keys: style_name and value: [positive_prompt, negative_prompt] + dict: Dictionary of styles, where each key is style_name and value is [positive_prompt, negative_prompt] """ - styles = {"Error loading styles.csv, check the console": ["",""]} + 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} - """) +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} + with open(styles_path, "r", encoding="utf-8") as f: + lines = f.readlines() + + styles = {} + for i, line in enumerate(lines[1:], start=2): # Skip header row + try: + # Split using regex that respects quoted fields + parts = [x.replace('"', '').replace('\n', '') for x in re.split(',(?=(?:[^"]*"[^"]*")*[^"]*$)', line)] + + if len(parts) != 3: + raise ValueError(f"Expected 3 columns, got {len(parts)}") + + style_name, positive_prompt, negative_prompt = parts + styles[style_name] = [positive_prompt, negative_prompt] + + except Exception as e: + print(f"Error parsing line {i}: {line.strip()} -> {e}") + 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} - """) +Your current root directory is: {folder_paths.base_path} +Error: {e} +""") + return styles - + @classmethod def INPUT_TYPES(cls): cls.styles_csv = cls.load_styles_csv(os.path.join(folder_paths.base_path, "styles.csv")) @@ -41,20 +59,23 @@ class StylesCSVLoader: "required": { "styles": (list(cls.styles_csv.keys()),), }, - } - - RETURN_TYPES = ("STRING","STRING") + + RETURN_TYPES = ("STRING", "STRING") RETURN_NAMES = ("positive prompt", "negative prompt") FUNCTION = "execute" - CATEGORY = "loaders" + CATEGORY = "loaders" def execute(self, styles): - return (self.styles_csv[styles][0], self.styles_csv[styles][1]) + 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" }