Steven Mackaay's stuff

StableDiffusion Permutation Script Update

So, like a week ago I wrote a script to make permutations for SD prompts. I’ve updated the script to allow for random terms as well. This allows one to add variance in the prompt but to not add to the number of permutations. Everything is explained in the code block comment. just change the filenames near the end of the script and run.

"""
Script: prompt_permutator.py

Description:
This script generates permutations of a given prompt with various modifiers.
The script reads an input file that contains a base prompt and lists of modifiers.
It creates all possible combinations of the modifiers and generates new prompts
based on these combinations. Additionally, it handles placeholders that are randomly
replaced with specified values and ensures these placeholders are not included in the final output.

Input File Format:
- The first line contains the base prompt.
- Subsequent lines contain comma-separated lists of modifiers.
- Lines starting with a placeholder (e.g., %1) are treated as random modifiers and are replaced
  with random values from the list provided.

Example Input File (test2.txt):

A %1 flower on a hill, photorealistic
on a hill, in a vase, on a bed
photorealistic, manga
%1, Red, Green, Blue

In this example:
- The base prompt is: "A %1 flower on a hill, photorealistic"
- The modifiers are: ["on a hill", "in a vase", "on a bed"] and ["photorealistic", "manga"]
- The placeholder %1 will be replaced with a random choice from ["Red", "Green", "Blue"]

Output:
The script generates all permutations of the prompt with the modifiers and replaces
the placeholder with a random value. The results are saved to an output file.

Usage:
1. Prepare an input file (e.g., 'test2.txt') following the described format.
2. Specify the input and output file paths in the script or pass them as arguments.
3. Run the script to generate the permutations and save them to the output file.

Example Execution:
$ python prompt_permutator.py

Dependencies:
- itertools
- random

Author:
Steven M

Date:
May 28, 2024

"""

import itertools
import random

def load_file(file_path):
    with open(file_path, 'r') as file:
        lines = file.readlines()
    return [line.strip() for line in lines]

def generate_permutations(prompt, modifiers, random_modifiers):
    # Create all combinations of modifiers
    all_combinations = list(itertools.product(*modifiers))
    
    permutations = []
    for combination in all_combinations:
        new_prompt = prompt
        for original, replacement in zip(modifiers, combination):
            new_prompt = replace_first(new_prompt, original[0], replacement)
        
        # Handle random modifiers
        for placeholder, values in random_modifiers.items():
            if placeholder in new_prompt:
                replacement = random.choice(values)
                new_prompt = new_prompt.replace(placeholder, replacement, 1)
        
        # Remove placeholders from the final prompt
        new_prompt = remove_placeholders(new_prompt, random_modifiers.keys())
        
        permutations.append(new_prompt)
    
    return permutations

def replace_first(text, search, replacement):
    # Helper function to replace only the first occurrence of a term
    if search not in text:
        raise ValueError(f"Term '{search}' not found in the prompt.")
    return text.replace(search, replacement, 1)

def remove_placeholders(text, placeholders):
    for placeholder in placeholders:
        text = text.replace(placeholder, "")
    return text

def save_to_file(output_path, permutations):
    with open(output_path, 'w') as file:
        for permutation in permutations:
            file.write(permutation + '\n')

def main(input_file_path, output_file_path):
    lines = load_file(input_file_path)
    if not lines:
        print("The input file is empty.")
        return

    prompt = lines[0]
    modifiers = [line.split(', ') for line in lines[1:] if not line.startswith('%')]
    random_modifiers = {}
    
    for line in lines[1:]:
        if line.startswith('%'):
            parts = line.split(', ')
            key = parts[0]
            values = parts[1:]
            random_modifiers[key] = values

    try:
        all_permutations = generate_permutations(prompt, modifiers, random_modifiers)
        save_to_file(output_file_path, all_permutations)
        print(f"Generated prompts have been saved to {output_file_path}")
        print(f"Total number of permutations: {len(all_permutations)}")
    except ValueError as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    input_file_path = 'test2.txt'  # Change this to the path of your input file
    output_file_path = 'output.txt'  # Change this to the desired path for the output file
    main(input_file_path, output_file_path)

So, in essence, anything with a % at the beginning of the line will be processed differently and the term will be matched up. As always, I guarantee nothing.