imaginAIry/imaginairy/cli/shared.py

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"