Add support for custom CSV file paths with fallback mechanism

Co-authored-by: theUpsider <25013640+theUpsider@users.noreply.github.com>

README.md

@@ -15,6 +15,15 @@ This extension allows users to load styles from a CSV file (styles.csv), primari
 ## Nodes Description
 Each style is represented as a dictionary with the keys being `style_name` and the values being a list containing `positive_prompt` and `negative_prompt`. The prompts are outputs of this Node.
 
+### Custom CSV Files
+The extension now supports loading styles from custom CSV files in addition to the default `styles.csv`. You can specify a custom CSV file path using the optional `csv_file_path` parameter:
+
+- **Relative paths**: Relative to the ComfyUI root directory (e.g., `"my_styles/portrait_styles.csv"`)
+- **Absolute paths**: Full file system paths (e.g., `"/path/to/my/custom_styles.csv"`)
+- **Default behavior**: When no custom path is specified, it uses `styles.csv` from the ComfyUI root directory
+
+If a selected style is not found in the custom CSV file, the node will fall back to the default styles loaded from `styles.csv`.
+
 ## CSV Format
 The CSV file should have the following format:
 ```csv
@@ -23,6 +32,23 @@ name,prompt,negative_prompt
 "Another Style","more positive text","more negative text"
 ```
 
+### Examples
+You can organize your styles in multiple CSV files for different purposes:
+
+**portraits.csv**:
+```csv
+name,prompt,negative_prompt
+"Professional Portrait","portrait photography, professional lighting, sharp focus","ugly, deformed, extra limbs"
+"Vintage Portrait","vintage portrait, sepia tone, classical","modern, digital, colorful"
+```
+
+**landscapes.csv**:
+```csv
+name,prompt,negative_prompt
+"Golden Hour","landscape photography, golden hour, wide angle","people, buildings, urban"
+"Mountain Vista","mountain landscape, dramatic clouds, epic view","flat, boring, low quality"
+```
+
 ## Development
 
 ### Setting up the development environment

styles_csv_loader.py

@@ -45,13 +45,15 @@ class StylesCSVLoader:
     @classmethod
     def INPUT_TYPES(cls):
         # Use os.path.normpath to ensure cross-platform compatibility
-        styles_path = os.path.normpath(os.path.join(folder_paths.base_path, "styles.csv"))
-        cls.styles_csv = cls.load_styles_csv(styles_path)
+        default_styles_path = os.path.normpath(os.path.join(folder_paths.base_path, "styles.csv"))
+        cls.styles_csv = cls.load_styles_csv(default_styles_path)
         return {
             "required": {
                 "styles": (list(cls.styles_csv.keys()),),
             },
-
+            "optional": {
+                "csv_file_path": ("STRING", {"default": "styles.csv", "multiline": False}),
+            }
         }
 
     RETURN_TYPES = ("STRING", "STRING")
@@ -59,7 +61,25 @@ class StylesCSVLoader:
     FUNCTION = "execute"
     CATEGORY = "loaders"
 
-    def execute(self, styles):
+    def execute(self, styles, csv_file_path="styles.csv"):
+        # Handle custom CSV file path
+        if csv_file_path != "styles.csv":
+            # Check if path is absolute or relative
+            if os.path.isabs(csv_file_path):
+                custom_styles_path = os.path.normpath(csv_file_path)
+            else:
+                # Relative path from ComfyUI root directory
+                custom_styles_path = os.path.normpath(os.path.join(folder_paths.base_path, csv_file_path))
+            
+            # Load styles from custom path
+            custom_styles = self.load_styles_csv(custom_styles_path)
+            if styles in custom_styles:
+                return (custom_styles[styles][0], custom_styles[styles][1])
+            else:
+                # Fall back to default styles if style not found in custom file
+                pass
+        
+        # Use default styles
         return (self.styles_csv[styles][0], self.styles_csv[styles][1])
 
 

tests/test_styles_csv_loader.py

@@ -152,6 +152,82 @@ class TestStylesCSVLoader(unittest.TestCase):
         
         self.assertEqual(positive, 'positive prompt here')
         self.assertEqual(negative, 'negative prompt here')
+    
+    def test_execute_method_with_custom_csv_path(self):
+        """Test the execute method with custom CSV file path."""
+        loader = StylesCSVLoader()
+        
+        # Mock the default styles_csv attribute
+        loader.styles_csv = {
+            'Default Style': ['default positive', 'default negative']
+        }
+        
+        # Test with custom CSV path
+        positive, negative = loader.execute('Cinematic', self.valid_csv_path)
+        
+        self.assertEqual(positive, 'cinematic shot, dramatic lighting, film grain')
+        self.assertEqual(negative, 'low quality, blurry, amateur')
+    
+    def test_execute_method_with_custom_csv_fallback(self):
+        """Test that execute method falls back to default when style not found in custom CSV."""
+        loader = StylesCSVLoader()
+        
+        # Mock the default styles_csv attribute
+        loader.styles_csv = {
+            'Default Style': ['default positive', 'default negative']
+        }
+        
+        # Test with custom CSV path but style that doesn't exist in custom file
+        positive, negative = loader.execute('Default Style', self.valid_csv_path)
+        
+        self.assertEqual(positive, 'default positive')
+        self.assertEqual(negative, 'default negative')
+    
+    def test_execute_method_with_absolute_path(self):
+        """Test execute method with absolute path to custom CSV."""
+        loader = StylesCSVLoader()
+        
+        # Mock the default styles_csv attribute
+        loader.styles_csv = {
+            'Default Style': ['default positive', 'default negative']
+        }
+        
+        # Test with absolute path
+        absolute_path = os.path.abspath(self.valid_csv_path)
+        positive, negative = loader.execute('Cinematic', absolute_path)
+        
+        self.assertEqual(positive, 'cinematic shot, dramatic lighting, film grain')
+        self.assertEqual(negative, 'low quality, blurry, amateur')
+    
+    @patch('styles_csv_loader.folder_paths')
+    def test_input_types_includes_csv_file_path(self, mock_folder_paths):
+        """Test that INPUT_TYPES includes the csv_file_path optional parameter."""
+        mock_folder_paths.base_path = self.fixtures_dir
+        
+        # Create a temporary styles.csv in the fixtures directory
+        temp_csv_path = os.path.join(self.fixtures_dir, 'styles.csv')
+        with open(temp_csv_path, 'w', encoding='utf-8') as f:
+            f.write('name,prompt,negative_prompt\n')
+            f.write('"Test Style","test prompt","test negative"\n')
+        
+        try:
+            input_types = StylesCSVLoader.INPUT_TYPES()
+            
+            # Check structure includes optional csv_file_path
+            self.assertIn('required', input_types)
+            self.assertIn('optional', input_types)
+            self.assertIn('csv_file_path', input_types['optional'])
+            
+            # Check csv_file_path is a STRING type with default
+            csv_param = input_types['optional']['csv_file_path']
+            self.assertEqual(csv_param[0], 'STRING')
+            self.assertIn('default', csv_param[1])
+            self.assertEqual(csv_param[1]['default'], 'styles.csv')
+        
+        finally:
+            # Clean up
+            if os.path.exists(temp_csv_path):
+                os.remove(temp_csv_path)
 
 
 if __name__ == '__main__':