mingchoi-ComfyUI-Prompt-Iterator/nodes.py

"""
ComfyUI Prompt Iterator Custom Nodes

This module provides two custom nodes for iterating through prompts and templates:
1. PromptIteratorNode - Iterates through multi-line prompts
2. PromptTemplateIteratorNode - Iterates through keywords and fills templates
"""

# Global counter shared between all iterator nodes
_iteration_counter = 0
# Track last content to detect changes
_last_prompt_content = ""
_last_keyword_content = ""
_last_template_content = ""


def get_counter():
    """Get the current iteration counter value."""
    global _iteration_counter
    return _iteration_counter


def increment_counter():
    """Increment the global iteration counter."""
    global _iteration_counter
    _iteration_counter += 1
    return _iteration_counter


def reset_counter():
    """Reset the global iteration counter to 0."""
    global _iteration_counter
    _iteration_counter = 0
    return _iteration_counter


def check_and_reset_if_content_changed(content, content_type="prompt"):
    """Check if content changed and reset counter if needed."""
    global _iteration_counter, _last_prompt_content, _last_keyword_content, _last_template_content
    
    if content_type == "prompt":
        if content != _last_prompt_content:
            _last_prompt_content = content
            _iteration_counter = 0
    elif content_type == "keyword":
        if content != _last_keyword_content:
            _last_keyword_content = content
            _iteration_counter = 0
    elif content_type == "template":
        if content != _last_template_content:
            _last_template_content = content
            _iteration_counter = 0
    
    return _iteration_counter


class PromptIteratorNode:
    """
    A node that iterates through multi-line prompts, outputting one line per execution.
    Uses a global counter to track the current position.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompts": ("STRING", {
                    "multiline": True,
                    "default": "prompt 1\nprompt 2\nprompt 3"
                }),
            },
        }
    
    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("prompt", "index")
    FUNCTION = "iterate"
    CATEGORY = "utils"
    OUTPUT_NODE = False
    
    @classmethod
    def IS_CHANGED(cls, prompts):
        """Force execution on every run by returning the current counter and content."""
        return (get_counter(), prompts)
    
    def iterate(self, prompts):
        """
        Iterate through prompts and return the current one based on global counter.
        Counter resets automatically when prompt content changes.
        
        Args:
            prompts: Multi-line string containing prompts (one per line)
            
        Returns:
            tuple: (current_prompt, index_info)
        """
        # Reset counter if prompt content changed
        check_and_reset_if_content_changed(prompts, "prompt")
        
        # Split prompts by newlines and filter out empty lines
        prompt_lines = [line.strip() for line in prompts.split('\n') if line.strip()]
        
        if not prompt_lines:
            return ("", "0/0")
        
        # Get current counter value
        counter = get_counter()
        
        # Calculate current index with wraparound
        current_index = counter % len(prompt_lines)
        current_prompt = prompt_lines[current_index]
        
        # Increment counter for next execution
        increment_counter()
        
        # Return the prompt and index info
        index_info = f"{current_index+1}/{len(prompt_lines)}"
        return (current_prompt, index_info)


class PromptTemplateIteratorNode:
    """
    A node that iterates through keywords and fills a template string.
    Uses the same global counter as PromptIteratorNode.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "keywords": ("STRING", {
                    "multiline": True,
                    "default": "cat\ndog\nbird"
                }),
                "template": ("STRING", {
                    "multiline": False,
                    "default": "A photo of a {keyword}"
                }),
            },
        }
    
    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("prompt", "index")
    FUNCTION = "iterate"
    CATEGORY = "utils"
    OUTPUT_NODE = False
    
    @classmethod
    def IS_CHANGED(cls, keywords, template):
        """Force execution on every run by returning the current counter and content."""
        return (get_counter(), keywords, template)
    
    def iterate(self, keywords, template):
        """
        Iterate through keywords and fill the template with the current keyword.
        Counter resets automatically when keyword or template content changes.
        
        Args:
            keywords: Multi-line string containing keywords (one per line)
            template: Template string with {keyword} placeholder
            
        Returns:
            tuple: (filled_template, index_info)
        """
        # Reset counter if keyword or template content changed
        check_and_reset_if_content_changed(keywords, "keyword")
        check_and_reset_if_content_changed(template, "template")
        
        # Split keywords by newlines and filter out empty lines
        keyword_lines = [line.strip() for line in keywords.split('\n') if line.strip()]
        
        if not keyword_lines:
            return ("", "0/0")
        
        # Get current counter value
        counter = get_counter()
        
        # Calculate current index with wraparound
        current_index = counter % len(keyword_lines)
        current_keyword = keyword_lines[current_index]
        
        # Fill template with current keyword
        output = template.replace("{keyword}", current_keyword)
        
        # Increment counter for next execution
        increment_counter()
        
        # Return the filled template and index info
        index_info = f"{current_index+1}/{len(keyword_lines)}"
        return (output, index_info)


# Node class mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
    "PromptIterator": PromptIteratorNode,
    "PromptTemplateIterator": PromptTemplateIteratorNode,
}

# Display name mappings for ComfyUI UI
NODE_DISPLAY_NAME_MAPPINGS = {
    "PromptIterator": "Prompt Iterator",
    "PromptTemplateIterator": "Prompt Template Iterator",
}