Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Dynamic prompt generation script for parameter scans (#2831)

Browse Source 
# Programatically generate a large number of images varying by prompt
and other image generation parameters

This is a little standalone script named `dynamic_prompting.py` that
enables the generation of dynamic prompts. Using YAML syntax, you
specify a template of prompt phrases and lists of generation parameters,
and the script will generate a cross product of prompts and generation
settings for you. You can save these prompts to disk for later use, or
pipe them to the invokeai CLI to generate the images on the fly.

Typical uses are testing step and CFG values systematically while
holding the seed and prompt constant, testing out various artist's
styles, and comparing the results of the same prompt across different
models.

A typical template will look like this:

```
model: stable-diffusion-1.5
steps: 30;50;10
seed: 50
dimensions: 512x512
cfg:
  - 7
  - 12
sampler:
  - k_euler_a
  - k_lms
prompt:
  style:
       - greg rutkowski
       - gustav klimt
  location:
       - the mountains
       - a desert
  object:
       - luxurious dwelling
       - crude tent
  template: a {object} in {location}, in the style of {style}
```

This will generate 96 different images, each of which varies by one of
the dimensions specified in the template. For example, the prompt axis
will generate a cross product list like:
```
a luxurious dwelling in the mountains, in the style of greg rutkowski
a luxurious dwelling in the mountains, in the style of gustav klimt
a luxious dwelling in a desert, in the style of greg rutkowski
... etc
```

A typical usage would be:
```
python scripts/dynamic_prompts.py --invoke --outdir=/tmp/scanning my_template.yaml
```
This will populate `/tmp/scanning` with each of the requested images,
and also generate a `log.md` file which you can open with an e-book
reader to show something like this:


![image](https://user-images.githubusercontent.com/111189/221970165-4bbd9070-3f32-4d89-8ff2-b03a82ada575.png)

Full instructions can be obtained using the `--instructions` switch, and
an example template can be printed out using `--example`:

```
python scripts/dynamic_prompts.py --instructions
python scripts/dynamic_prompts.py --example > my_first_template.yaml
```
This commit is contained in:
Lincoln Stein 2023-03-09 20:18:28 -05:00 committed by GitHub
commit 8323169864
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 588 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
invokeai/frontend/stats.html
View File

File diff suppressed because one or more lines are too long

2
ldm/generate.py
View File

@ -200,6 +200,8 @@ class Generate:
# it wasn't actually doing anything. This logic could be reinstated.
self.device = torch.device(choose_torch_device())
print(f">> Using device_type {self.device.type}")
if self.device.type == 'cuda':
print(f">> CUDA device '{torch.cuda.get_device_name(torch.cuda.current_device())}' (GPU {os.environ.get('CUDA_VISIBLE_DEVICES') or 0})")
if full_precision:
if self.precision != "auto":
raise ValueError("Remove --full_precision / -F if using --precision")

28
ldm/invoke/CLI.py
View File

@ -389,6 +389,7 @@ def main_loop(gen, opt):
prior_variations,
postprocessed,
first_seed,
gen.model_name,
)
path = file_writer.save_image_and_prompt_to_png(
image=image,
@ -402,6 +403,7 @@ def main_loop(gen, opt):
else first_seed
],
model_hash=gen.model_hash,
model_id=gen.model_name,
),
name=filename,
compress_level=opt.png_compression,
@ -941,13 +943,14 @@ def add_postprocessing_to_metadata(opt, original_file, new_file, tool, command):
def prepare_image_metadata(
opt,
prefix,
seed,
operation="generate",
prior_variations=[],
postprocessed=False,
first_seed=None,
opt,
prefix,
seed,
operation="generate",
prior_variations=[],
postprocessed=False,
first_seed=None,
model_id='unknown',
):
if postprocessed and opt.save_original:
filename = choose_postprocess_name(opt, prefix, seed)
@ -955,7 +958,9 @@ def prepare_image_metadata(
wildcards = dict(opt.__dict__)
wildcards["prefix"] = prefix
wildcards["seed"] = seed
wildcards["model_id"] = model_id
try:
print(f'DEBUG: fnformat={opt.fnformat}')
filename = opt.fnformat.format(**wildcards)
except KeyError as e:
print(
@ -972,18 +977,17 @@ def prepare_image_metadata(
first_seed = first_seed or seed
this_variation = [[seed, opt.variation_amount]]
opt.with_variations = prior_variations + this_variation
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed)
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed,model_id=model_id)
elif len(prior_variations) > 0:
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed)
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed,model_id=model_id)
elif operation == "postprocess":
formatted_dream_prompt = "!fix " + opt.dream_prompt_str(
seed=seed, prompt=opt.input_file_path
seed=seed, prompt=opt.input_file_path, model_id=model_id,
)
else:
formatted_dream_prompt = opt.dream_prompt_str(seed=seed)
formatted_dream_prompt = opt.dream_prompt_str(seed=seed,model_id=model_id)
return filename, formatted_dream_prompt
def choose_postprocess_name(opt, prefix, seed) -> str:
match = re.search("postprocess:(\w+)", opt.last_operation)
if match:

7
ldm/invoke/args.py
View File

@ -333,7 +333,7 @@ class Args(object):
switches.append(f'-V {formatted_variations}')
if 'variations' in a and len(a['variations'])>0:
switches.append(f'-V {a["variations"]}')
return ' '.join(switches)
return ' '.join(switches) + f' # model_id={kwargs.get("model_id","unknown model")}'
def __getattribute__(self,name):
'''
@ -878,7 +878,7 @@ class Args(object):
)
render_group.add_argument(
'--fnformat',
default='{prefix}.{seed}.png',
default=None,
type=str,
help='Overwrite the filename format. You can use any argument as wildcard enclosed in curly braces. Default is {prefix}.{seed}.png',
)
@ -1155,6 +1155,7 @@ def format_metadata(**kwargs):
def metadata_dumps(opt,
seeds=[],
model_hash=None,
model_id=None,
postprocessing=None):
'''
Given an Args object, returns a dict containing the keys and
@ -1167,7 +1168,7 @@ def metadata_dumps(opt,
# top-level metadata minus `image` or `images`
metadata = {
'model' : 'stable diffusion',
'model_id' : opt.model,
'model_id' : model_id or opt.model,
'model_hash' : model_hash,
'app_id' : ldm.invoke.__app_id__,
'app_version' : ldm.invoke.__version__,

535
ldm/invoke/dynamic_prompts.py Executable file
View File

@ -0,0 +1,535 @@
#!/usr/bin/env python
"""
Simple script to generate a file of InvokeAI prompts and settings
that scan across steps and other parameters.
"""
import argparse
import io
import json
import os
import pydoc
import re
import shutil
import sys
import time
from contextlib import redirect_stderr
from io import TextIOBase
from itertools import product
from multiprocessing import Process
from multiprocessing.connection import Connection, Pipe
from pathlib import Path
from tempfile import gettempdir
from typing import Callable, Iterable, List
import numpy as np
import yaml
from omegaconf import OmegaConf, dictconfig, listconfig
def expand_prompts(
template_file: Path,
run_invoke: bool = False,
invoke_model: str = None,
invoke_outdir: Path = None,
processes_per_gpu: int = 1,
):
"""
:param template_file: A YAML file containing templated prompts and args
:param run_invoke: A boolean which if True will pass expanded prompts to invokeai CLI
:param invoke_model: Name of the model to load when run_invoke is true; otherwise uses default
:param invoke_outdir: Directory for outputs when run_invoke is true; otherwise uses default
"""
if template_file.name.endswith(".json"):
with open(template_file, "r") as file:
with io.StringIO(yaml.dump(json.load(file))) as fh:
conf = OmegaConf.load(fh)
else:
conf = OmegaConf.load(template_file)
# loading here to avoid long wait for help message
import torch
torch.multiprocessing.set_start_method("spawn")
gpu_count = torch.cuda.device_count() if torch.cuda.is_available() else 1
commands = expanded_invokeai_commands(conf, run_invoke)
children = list()
try:
if run_invoke:
invokeai_args = [shutil.which("invokeai"), "--from_file", "-"]
if invoke_model:
invokeai_args.extend(("--model", invoke_model))
if invoke_outdir:
outdir = os.path.expanduser(invoke_outdir)
invokeai_args.extend(("--outdir", outdir))
else:
outdir = gettempdir()
logdir = Path(outdir, "invokeai-batch-logs")
processes_to_launch = gpu_count * processes_per_gpu
print(
f">> Spawning {processes_to_launch} invokeai processes across {gpu_count} CUDA gpus",
file=sys.stderr,
)
print(
f'>> Outputs will be written into {invoke_outdir or "default InvokeAI outputs directory"}, and error logs will be written to {logdir}',
file=sys.stderr,
)
import ldm.invoke.CLI
parent_conn, child_conn = Pipe()
children = set()
for i in range(processes_to_launch):
p = Process(
target=_run_invoke,
kwargs=dict(
entry_point=ldm.invoke.CLI.main,
conn_in=child_conn,
conn_out=parent_conn,
args=invokeai_args,
gpu=i % gpu_count,
logdir=logdir,
),
)
p.start()
children.add(p)
child_conn.close()
sequence = 0
for command in commands:
sequence += 1
parent_conn.send(
command + f' --fnformat="dp.{sequence:04}.{{prompt}}.png"'
)
parent_conn.close()
else:
for command in commands:
print(command)
except KeyboardInterrupt:
for p in children:
p.terminate()
class MessageToStdin(object):
def __init__(self, connection: Connection):
self.connection = connection
self.linebuffer = list()
def readline(self) -> str:
try:
if len(self.linebuffer) == 0:
message = self.connection.recv()
self.linebuffer = message.split("\n")
result = self.linebuffer.pop(0)
return result
except EOFError:
return None
class FilterStream(object):
def __init__(
self, stream: TextIOBase, include: re.Pattern = None, exclude: re.Pattern = None
):
self.stream = stream
self.include = include
self.exclude = exclude
def write(self, data: str):
if self.include and self.include.match(data):
self.stream.write(data)
self.stream.flush()
elif self.exclude and not self.exclude.match(data):
self.stream.write(data)
self.stream.flush()
def flush(self):
self.stream.flush()
def _run_invoke(
entry_point: Callable,
conn_in: Connection,
conn_out: Connection,
args: List[str],
logdir: Path,
gpu: int = 0,
):
pid = os.getpid()
logdir.mkdir(parents=True, exist_ok=True)
logfile = Path(logdir, f'{time.strftime("%Y-%m-%d-%H:%M:%S")}-pid={pid}.txt')
print(
f">> Process {pid} running on GPU {gpu}; logging to {logfile}", file=sys.stderr
)
conn_out.close()
os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu}"
sys.argv = args
sys.stdin = MessageToStdin(conn_in)
sys.stdout = FilterStream(sys.stdout, include=re.compile("^\[\d+\]"))
with open(logfile, "w") as stderr, redirect_stderr(stderr):
entry_point()
def _filter_output(stream: TextIOBase):
while line := stream.readline():
if re.match("^\[\d+\]", line):
print(line)
def main():
parser = argparse.ArgumentParser(
description=HELP,
)
parser.add_argument(
"template_file",
type=Path,
nargs="?",
help="path to a template file, use --example to generate an example file",
)
parser.add_argument(
"--example",
action="store_true",
default=False,
help=f'Print an example template file in YAML format. Use "{sys.argv[0]} --example > example.yaml" to save output to a file',
)
parser.add_argument(
"--json-example",
action="store_true",
default=False,
help=f'Print an example template file in json format. Use "{sys.argv[0]} --json-example > example.json" to save output to a file',
)
parser.add_argument(
"--instructions",
"-i",
dest="instructions",
action="store_true",
default=False,
help="Print verbose instructions.",
)
parser.add_argument(
"--invoke",
action="store_true",
help="Execute invokeai using specified optional --model, --processes_per_gpu and --outdir",
)
parser.add_argument(
"--model",
help="Feed the generated prompts to the invokeai CLI using the indicated model. Will be overriden by a model: section in template file.",
)
parser.add_argument(
"--outdir", type=Path, help="Write images and log into indicated directory"
)
parser.add_argument(
"--processes_per_gpu",
type=int,
default=1,
help="When executing invokeai, how many parallel processes to execute per CUDA GPU.",
)
opt = parser.parse_args()
if opt.example:
print(EXAMPLE_TEMPLATE_FILE)
sys.exit(0)
if opt.json_example:
print(_yaml_to_json(EXAMPLE_TEMPLATE_FILE))
sys.exit(0)
if opt.instructions:
pydoc.pager(INSTRUCTIONS)
sys.exit(0)
if not opt.template_file:
parser.print_help()
sys.exit(-1)
expand_prompts(
template_file=opt.template_file,
run_invoke=opt.invoke,
invoke_model=opt.model,
invoke_outdir=opt.outdir,
processes_per_gpu=opt.processes_per_gpu,
)
def expanded_invokeai_commands(
conf: OmegaConf, always_switch_models: bool = False
) -> List[List[str]]:
models = expand_values(conf.get("model"))
steps = expand_values(conf.get("steps")) or [30]
cfgs = expand_values(conf.get("cfg")) or [7.5]
samplers = expand_values(conf.get("sampler")) or ["ddim"]
seeds = expand_values(conf.get("seed")) or [0]
dimensions = expand_values(conf.get("dimensions")) or ["512x512"]
init_img = expand_values(conf.get("init_img")) or [""]
perlin = expand_values(conf.get("perlin")) or [0]
threshold = expand_values(conf.get("threshold")) or [0]
strength = expand_values(conf.get("strength")) or [0.75]
prompts = expand_prompt(conf.get("prompt")) or ["banana sushi"]
cross_product = product(
*[
models,
seeds,
prompts,
samplers,
cfgs,
steps,
perlin,
threshold,
init_img,
strength,
dimensions,
]
)
previous_model = None
result = list()
for p in cross_product:
(
model,
seed,
prompt,
sampler,
cfg,
step,
perlin,
threshold,
init_img,
strength,
dimensions,
) = tuple(p)
(width, height) = dimensions.split("x")
switch_args = (
f"!switch {model}\n"
if always_switch_models or previous_model != model
else ""
)
image_args = f"-I{init_img} -f{strength}" if init_img else ""
command = f"{switch_args}{prompt} -S{seed} -A{sampler} -C{cfg} -s{step} {image_args} --perlin={perlin} --threshold={threshold} -W{width} -H{height}"
result.append(command)
previous_model = model
return result
def expand_prompt(
stanza: str | dict | listconfig.ListConfig | dictconfig.DictConfig,
) -> list | range:
if not stanza:
return None
if isinstance(stanza, listconfig.ListConfig):
return stanza
if isinstance(stanza, str):
return [stanza]
if not isinstance(stanza, dictconfig.DictConfig):
raise ValueError(f"Unrecognized template: {stanza}")
if not (template := stanza.get("template")):
raise KeyError('"prompt" section must contain a "template" definition')
fragment_labels = re.findall("{([^{}]+?)}", template)
if len(fragment_labels) == 0:
return [template]
fragments = [[{x: y} for y in stanza.get(x)] for x in fragment_labels]
dicts = merge(product(*fragments))
return [template.format(**x) for x in dicts]
def merge(dicts: Iterable) -> List[dict]:
result = list()
for x in dicts:
to_merge = dict()
for item in x:
to_merge = to_merge | item
result.append(to_merge)
return result
def expand_values(stanza: str | dict | listconfig.ListConfig) -> list | range:
if not stanza:
return None
if isinstance(stanza, listconfig.ListConfig):
return stanza
elif match := re.match("^(-?\d+);(-?\d+)(;(\d+))?", str(stanza)):
(start, stop, step) = (
int(match.group(1)),
int(match.group(2)),
int(match.group(4)) or 1,
)
return range(start, stop + step, step)
elif match := re.match("^(-?[\d.]+);(-?[\d.]+)(;([\d.]+))?", str(stanza)):
(start, stop, step) = (
float(match.group(1)),
float(match.group(2)),
float(match.group(4)) or 1.0,
)
return np.arange(start, stop + step, step).tolist()
else:
return [stanza]
def _yaml_to_json(yaml_input: str) -> str:
"""
Converts a yaml string into a json string. Used internally
to generate the example template file.
"""
with io.StringIO(yaml_input) as yaml_in:
data = yaml.safe_load(yaml_in)
return json.dumps(data, indent=2)
HELP = """
This script takes a prompt template file that contains multiple
alternative values for the prompt and its generation arguments (such
as steps). It then expands out the prompts using all combinations of
arguments and either prints them to the terminal's standard output, or
feeds the prompts directly to the invokeai command-line interface.
Call this script again with --instructions (-i) for verbose instructions.
"""
INSTRUCTIONS = f"""
== INTRODUCTION ==
This script takes a prompt template file that contains multiple
alternative values for the prompt and its generation arguments (such
as steps). It then expands out the prompts using all combinations of
arguments and either prints them to the terminal's standard output, or
feeds the prompts directly to the invokeai command-line interface.
If the optional --invoke argument is provided, then the generated
prompts will be fed directly to invokeai for image generation. You
will likely want to add the --outdir option in order to save the image
files to their own folder.
{sys.argv[0]} --invoke --outdir=/tmp/outputs my_template.yaml
If --invoke isn't specified, the expanded prompts will be printed to
output. You can capture them to a file for inspection and editing this
way:
{sys.argv[0]} my_template.yaml > prompts.txt
And then feed them to invokeai this way:
invokeai --outdir=/tmp/outputs < prompts.txt
Note that after invokeai finishes processing the list of prompts, the
output directory will contain a markdown file named `log.md`
containing annotated images. You can open this file using an e-book
reader such as the cross-platform Calibre eBook reader
(https://calibre-ebook.com/).
== FORMAT OF THE TEMPLATES FILE ==
This will generate an example template file that you can get
started with:
{sys.argv[0]} --example > example.yaml
An excerpt from the top of this file looks like this:
model:
- stable-diffusion-1.5
- stable-diffusion-2.1-base
steps: 30;50;1 # start steps at 30 and go up to 50, incrementing by 1 each time
seed: 50 # fixed constant, seed=50
cfg: # list of CFG values to try
- 7
- 8
- 12
prompt: a walk in the park # constant value
In more detail, the template file can one or more of the
following sections:
- model:
- steps:
- seed:
- cfg:
- sampler:
- prompt:
- init_img:
- perlin:
- threshold:
- strength
- Each section can have a constant value such as this:
steps: 50
- Or a range of numeric values in the format:
steps: <start>;<stop>;<step> (note semicolon, not colon!)
- Or a list of values in the format:
- value1
- value2
- value3
The "prompt:" section is special. It can accept a constant value:
prompt: a walk in the woods in the style of donatello
Or it can accept a list of prompts:
prompt:
- a walk in the woods
- a walk on the beach
Or it can accept a templated list of prompts. These allow you to
define a series of phrases, each of which is a list. You then combine
them together into a prompt template in this way:
prompt:
style:
- oil painting
- watercolor
- comic book
- studio photography
subject:
- sunny meadow in the mountains
- gathering storm in the mountains
template: a {{subject}} in the style of {{style}}
In the example above, the phrase names "style" and "subject" are
examples only. You can use whatever you like. However, the "template:"
field is required. The output will be:
"a sunny meadow in the mountains in the style of an oil painting"
"a sunny meadow in the mountains in the style of watercolor masterpiece"
...
"a gathering storm in the mountains in the style of an ink sketch"
== SUPPORT FOR JSON FORMAT ==
For those who prefer the JSON format, this script will accept JSON
template files as well. Please run "{sys.argv[0]} --json-example"
to print out a version of the example template file in json format.
You may save it to disk and use it as a starting point for your own
template this way:
{sys.argv[0]} --json-example > template.json
"""
EXAMPLE_TEMPLATE_FILE = """
model: stable-diffusion-1.5
steps: 30;50;10
seed: 50
dimensions: 512x512
perlin: 0.0
threshold: 0
cfg:
- 7
- 12
sampler:
- k_euler_a
- k_lms
prompt:
style:
- oil painting
- watercolor
location:
- the mountains
- a desert
object:
- luxurious dwelling
- crude tent
template: a {object} in {location}, in the style of {style}
"""
if __name__ == "__main__":
main()

25
ldm/invoke/pngwriter.py
View File

@ -9,6 +9,8 @@ Exports function retrieve_metadata(path)
import os
import re
import json
from pathlib import Path
from filelock import FileLock
from PIL import PngImagePlugin, Image
# -------------------image generation utils-----
@ -19,8 +21,23 @@ class PngWriter:
self.outdir = outdir
os.makedirs(outdir, exist_ok=True)
def unique_prefix(self)->str:
next_prefix_file = Path(self.outdir,'.next_prefix')
next_prefix_lock = Path(self.outdir,'.next_prefix.lock')
prefix = 0
with FileLock(next_prefix_lock):
if not next_prefix_file.exists():
prefix = self._unused_prefix()
else:
with open(next_prefix_file,'r') as file:
prefix=int(file.readline() or int(self._unused_prefix())-1)
prefix+=1
with open(next_prefix_file,'w') as file:
file.write(str(prefix))
return f'{prefix:06}'
# gives the next unique prefix in outdir
def unique_prefix(self):
def _unused_prefix(self):
# sort reverse alphabetically until we find max+1
dirlist = sorted(os.listdir(self.outdir), reverse=True)
# find the first filename that matches our pattern or return 000000.0.png
@ -91,14 +108,12 @@ class PromptFormatter:
switches.append(f'-H{opt.height or t2i.height}')
switches.append(f'-C{opt.cfg_scale or t2i.cfg_scale}')
switches.append(f'-A{opt.sampler_name or t2i.sampler_name}')
# to do: put model name into the t2i object
# switches.append(f'--model{t2i.model_name}')
if opt.seamless or t2i.seamless:
switches.append(f'--seamless')
switches.append('--seamless')
if opt.init_img:
switches.append(f'-I{opt.init_img}')
if opt.fit:
switches.append(f'--fit')
switches.append('--fit')
if opt.strength and opt.init_img is not None:
switches.append(f'-f{opt.strength or t2i.strength}')
if opt.gfpgan_strength:

1
pyproject.toml
View File

@ -118,6 +118,7 @@ requires-python = ">=3.9, <3.11"
"invokeai-configure" = "ldm.invoke.config.invokeai_configure:main"
"invokeai-merge" = "ldm.invoke.merge_diffusers:main"
"invokeai-ti" = "ldm.invoke.training.textual_inversion:main"
"invokeai-batch" = "ldm.invoke.dynamic_prompts:main"
[project.urls]
"Bug Reports" = "https://github.com/invoke-ai/InvokeAI/issues"

9
scripts/dynamic_prompts.py Executable file
View File

@ -0,0 +1,9 @@
#!/usr/bin/env python
"""
Simple script to generate a file of InvokeAI prompts and settings
that scan across steps and other parameters.
"""
import ldm.invoke.dynamic_prompts
ldm.invoke.dynamic_prompts.main()