import datetime import json import os import random from typing import Dict, List, Tuple, Optional import torch from comfy.cli_args import args import pandas as pd import folder_paths # Import folder_paths from datasets import load_dataset from datasets.arrow_dataset import Dataset from .utils import log # Import log import logging from server import PromptServer logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') log = logging.getLogger(__name__) log.setLevel(logging.DEBUG) # --- Configuration (Adjust these in your workflow or node inputs) --- OUTPUT_FOLDER = folder_paths.get_output_directory() # Get ComfyUI output directory METADATA_SUBFOLDER = "dataset_batch_metadata" # Subfolder for metadata within the output directory # --- End Configuration --- class DatasetBatchNode: """ Custom ComfyUI node to handle batch processing of a dataset. """ def __init__(self): self.output_dir = os.path.join(OUTPUT_FOLDER, METADATA_SUBFOLDER) # Metadata subfolder path self.type = "output" self.metadata = [] self.current_row_index = 0 self.rows_to_process = [] self.prompts = [] self.total_rows = 0 # Track the total number of rows in the dataset self.dataset_length = 0 self.dataset = None # Store the loaded dataset self.magic_number = 0 # Initialise magic number @classmethod def INPUT_TYPES(cls): return { "required": { "dataset_path": ("STRING", {"default": ""}), "prompt_field": ("STRING", {"default": "text"}), "num_rows": ("INT", {"default": -1, "min": -1, "max": 10000, "step": 1}), # -1 for all rows "start_row": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1}), "random_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "shuffle": ("BOOLEAN", {"default": False}), }, "optional": { "delimiter": ("STRING", {"default": "\\n"}), "magic_number": ("INT", {"default": 0, "min": 0, "max": 9999999, "forceInput": True}), "text_input": ("STRING", {"forceInput": True, "default": ""}), "mixed_fields_config_json": ("STRING", {"multiline": True, "default": '[\n {\n "field": "dense_caption",\n "filter": ""\n },\n {\n "field": "style_caption",\n "filter": ""\n }\n]'}) } } RETURN_TYPES = ("STRING", "INT",) RETURN_NAMES = ("prompt", "seed",) FUNCTION = "process_batch" CATEGORY = "Dataset" OUTPUT_NODE = True def process_batch( self, dataset_path: str, prompt_field: str, num_rows: int, start_row: int, random_seed: int, shuffle: bool, magic_number: int = 0, delimiter: str = "\\n", text_input: str = "", mixed_fields_config_json: str = None, **kwargs, ) -> Tuple[str, int]: from datasets import load_dataset # Load the dataset only once if self.dataset is None: try: self.dataset = self.load_dataset(dataset_path) except Exception as e: log.error(f"Error loading dataset: {e}") return ("", 0,) # Check if the specified prompt field exists if prompt_field not in self.dataset.column_names: raise ValueError(f"Prompt field '{prompt_field}' not found in dataset. Available fields: {self.dataset.column_names}") # Determine the rows to process if num_rows == -1: num_rows = len(self.dataset) self.dataset_length = len(self.dataset) effective_start_row = start_row + self.magic_number end_row = min(effective_start_row + num_rows, self.dataset_length) if shuffle: random.seed(random_seed) shuffled_indices = list(range(len(self.dataset))) random.shuffle(shuffled_indices) self.rows_to_process = shuffled_indices[effective_start_row:end_row] else: self.rows_to_process = list(range(effective_start_row, end_row)) log.debug(f"DatasetBatchNode: process_batch called. Current magic_number: {self.magic_number}, start_row: {start_row}, num_rows: {num_rows}") # Only process if there are rows left to process if self.current_row_index < len(self.rows_to_process): # Process the current row current_row_index = self.rows_to_process[self.current_row_index] row = self.dataset[current_row_index] seed = random.randint(0, 2**32 - 1) combined_prompt_segments = [] if mixed_fields_config_json: # Parse mixed fields config from JSON string try: mixed_fields_config = json.loads(mixed_fields_config_json) if not isinstance(mixed_fields_config, list): raise ValueError("mixed_fields_config must be a JSON list of field configurations.") except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON in mixed_fields_config: {e}") # Process each field config for field_config in mixed_fields_config: field_name = field_config.get("field") filter_query = field_config.get("filter", "") if not field_name: log.warning("Skipping field config due to missing 'field' name.") continue if filter_query.strip(): # Apply filter filtered_dataset = self.dataset.filter(lambda example: eval(filter_query, {}, {"example": example})) if not filtered_dataset: log.warning(f"No rows found matching filter: '{filter_query}' for field '{field_name}'. Using current row instead.") selected_row = row else: selected_row = random.choice(filtered_dataset) else: selected_row = row caption_segment = selected_row.get(field_name, "") combined_prompt_segments.append(caption_segment) # Combine segments with newlines prompt_text = delimiter.join([segment.strip() for segment in combined_prompt_segments if segment]) else: prompt_text = row[prompt_field] # Combine with text_input if provided if text_input: prompt_text = text_input + delimiter + prompt_text # Trim whitespace from the start and end of the prompt prompt_text = prompt_text.strip() # Metadata entry timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") metadata_entry = { "row_index": current_row_index, "seed": seed, "magic_number": self.magic_number, "output_filename": f"row_{current_row_index}_{seed}_{timestamp}", "prompt": prompt_text, "mixed_fields_config": mixed_fields_config_json, **row, } self.metadata.append(metadata_entry) log.info(f"Processing row {current_row_index} with seed {seed}...") # Save metadata to a JSONL file after processing each row os.makedirs(self.output_dir, exist_ok=True) timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") metadata_filename = os.path.join(self.output_dir, f"metadata_{timestamp}.jsonl") with open(metadata_filename, "a", encoding='utf-8') as f: json.dump(metadata_entry, f, ensure_ascii=False) f.write("\n") log.info(f"Metadata for row {current_row_index} saved to {metadata_filename}") # Increment the row index and magic_number self.current_row_index += 1 self.magic_number += 1 # Increment magic_number for the next row # Dispatch custom event after processing each row PromptServer.instance.send_sync("dataset_row_processed", { "node_id": id(self), # Pass the unique ID of the node instance "magic_number": self.magic_number, # Pass the updated magic_number "row_index": current_row_index, }) # Check if all rows for this batch are processed if self.current_row_index >= len(self.rows_to_process): self.current_row_index = 0 # Reset for the next batch self.rows_to_process = [] # Clear the processed rows log.info(f"DatasetBatchNode: Batch complete. magic_number updated to {self.magic_number}") return (prompt_text, seed,) else: # If there are no more rows to process, you might want to handle this case. # For example, you can reset the magic_number or log a message. log.info("Dataset processing complete.") self.magic_number = 0 # Reset magic_number self.current_row_index = 0 # Reset current_row_index self.rows_to_process = [] # Clear the processed rows return ("", 0) # Indicate end of processing def load_dataset(self, dataset_path: str) -> Dataset: """Loads a dataset from a file (CSV, JSONL) or a Hugging Face dataset.""" from datasets import load_dataset, load_from_disk if os.path.isdir(dataset_path): # Load from disk return load_from_disk(dataset_path) elif os.path.isfile(dataset_path): # Load from a local file file_ext = os.path.splitext(dataset_path)[1].lower() if file_ext == ".jsonl": data = [] with open(dataset_path, 'r', encoding='utf-8') as f: for line in f: try: data.append(json.loads(line)) except json.JSONDecodeError as e: log.error(f"Error decoding JSON line: {e}") continue return Dataset.from_list(data) elif file_ext == ".csv": df = pd.read_csv(dataset_path) return Dataset.from_pandas(df) else: raise ValueError(f"Unsupported file format: {file_ext}") else: # Assume it's a Hugging Face dataset path try: # Try to load the dataset with a 'train' split return load_dataset(dataset_path, split="train") except ValueError: # If 'train' split not found, try loading without specifying the split return load_dataset(dataset_path, split=None) WEB_DIRECTORY = "./js" NODE_CLASS_MAPPINGS = { "DatasetBatchNode": DatasetBatchNode } NODE_DISPLAY_NAME_MAPPINGS = { "DatasetBatchNode": "Dataset Batch Node" }