refactor: split up command line tool
Bryce
1 year ago
committed by
Bryce Drennan
parent
c0fad5b2a3
commit
e3b23584ed
@ -1,6 +1,7 @@
|
||||
#!/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
import imaginairy.cli.imagine
|
||||
|
||||
if __name__ == "__main__":
|
||||
from imaginairy import cmds
|
||||
cmds.imagine_cmd()
|
||||
import imaginairy.cli.main
|
||||
imaginairy.cli.imagine.imagine_cmd()
|
||||
@ -0,0 +1,57 @@
|
||||
import logging
|
||||
|
||||
import click
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@click.argument("image_filepaths", nargs=-1)
|
||||
@click.option(
|
||||
"--outdir",
|
||||
default="./outputs/colorized",
|
||||
show_default=True,
|
||||
type=click.Path(),
|
||||
help="Where to write results to.",
|
||||
)
|
||||
@click.option(
|
||||
"-r",
|
||||
"--repeats",
|
||||
default=1,
|
||||
show_default=True,
|
||||
type=int,
|
||||
help="How many times to repeat the renders. If you provide two prompts and --repeat=3 then six images will be generated.",
|
||||
)
|
||||
@click.command("colorize")
|
||||
def colorize_cmd(image_filepaths, outdir, repeats):
|
||||
"""
|
||||
Colorize images using AI. Doesn't work very well yet.
|
||||
"""
|
||||
import os.path
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from imaginairy import LazyLoadingImage
|
||||
from imaginairy.colorize import colorize_img
|
||||
from imaginairy.log_utils import configure_logging
|
||||
|
||||
configure_logging()
|
||||
|
||||
os.makedirs(outdir, exist_ok=True)
|
||||
base_count = len(os.listdir(outdir))
|
||||
for _ in range(repeats):
|
||||
for p in tqdm(image_filepaths):
|
||||
base_count += 1
|
||||
filename = f"{base_count:06d}_{os.path.basename(p)}".lower()
|
||||
savepath = os.path.join(outdir, filename)
|
||||
if p.startswith("http"):
|
||||
img = LazyLoadingImage(url=p)
|
||||
elif os.path.isdir(p):
|
||||
print(f"Skipping directory: {p}")
|
||||
continue
|
||||
else:
|
||||
img = LazyLoadingImage(filepath=p)
|
||||
logger.info(f"Colorizing {p} and saving it to {savepath}")
|
||||
|
||||
img = colorize_img(img)
|
||||
|
||||
img.save(savepath)
|
||||
@ -0,0 +1,25 @@
|
||||
import click
|
||||
|
||||
|
||||
@click.argument("image_filepaths", nargs=-1)
|
||||
@click.command()
|
||||
def describe_cmd(image_filepaths):
|
||||
"""Generate text descriptions of images."""
|
||||
import os
|
||||
|
||||
from imaginairy import LazyLoadingImage
|
||||
from imaginairy.enhancers.describe_image_blip import generate_caption
|
||||
|
||||
imgs = []
|
||||
for p in image_filepaths:
|
||||
|
||||
if p.startswith("http"):
|
||||
img = LazyLoadingImage(url=p)
|
||||
elif os.path.isdir(p):
|
||||
print(f"Skipping directory: {p}")
|
||||
continue
|
||||
else:
|
||||
img = LazyLoadingImage(filepath=p)
|
||||
imgs.append(img)
|
||||
for img in imgs:
|
||||
print(generate_caption(img.copy()))
|
||||
@ -0,0 +1,139 @@
|
||||
import click
|
||||
|
||||
from imaginairy import config
|
||||
from imaginairy.cli.shared import (
|
||||
_imagine_cmd,
|
||||
add_options,
|
||||
common_options,
|
||||
remove_option,
|
||||
)
|
||||
|
||||
edit_options = common_options.copy()
|
||||
|
||||
|
||||
remove_option(edit_options, "model_weights_path")
|
||||
remove_option(edit_options, "init_image")
|
||||
remove_option(edit_options, "init_image_strength")
|
||||
remove_option(edit_options, "negative_prompt")
|
||||
remove_option(edit_options, "allow_compose_phase")
|
||||
|
||||
|
||||
@click.command("edit")
|
||||
@click.argument("image_paths", metavar="PATH|URL", required=True, nargs=-1)
|
||||
@click.option(
|
||||
"--image-strength",
|
||||
default=1,
|
||||
show_default=False,
|
||||
type=float,
|
||||
help="Starting image strength. Between 0 and 1.",
|
||||
)
|
||||
@click.option("--prompt", "-p", required=True, multiple=True)
|
||||
@click.option(
|
||||
"--model-weights-path",
|
||||
"--model",
|
||||
help=f"Model to use. Should be one of {', '.join(config.MODEL_SHORT_NAMES)}, or a path to custom weights.",
|
||||
show_default=True,
|
||||
default="edit",
|
||||
)
|
||||
@click.option(
|
||||
"--negative-prompt",
|
||||
default=None,
|
||||
show_default=False,
|
||||
help="Negative prompt. Things to try and exclude from images. Same negative prompt will be used for all images. A default negative prompt is used if none is selected.",
|
||||
)
|
||||
@add_options(edit_options)
|
||||
@click.pass_context
|
||||
def edit_cmd( # noqa
|
||||
ctx,
|
||||
image_paths,
|
||||
image_strength,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
prompt_strength,
|
||||
outdir,
|
||||
output_file_extension,
|
||||
repeats,
|
||||
height,
|
||||
width,
|
||||
steps,
|
||||
seed,
|
||||
upscale,
|
||||
fix_faces,
|
||||
fix_faces_fidelity,
|
||||
sampler_type,
|
||||
log_level,
|
||||
quiet,
|
||||
show_work,
|
||||
tile,
|
||||
tile_x,
|
||||
tile_y,
|
||||
mask_image,
|
||||
mask_prompt,
|
||||
mask_mode,
|
||||
mask_modify_original,
|
||||
outpaint,
|
||||
caption,
|
||||
precision,
|
||||
model_weights_path,
|
||||
model_config_path,
|
||||
prompt_library_path,
|
||||
version, # noqa
|
||||
make_gif,
|
||||
make_compare_gif,
|
||||
arg_schedules,
|
||||
make_compilation_animation,
|
||||
caption_text,
|
||||
):
|
||||
"""
|
||||
Edit an image via AI.
|
||||
|
||||
Provide paths or URLs to images and directions on how to alter them.
|
||||
|
||||
Example: aimg edit --prompt "make the dog red" my-dog.jpg my-dog2.jpg
|
||||
|
||||
Same as calling `aimg imagine --model edit --init-image my-dog.jpg --init-image-strength 1` except this command
|
||||
can batch edit images.
|
||||
"""
|
||||
allow_compose_phase = False
|
||||
return _imagine_cmd(
|
||||
ctx,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
prompt_strength,
|
||||
image_paths,
|
||||
image_strength,
|
||||
outdir,
|
||||
output_file_extension,
|
||||
repeats,
|
||||
height,
|
||||
width,
|
||||
steps,
|
||||
seed,
|
||||
upscale,
|
||||
fix_faces,
|
||||
fix_faces_fidelity,
|
||||
sampler_type,
|
||||
log_level,
|
||||
quiet,
|
||||
show_work,
|
||||
tile,
|
||||
tile_x,
|
||||
tile_y,
|
||||
allow_compose_phase,
|
||||
mask_image,
|
||||
mask_prompt,
|
||||
mask_mode,
|
||||
mask_modify_original,
|
||||
outpaint,
|
||||
caption,
|
||||
precision,
|
||||
model_weights_path,
|
||||
model_config_path,
|
||||
prompt_library_path,
|
||||
version, # noqa
|
||||
make_gif,
|
||||
make_compare_gif,
|
||||
arg_schedules,
|
||||
make_compilation_animation,
|
||||
caption_text,
|
||||
)
|
||||
@ -0,0 +1,39 @@
|
||||
import click
|
||||
|
||||
|
||||
@click.command("edit-demo")
|
||||
@click.argument("image_paths", metavar="PATH|URL", required=True, nargs=-1)
|
||||
@click.option(
|
||||
"--outdir",
|
||||
default="./outputs",
|
||||
show_default=True,
|
||||
type=click.Path(),
|
||||
help="Where to write results to.",
|
||||
)
|
||||
@click.option(
|
||||
"-h",
|
||||
"--height",
|
||||
default=512,
|
||||
show_default=True,
|
||||
type=int,
|
||||
help="Image height. Should be multiple of 8.",
|
||||
)
|
||||
@click.option(
|
||||
"-w",
|
||||
"--width",
|
||||
default=512,
|
||||
show_default=True,
|
||||
type=int,
|
||||
help="Image width. Should be multiple of 8.",
|
||||
)
|
||||
def edit_demo_cmd(image_paths, outdir, height, width):
|
||||
"""Make some fun pre-set edits to input photos."""
|
||||
|
||||
from imaginairy.log_utils import configure_logging
|
||||
from imaginairy.surprise_me import create_surprise_me_images
|
||||
|
||||
configure_logging()
|
||||
for image_path in image_paths:
|
||||
create_surprise_me_images(
|
||||
image_path, outdir=outdir, make_gif=True, width=width, height=height
|
||||
)
|
||||
@ -0,0 +1,137 @@
|
||||
import click
|
||||
|
||||
from imaginairy.cli.shared import (
|
||||
ImagineColorsCommand,
|
||||
_imagine_cmd,
|
||||
add_options,
|
||||
common_options,
|
||||
)
|
||||
|
||||
|
||||
@click.command(context_settings={"max_content_width": 140}, cls=ImagineColorsCommand)
|
||||
@click.argument("prompt_texts", nargs=-1)
|
||||
@add_options(common_options)
|
||||
@click.option(
|
||||
"--control-image",
|
||||
metavar="PATH|URL",
|
||||
help=(
|
||||
"Image used for control signal in image generation. "
|
||||
"For example if control-mode is depth, then the generated image will match the depth map "
|
||||
"extracted from the control image. "
|
||||
"Defaults to the `--init-image`"
|
||||
),
|
||||
multiple=False,
|
||||
)
|
||||
@click.option(
|
||||
"--control-image-raw",
|
||||
metavar="PATH|URL",
|
||||
help=(
|
||||
"Preprocessed image used for control signal in image generation. Like `--control-image` but "
|
||||
" expects the already extracted signal. For example the raw control image would be a depth map or"
|
||||
"pose information."
|
||||
),
|
||||
multiple=False,
|
||||
)
|
||||
@click.option(
|
||||
"--control-mode",
|
||||
default=None,
|
||||
show_default=False,
|
||||
type=click.Choice(["", "canny", "depth", "normal", "hed", "openpose"]),
|
||||
help="how the control image is used as signal",
|
||||
)
|
||||
@click.pass_context
|
||||
def imagine_cmd(
|
||||
ctx,
|
||||
prompt_texts,
|
||||
negative_prompt,
|
||||
prompt_strength,
|
||||
init_image,
|
||||
init_image_strength,
|
||||
outdir,
|
||||
output_file_extension,
|
||||
repeats,
|
||||
height,
|
||||
width,
|
||||
steps,
|
||||
seed,
|
||||
upscale,
|
||||
fix_faces,
|
||||
fix_faces_fidelity,
|
||||
sampler_type,
|
||||
log_level,
|
||||
quiet,
|
||||
show_work,
|
||||
tile,
|
||||
tile_x,
|
||||
tile_y,
|
||||
allow_compose_phase,
|
||||
mask_image,
|
||||
mask_prompt,
|
||||
mask_mode,
|
||||
mask_modify_original,
|
||||
outpaint,
|
||||
caption,
|
||||
precision,
|
||||
model_weights_path,
|
||||
model_config_path,
|
||||
prompt_library_path,
|
||||
version, # noqa
|
||||
make_gif,
|
||||
make_compare_gif,
|
||||
arg_schedules,
|
||||
make_compilation_animation,
|
||||
caption_text,
|
||||
control_image,
|
||||
control_image_raw,
|
||||
control_mode,
|
||||
):
|
||||
"""
|
||||
Generate images via AI.
|
||||
|
||||
Can be invoked via either `aimg imagine` or just `imagine`.
|
||||
"""
|
||||
return _imagine_cmd(
|
||||
ctx,
|
||||
prompt_texts,
|
||||
negative_prompt,
|
||||
prompt_strength,
|
||||
init_image,
|
||||
init_image_strength,
|
||||
outdir,
|
||||
output_file_extension,
|
||||
repeats,
|
||||
height,
|
||||
width,
|
||||
steps,
|
||||
seed,
|
||||
upscale,
|
||||
fix_faces,
|
||||
fix_faces_fidelity,
|
||||
sampler_type,
|
||||
log_level,
|
||||
quiet,
|
||||
show_work,
|
||||
tile,
|
||||
tile_x,
|
||||
tile_y,
|
||||
allow_compose_phase,
|
||||
mask_image,
|
||||
mask_prompt,
|
||||
mask_mode,
|
||||
mask_modify_original,
|
||||
outpaint,
|
||||
caption,
|
||||
precision,
|
||||
model_weights_path,
|
||||
model_config_path,
|
||||
prompt_library_path,
|
||||
version, # noqa
|
||||
make_gif,
|
||||
make_compare_gif,
|
||||
arg_schedules,
|
||||
make_compilation_animation,
|
||||
caption_text,
|
||||
control_image,
|
||||
control_image_raw,
|
||||
control_mode,
|
||||
)
|
||||
@ -0,0 +1,70 @@
|
||||
import logging
|
||||
|
||||
import click
|
||||
|
||||
from imaginairy.cli.colorize import colorize_cmd
|
||||
from imaginairy.cli.describe import describe_cmd
|
||||
from imaginairy.cli.edit import edit_cmd
|
||||
from imaginairy.cli.edit_demo import edit_demo_cmd
|
||||
from imaginairy.cli.imagine import imagine_cmd
|
||||
from imaginairy.cli.shared import ColorShell, ImagineColorsCommand
|
||||
from imaginairy.cli.train import prep_images_cmd, prune_ckpt_cmd, train_concept_cmd
|
||||
from imaginairy.cli.upscale import upscale_cmd
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@click.command(
|
||||
prompt="🤖🧠> ",
|
||||
intro="Starting imaginAIry shell...",
|
||||
help_headers_color="yellow",
|
||||
help_options_color="green",
|
||||
context_settings={"max_content_width": 140},
|
||||
cls=ColorShell,
|
||||
)
|
||||
@click.pass_context
|
||||
def aimg(ctx):
|
||||
"""
|
||||
🤖🧠 ImaginAIry.
|
||||
|
||||
Pythonic generation of images via AI
|
||||
"""
|
||||
import sys
|
||||
|
||||
is_shell = len(sys.argv) == 1
|
||||
if is_shell:
|
||||
print(ctx.get_help())
|
||||
|
||||
|
||||
aimg.command_class = ImagineColorsCommand
|
||||
|
||||
|
||||
aimg.add_command(colorize_cmd, name="colorize")
|
||||
aimg.add_command(describe_cmd, name="describe")
|
||||
aimg.add_command(edit_cmd, name="edit")
|
||||
aimg.add_command(edit_demo_cmd, name="edit-demo")
|
||||
aimg.add_command(imagine_cmd, name="imagine")
|
||||
aimg.add_command(prep_images_cmd, name="prep-images")
|
||||
aimg.add_command(prune_ckpt_cmd, name="prune-ckpt")
|
||||
aimg.add_command(train_concept_cmd, name="train-concept")
|
||||
aimg.add_command(upscale_cmd, name="upscale")
|
||||
|
||||
|
||||
@aimg.command()
|
||||
def version():
|
||||
"""Print the version."""
|
||||
from imaginairy.version import get_version
|
||||
|
||||
print(get_version())
|
||||
|
||||
|
||||
@aimg.command("system-info")
|
||||
def system_info():
|
||||
"""
|
||||
Display system information. Submit this when reporting bugs.
|
||||
"""
|
||||
from imaginairy.debug_info import get_debug_info
|
||||
|
||||
for k, v in get_debug_info().items():
|
||||
k += ":"
|
||||
click.secho(f"{k: <30} {v}")
|
||||
@ -0,0 +1,501 @@
|
||||
import logging
|
||||
import math
|
||||
|
||||
import click
|
||||
from click_help_colors import HelpColorsCommand, HelpColorsMixin
|
||||
from click_shell import Shell
|
||||
|
||||
from imaginairy import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _imagine_cmd(
|
||||
ctx,
|
||||
prompt_texts,
|
||||
negative_prompt,
|
||||
prompt_strength,
|
||||
init_image,
|
||||
init_image_strength,
|
||||
outdir,
|
||||
output_file_extension,
|
||||
repeats,
|
||||
height,
|
||||
width,
|
||||
steps,
|
||||
seed,
|
||||
upscale,
|
||||
fix_faces,
|
||||
fix_faces_fidelity,
|
||||
sampler_type,
|
||||
log_level,
|
||||
quiet,
|
||||
show_work,
|
||||
tile,
|
||||
tile_x,
|
||||
tile_y,
|
||||
allow_compose_phase,
|
||||
mask_image,
|
||||
mask_prompt,
|
||||
mask_mode,
|
||||
mask_modify_original,
|
||||
outpaint,
|
||||
caption,
|
||||
precision,
|
||||
model_weights_path,
|
||||
model_config_path,
|
||||
prompt_library_path,
|
||||
version=False, # noqa
|
||||
make_gif=False,
|
||||
make_compare_gif=False,
|
||||
arg_schedules=None,
|
||||
make_compilation_animation=False,
|
||||
caption_text="",
|
||||
control_image=None,
|
||||
control_image_raw=None,
|
||||
control_mode="",
|
||||
):
|
||||
"""Have the AI generate images. alias:imagine."""
|
||||
|
||||
if ctx.invoked_subcommand is not None:
|
||||
return
|
||||
|
||||
if version:
|
||||
from imaginairy.version import get_version
|
||||
|
||||
print(get_version())
|
||||
return
|
||||
|
||||
if quiet:
|
||||
log_level = "ERROR"
|
||||
|
||||
import sys
|
||||
|
||||
if len(sys.argv) > 1:
|
||||
msg = (
|
||||
"✨ Generate images faster using a persistent shell session. Just run `aimg` to start. "
|
||||
"This makes generation and editing much quicker since the model can stay loaded in memory.\n"
|
||||
)
|
||||
print(msg)
|
||||
|
||||
from imaginairy.log_utils import configure_logging
|
||||
|
||||
configure_logging(log_level)
|
||||
|
||||
if isinstance(init_image, str):
|
||||
init_images = [init_image]
|
||||
else:
|
||||
init_images = init_image
|
||||
total_image_count = len(prompt_texts) * max(len(init_images), 1) * repeats
|
||||
logger.info(
|
||||
f"Received {len(prompt_texts)} prompt(s) and {len(init_images)} input image(s). Will repeat the generations {repeats} times to create {total_image_count} images."
|
||||
)
|
||||
|
||||
from imaginairy import ImaginePrompt, LazyLoadingImage, imagine_image_files
|
||||
|
||||
if control_image and control_image.startswith("http"):
|
||||
control_image = LazyLoadingImage(url=control_image)
|
||||
|
||||
if control_image_raw and control_image_raw.startswith("http"):
|
||||
control_image_raw = LazyLoadingImage(url=control_image_raw)
|
||||
|
||||
new_init_images = []
|
||||
for _init_image in init_images:
|
||||
if _init_image and _init_image.startswith("http"):
|
||||
_init_image = LazyLoadingImage(url=_init_image)
|
||||
new_init_images.append(_init_image)
|
||||
init_images = new_init_images
|
||||
if not init_images:
|
||||
init_images = [None]
|
||||
|
||||
if mask_image and mask_image.startswith("http"):
|
||||
mask_image = LazyLoadingImage(url=mask_image)
|
||||
|
||||
prompts = []
|
||||
prompt_expanding_iterators = {}
|
||||
from imaginairy.enhancers.prompt_expansion import expand_prompts
|
||||
|
||||
for _ in range(repeats):
|
||||
for prompt_text in prompt_texts:
|
||||
if prompt_text not in prompt_expanding_iterators:
|
||||
prompt_expanding_iterators[prompt_text] = expand_prompts(
|
||||
n=math.inf,
|
||||
prompt_text=prompt_text,
|
||||
prompt_library_paths=prompt_library_path,
|
||||
)
|
||||
prompt_iterator = prompt_expanding_iterators[prompt_text]
|
||||
if tile:
|
||||
_tile_mode = "xy"
|
||||
elif tile_x:
|
||||
_tile_mode = "x"
|
||||
elif tile_y:
|
||||
_tile_mode = "y"
|
||||
else:
|
||||
_tile_mode = ""
|
||||
for _init_image in init_images:
|
||||
prompt = ImaginePrompt(
|
||||
next(prompt_iterator),
|
||||
negative_prompt=negative_prompt,
|
||||
prompt_strength=prompt_strength,
|
||||
init_image=_init_image,
|
||||
init_image_strength=init_image_strength,
|
||||
control_image=control_image,
|
||||
control_image_raw=control_image_raw,
|
||||
control_mode=control_mode,
|
||||
seed=seed,
|
||||
sampler_type=sampler_type,
|
||||
steps=steps,
|
||||
height=height,
|
||||
width=width,
|
||||
mask_image=mask_image,
|
||||
mask_prompt=mask_prompt,
|
||||
mask_mode=mask_mode,
|
||||
mask_modify_original=mask_modify_original,
|
||||
outpaint=outpaint,
|
||||
upscale=upscale,
|
||||
fix_faces=fix_faces,
|
||||
fix_faces_fidelity=fix_faces_fidelity,
|
||||
tile_mode=_tile_mode,
|
||||
allow_compose_phase=allow_compose_phase,
|
||||
model=model_weights_path,
|
||||
model_config_path=model_config_path,
|
||||
caption_text=caption_text,
|
||||
)
|
||||
from imaginairy.prompt_schedules import (
|
||||
parse_schedule_strs,
|
||||
prompt_mutator,
|
||||
)
|
||||
|
||||
if arg_schedules:
|
||||
schedules = parse_schedule_strs(arg_schedules)
|
||||
for new_prompt in prompt_mutator(prompt, schedules):
|
||||
prompts.append(new_prompt)
|
||||
else:
|
||||
prompts.append(prompt)
|
||||
|
||||
filenames = imagine_image_files(
|
||||
prompts,
|
||||
outdir=outdir,
|
||||
record_step_images=show_work,
|
||||
output_file_extension=output_file_extension,
|
||||
print_caption=caption,
|
||||
precision=precision,
|
||||
make_gif=make_gif,
|
||||
make_compare_gif=make_compare_gif,
|
||||
)
|
||||
if make_compilation_animation:
|
||||
import os.path
|
||||
|
||||
ext = make_compilation_animation
|
||||
|
||||
compilation_outdir = os.path.join(outdir, "compilations")
|
||||
os.makedirs(compilation_outdir, exist_ok=True)
|
||||
base_count = len(os.listdir(compilation_outdir))
|
||||
new_filename = os.path.join(
|
||||
compilation_outdir, f"{base_count:04d}_compilation.{ext}"
|
||||
)
|
||||
comp_imgs = [LazyLoadingImage(filepath=f) for f in filenames]
|
||||
comp_imgs.reverse()
|
||||
|
||||
from imaginairy.animations import make_bounce_animation
|
||||
|
||||
make_bounce_animation(
|
||||
outpath=new_filename,
|
||||
imgs=comp_imgs,
|
||||
start_pause_duration_ms=1500,
|
||||
end_pause_duration_ms=1000,
|
||||
)
|
||||
|
||||
logger.info(f"[compilation] saved to: {new_filename}")
|
||||
|
||||
|
||||
def add_options(options):
|
||||
def _add_options(func):
|
||||
for option in reversed(options):
|
||||
func = option(func)
|
||||
return func
|
||||
|
||||
return _add_options
|
||||
|
||||
|
||||
def replace_option(options, option_name, new_option):
|
||||
for i, option in enumerate(options):
|
||||
|
||||
if option.name == option_name:
|
||||
options[i] = new_option
|
||||
return
|
||||
raise ValueError(f"Option {option_name} not found")
|
||||
|
||||
|
||||
def remove_option(options, option_name):
|
||||
for i, option_dec in enumerate(options):
|
||||
|
||||
def temp_f():
|
||||
return True
|
||||
|
||||
temp_f = option_dec(temp_f)
|
||||
option = temp_f.__click_params__[0]
|
||||
|
||||
if option.name == option_name:
|
||||
del options[i]
|
||||
return
|
||||
raise ValueError(f"Option {option_name} not found")
|
||||
|
||||
|
||||
common_options = [
|
||||
click.option(
|
||||
"--negative-prompt",
|
||||
default=None,
|
||||
show_default=False,
|
||||
help="Negative prompt. Things to try and exclude from images. Same negative prompt will be used for all images.",
|
||||
),
|
||||
click.option(
|
||||
"--prompt-strength",
|
||||
default=7.5,
|
||||
show_default=True,
|
||||
help="How closely to follow the prompt. Image looks unnatural at higher values",
|
||||
),
|
||||
click.option(
|
||||
"--init-image",
|
||||
metavar="PATH|URL",
|
||||
help="Starting image.",
|
||||
multiple=True,
|
||||
),
|
||||
click.option(
|
||||
"--init-image-strength",
|
||||
default=None,
|
||||
show_default=False,
|
||||
type=float,
|
||||
help="Starting image strength. Between 0 and 1.",
|
||||
),
|
||||
click.option(
|
||||
"--outdir",
|
||||
default="./outputs",
|
||||
show_default=True,
|
||||
type=click.Path(),
|
||||
help="Where to write results to.",
|
||||
),
|
||||
click.option(
|
||||
"--output-file-extension",
|
||||
default="jpg",
|
||||
show_default=True,
|
||||
type=click.Choice(["jpg", "png"]),
|
||||
help="Where to write results to.",
|
||||
),
|
||||
click.option(
|
||||
"-r",
|
||||
"--repeats",
|
||||
default=1,
|
||||
show_default=True,
|
||||
type=int,
|
||||
help="How many times to repeat the renders. If you provide two prompts and --repeat=3 then six images will be generated.",
|
||||
),
|
||||
click.option(
|
||||
"-h",
|
||||
"--height",
|
||||
default=None,
|
||||
show_default=True,
|
||||
type=int,
|
||||
help="Image height. Should be multiple of 8.",
|
||||
),
|
||||
click.option(
|
||||
"-w",
|
||||
"--width",
|
||||
default=None,
|
||||
show_default=True,
|
||||
type=int,
|
||||
help="Image width. Should be multiple of 8.",
|
||||
),
|
||||
click.option(
|
||||
"--steps",
|
||||
default=None,
|
||||
type=int,
|
||||
show_default=True,
|
||||
help="How many diffusion steps to run. More steps, more detail, but with diminishing returns.",
|
||||
),
|
||||
click.option(
|
||||
"--seed",
|
||||
default=None,
|
||||
type=int,
|
||||
help="What seed to use for randomness. Allows reproducible image renders.",
|
||||
),
|
||||
click.option("--upscale", is_flag=True),
|
||||
click.option("--fix-faces", is_flag=True),
|
||||
click.option(
|
||||
"--fix-faces-fidelity",
|
||||
default=None,
|
||||
type=float,
|
||||
help="How faithful to the original should face enhancement be. 1 = best fidelity, 0 = best looking face.",
|
||||
),
|
||||
click.option(
|
||||
"--sampler-type",
|
||||
"--sampler",
|
||||
default=config.DEFAULT_SAMPLER,
|
||||
show_default=True,
|
||||
type=click.Choice(config.SAMPLER_TYPE_OPTIONS),
|
||||
help="What sampling strategy to use.",
|
||||
),
|
||||
click.option(
|
||||
"--log-level",
|
||||
default="INFO",
|
||||
show_default=True,
|
||||
type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR"]),
|
||||
help="What level of logs to show.",
|
||||
),
|
||||
click.option(
|
||||
"--quiet",
|
||||
"-q",
|
||||
is_flag=True,
|
||||
help="Suppress logs. Alias of `--log-level ERROR`.",
|
||||
),
|
||||
click.option(
|
||||
"--show-work",
|
||||
default=False,
|
||||
is_flag=True,
|
||||
help="Output a debug images to `steps` folder.",
|
||||
),
|
||||
click.option(
|
||||
"--tile",
|
||||
is_flag=True,
|
||||
help="Any images rendered will be tileable in both X and Y directions.",
|
||||
),
|
||||
click.option(
|
||||
"--tile-x",
|
||||
is_flag=True,
|
||||
help="Any images rendered will be tileable in the X direction.",
|
||||
),
|
||||
click.option(
|
||||
"--tile-y",
|
||||
is_flag=True,
|
||||
help="Any images rendered will be tileable in the Y direction.",
|
||||
),
|
||||
click.option(
|
||||
"--allow-compose-phase/--no-compose-phase",
|
||||
default=True,
|
||||
help="Allow the image to be composed at a lower resolution.",
|
||||
),
|
||||
click.option(
|
||||
"--mask-image",
|
||||
metavar="PATH|URL",
|
||||
help="A mask to use for inpainting. White gets painted, Black is left alone.",
|
||||
),
|
||||
click.option(
|
||||
"--mask-prompt",
|
||||
help=(
|
||||
"Describe what you want masked and the AI will mask it for you. "
|
||||
"You can describe complex masks with AND, OR, NOT keywords and parentheses. "
|
||||
"The strength of each mask can be modified with {*1.5} notation. \n\n"
|
||||
"Examples: \n"
|
||||
"car AND (wheels{*1.1} OR trunk OR engine OR windows OR headlights) AND NOT (truck OR headlights){*10}\n"
|
||||
"fruit|fruit stem"
|
||||
),
|
||||
),
|
||||
click.option(
|
||||
"--mask-mode",
|
||||
default="replace",
|
||||
show_default=True,
|
||||
type=click.Choice(["keep", "replace"]),
|
||||
help="Should we replace the masked area or keep it?",
|
||||
),
|
||||
click.option(
|
||||
"--mask-modify-original",
|
||||
default=True,
|
||||
is_flag=True,
|
||||
help="After the inpainting is done, apply the changes to a copy of the original image.",
|
||||
),
|
||||
click.option(
|
||||
"--outpaint",
|
||||
help=(
|
||||
"Specify in what directions to expand the image. Values will be snapped such that output image size is multiples of 8. Examples\n"
|
||||
" `--outpaint up10,down300,left50,right50`\n"
|
||||
" `--outpaint u10,d300,l50,r50`\n"
|
||||
" `--outpaint all200`\n"
|
||||
" `--outpaint a200`\n"
|
||||
),
|
||||
default="",
|
||||
),
|
||||
click.option(
|
||||
"--caption",
|
||||
default=False,
|
||||
is_flag=True,
|
||||
help="Generate a text description of the generated image.",
|
||||
),
|
||||
click.option(
|
||||
"--precision",
|
||||
help="Evaluate at this precision.",
|
||||
type=click.Choice(["full", "autocast"]),
|
||||
default="autocast",
|
||||
show_default=True,
|
||||
),
|
||||
click.option(
|
||||
"--model-weights-path",
|
||||
"--model",
|
||||
help=f"Model to use. Should be one of {', '.join(config.MODEL_SHORT_NAMES)}, or a path to custom weights.",
|
||||
show_default=True,
|
||||
default=config.DEFAULT_MODEL,
|
||||
),
|
||||
click.option(
|
||||
"--model-config-path",
|
||||
help="Model config file to use. If a model name is specified, the appropriate config will be used.",
|
||||
show_default=True,
|
||||
default=None,
|
||||
),
|
||||
click.option(
|
||||
"--prompt-library-path",
|
||||
help="Path to folder containing phrase lists in txt files. Use txt filename in prompt: {_filename_}.",
|
||||
type=click.Path(exists=True),
|
||||
default=None,
|
||||
multiple=True,
|
||||
),
|
||||
click.option(
|
||||
"--version",
|
||||
default=False,
|
||||
is_flag=True,
|
||||
help="Print the version and exit.",
|
||||
),
|
||||
click.option(
|
||||
"--gif",
|
||||
"make_gif",
|
||||
default=False,
|
||||
is_flag=True,
|
||||
help="Create a gif of the generation.",
|
||||
),
|
||||
click.option(
|
||||
"--compare-gif",
|
||||
"make_compare_gif",
|
||||
default=False,
|
||||
is_flag=True,
|
||||
help="Create a gif comparing the original image to the modified one.",
|
||||
),
|
||||
click.option(
|
||||
"--arg-schedule",
|
||||
"arg_schedules",
|
||||
multiple=True,
|
||||
help="Schedule how an argument should change over several generations. Format: `--arg-schedule arg_name[start:end:increment]` or `--arg-schedule arg_name[val,val2,val3]`",
|
||||
),
|
||||
click.option(
|
||||
"--compilation-anim",
|
||||
"make_compilation_animation",
|
||||
default=None,
|
||||
type=click.Choice(["gif", "mp4"]),
|
||||
help="Generate an animation composed of all the images generated in this run. Defaults to gif but `--compilation-anim mp4` will generate an mp4 instead.",
|
||||
),
|
||||
click.option(
|
||||
"--caption-text",
|
||||
"caption_text",
|
||||
default=None,
|
||||
help="Specify the text to write onto the image",
|
||||
type=str,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ColorShell(HelpColorsMixin, Shell):
|
||||
pass
|
||||
|
||||
|
||||
class ImagineColorsCommand(HelpColorsCommand):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.help_headers_color = "yellow"
|
||||
self.help_options_color = "green"
|
||||
@ -0,0 +1,261 @@
|
||||
import logging
|
||||
|
||||
import click
|
||||
|
||||
from imaginairy import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@click.option(
|
||||
"--concept-label",
|
||||
help=(
|
||||
'The concept you are training on. Usually "a photo of [person or thing] [classname]" is what you should use.'
|
||||
),
|
||||
required=True,
|
||||
)
|
||||
@click.option(
|
||||
"--concept-images-dir",
|
||||
type=click.Path(),
|
||||
required=True,
|
||||
help="Where to find the pre-processed concept images to train on.",
|
||||
)
|
||||
@click.option(
|
||||
"--class-label",
|
||||
help=(
|
||||
'What class of things does the concept belong to. For example, if you are training on "a painting of a George Washington", '
|
||||
'you might use "a painting of a man" as the class label. We use this to prevent the model from overfitting.'
|
||||
),
|
||||
default="a photo of *",
|
||||
)
|
||||
@click.option(
|
||||
"--class-images-dir",
|
||||
type=click.Path(),
|
||||
required=True,
|
||||
help="Where to find the pre-processed class images to train on.",
|
||||
)
|
||||
@click.option(
|
||||
"--n-class-images",
|
||||
type=int,
|
||||
default=300,
|
||||
help="Number of class images to generate.",
|
||||
)
|
||||
@click.option(
|
||||
"--model-weights-path",
|
||||
"--model",
|
||||
"model",
|
||||
help=f"Model to use. Should be one of {', '.join(config.MODEL_SHORT_NAMES)}, or a path to custom weights.",
|
||||
show_default=True,
|
||||
default=config.DEFAULT_MODEL,
|
||||
)
|
||||
@click.option(
|
||||
"--person",
|
||||
"is_person",
|
||||
is_flag=True,
|
||||
help="Set if images are of a person. Will use face detection and enhancement.",
|
||||
)
|
||||
@click.option(
|
||||
"-y",
|
||||
"preconfirmed",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Bypass input confirmations.",
|
||||
)
|
||||
@click.option(
|
||||
"--skip-prep",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Skip the image preparation step.",
|
||||
)
|
||||
@click.option(
|
||||
"--skip-class-img-gen",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Skip the class image generation step.",
|
||||
)
|
||||
@click.command("train-concept")
|
||||
def train_concept_cmd(
|
||||
concept_label,
|
||||
concept_images_dir,
|
||||
class_label,
|
||||
class_images_dir,
|
||||
n_class_images,
|
||||
model,
|
||||
is_person,
|
||||
preconfirmed,
|
||||
skip_prep,
|
||||
skip_class_img_gen,
|
||||
):
|
||||
"""
|
||||
Teach the model a new concept (a person, thing, style, etc).
|
||||
|
||||
Provided a directory of concept images, a concept token, and a class token, this command will train the model
|
||||
to generate images of that concept.
|
||||
|
||||
\b
|
||||
This happens in a 3-step process:
|
||||
1. Cropping and resizing your training images. If --person is set we crop to include the face.
|
||||
2. Generating a set of class images to train on. This helps prevent overfitting.
|
||||
3. Training the model on the concept and class images.
|
||||
|
||||
The output of this command is a new model weights file that you can use with the --model option.
|
||||
|
||||
\b
|
||||
## Instructions
|
||||
1. Gather a set of images of the concept you want to train on. They should show the subject from a variety of angles
|
||||
and in a variety of situations.
|
||||
2. Train the model.
|
||||
- Concept label: For a person, firstnamelastname should be fine.
|
||||
- If all the training images are photos you should add "a photo of" to the beginning of the concept label.
|
||||
- Class label: This is the category of the things beings trained on. For people this is typically "person", "man"
|
||||
or "woman".
|
||||
- If all the training images are photos you should add "a photo of" to the beginning of the class label.
|
||||
- CLass images will be generated for you if you do not provide them.
|
||||
3. Stop training before it overfits. I haven't figured this out yet.
|
||||
|
||||
|
||||
For example, if you were training on photos of a man named bill hamilton you could run the following:
|
||||
|
||||
\b
|
||||
aimg train-concept \\
|
||||
--person \\
|
||||
--concept-label "photo of billhamilton man" \\
|
||||
--concept-images-dir ./images/billhamilton \\
|
||||
--class-label "photo of a man" \\
|
||||
--class-images-dir ./images/man
|
||||
|
||||
When you use the model you should prompt with `firstnamelastname classname` (e.g. `billhamilton man`).
|
||||
|
||||
You can find a lot of relevant instructions here: https://github.com/JoePenna/Dreambooth-Stable-Diffusion
|
||||
"""
|
||||
from imaginairy.utils import get_device
|
||||
|
||||
if "mps" in get_device():
|
||||
click.secho(
|
||||
"⚠️ MPS (MacOS) is not supported for training. Please use a GPU or CPU.",
|
||||
fg="yellow",
|
||||
)
|
||||
return
|
||||
|
||||
import os.path
|
||||
|
||||
from imaginairy.train import train_diffusion_model
|
||||
from imaginairy.training_tools.image_prep import (
|
||||
create_class_images,
|
||||
get_image_filenames,
|
||||
prep_images,
|
||||
)
|
||||
|
||||
target_size = 512
|
||||
# Step 1. Crop and enhance the training images
|
||||
prepped_images_path = os.path.join(concept_images_dir, "prepped-images")
|
||||
image_filenames = get_image_filenames(concept_images_dir)
|
||||
click.secho(
|
||||
f'\n🤖🧠 Training "{concept_label}" based on {len(image_filenames)} images.\n'
|
||||
)
|
||||
|
||||
if not skip_prep:
|
||||
msg = (
|
||||
f"Creating cropped copies of the {len(image_filenames)} concept images\n"
|
||||
f" Is Person: {is_person}\n"
|
||||
f" Source: {concept_images_dir}\n"
|
||||
f" Dest: {prepped_images_path}\n"
|
||||
)
|
||||
logger.info(msg)
|
||||
if not is_person:
|
||||
click.secho("⚠️ the --person flag was not set. ", fg="yellow")
|
||||
|
||||
if not preconfirmed and not click.confirm("Continue?"):
|
||||
return
|
||||
|
||||
prep_images(
|
||||
images_dir=concept_images_dir, is_person=is_person, target_size=target_size
|
||||
)
|
||||
concept_images_dir = prepped_images_path
|
||||
|
||||
if not skip_class_img_gen:
|
||||
# Step 2. Generate class images
|
||||
class_image_filenames = get_image_filenames(class_images_dir)
|
||||
images_needed = max(n_class_images - len(class_image_filenames), 0)
|
||||
logger.info(f"Generating {n_class_images} class images in {class_images_dir}")
|
||||
logger.info(
|
||||
f"{len(class_image_filenames)} existing class images found so only generating {images_needed}."
|
||||
)
|
||||
if not preconfirmed and not click.confirm("Continue?"):
|
||||
return
|
||||
create_class_images(
|
||||
class_description=class_label,
|
||||
output_folder=class_images_dir,
|
||||
num_images=n_class_images,
|
||||
)
|
||||
|
||||
logger.info("Training the model...")
|
||||
if not preconfirmed and not click.confirm("Continue?"):
|
||||
return
|
||||
|
||||
# Step 3. Train the model
|
||||
train_diffusion_model(
|
||||
concept_label=concept_label,
|
||||
concept_images_dir=concept_images_dir,
|
||||
class_label=class_label,
|
||||
class_images_dir=class_images_dir,
|
||||
weights_location=model,
|
||||
logdir="logs",
|
||||
learning_rate=1e-6,
|
||||
accumulate_grad_batches=32,
|
||||
)
|
||||
|
||||
|
||||
@click.argument("ckpt_paths", nargs=-1)
|
||||
@click.command("prune-ckpt")
|
||||
def prune_ckpt_cmd(ckpt_paths):
|
||||
"""
|
||||
Prune a checkpoint file.
|
||||
|
||||
This will remove the optimizer state from the checkpoint file.
|
||||
This is useful if you want to use the checkpoint file for inference and save a lot of disk space
|
||||
|
||||
Example:
|
||||
aimg prune-ckpt ./path/to/checkpoint.ckpt
|
||||
"""
|
||||
from imaginairy.training_tools.prune_model import prune_diffusion_ckpt
|
||||
|
||||
click.secho("Pruning checkpoint files...")
|
||||
for p in ckpt_paths:
|
||||
prune_diffusion_ckpt(p)
|
||||
|
||||
|
||||
@click.argument(
|
||||
"images_dir",
|
||||
required=True,
|
||||
)
|
||||
@click.option(
|
||||
"--person",
|
||||
"is_person",
|
||||
is_flag=True,
|
||||
help="Set if images are of a person. Will use face detection and enhancement.",
|
||||
)
|
||||
@click.option(
|
||||
"--target-size",
|
||||
default=512,
|
||||
type=int,
|
||||
show_default=True,
|
||||
)
|
||||
@click.command("prep-images")
|
||||
def prep_images_cmd(images_dir, is_person, target_size):
|
||||
"""
|
||||
Prepare a folder of images for training.
|
||||
|
||||
Prepped images will be written to the `prepped-images` subfolder.
|
||||
|
||||
All images will be cropped and resized to (default) 512x512.
|
||||
Upscaling and face enhancement will be applied as needed to smaller images.
|
||||
|
||||
Examples:
|
||||
aimg prep-images --person ./images/selfies
|
||||
aimg prep-images ./images/toy-train
|
||||
"""
|
||||
|
||||
from imaginairy.training_tools.image_prep import prep_images
|
||||
|
||||
prep_images(images_dir=images_dir, is_person=is_person, target_size=target_size)
|
||||
@ -0,0 +1,42 @@
|
||||
import logging
|
||||
|
||||
import click
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@click.argument("image_filepaths", nargs=-1)
|
||||
@click.option(
|
||||
"--outdir",
|
||||
default="./outputs/upscaled",
|
||||
show_default=True,
|
||||
type=click.Path(),
|
||||
help="Where to write results to.",
|
||||
)
|
||||
@click.command("upscale")
|
||||
def upscale_cmd(image_filepaths, outdir):
|
||||
"""
|
||||
Upscale an image 4x using AI.
|
||||
"""
|
||||
import os.path
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from imaginairy import LazyLoadingImage
|
||||
from imaginairy.enhancers.upscale_realesrgan import upscale_image
|
||||
|
||||
os.makedirs(outdir, exist_ok=True)
|
||||
|
||||
for p in tqdm(image_filepaths):
|
||||
savepath = os.path.join(outdir, os.path.basename(p))
|
||||
if p.startswith("http"):
|
||||
img = LazyLoadingImage(url=p)
|
||||
else:
|
||||
img = LazyLoadingImage(filepath=p)
|
||||
logger.info(
|
||||
f"Upscaling {p} from {img.width}x{img.height} to {img.width * 4}x{img.height * 4} and saving it to {savepath}"
|
||||
)
|
||||
|
||||
img = upscale_image(img)
|
||||
|
||||
img.save(os.path.join(outdir, os.path.basename(p)))
|
||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue