diff --git a/README.md b/README.md index ea0b9f3..a79d92a 100644 --- a/README.md +++ b/README.md @@ -381,6 +381,7 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface - ## ChangeLog - feature: 🎉 ControlNet integration! Control the structure of generated images. +- feature: `aimg colorize` attempts to use controlnet to colorize images **10.0.1** - fix: `edit` was broken diff --git a/imaginairy/cmds.py b/imaginairy/cmds.py index 7eb1e31..e0a9f02 100644 --- a/imaginairy/cmds.py +++ b/imaginairy/cmds.py @@ -499,6 +499,8 @@ 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") + + # remove_option(edit_options, "control_mode") # remove_option(edit_options, "control_image") # remove_option(edit_options, "control_image_raw") @@ -869,18 +871,75 @@ def upscale_cmd(image_filepaths, outdir): img.save(os.path.join(outdir, os.path.basename(p))) +@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.", +) +@aimg.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) + + @click.argument("image_filepaths", nargs=-1) @aimg.command() def describe(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) @@ -1123,6 +1182,19 @@ def prepare_images(images_dir, is_person, target_size): prep_images(images_dir=images_dir, is_person=is_person, target_size=target_size) +@click.option( + "--target-size", + default=512, + type=int, + show_default=True, +) +@aimg.command("prep-images") +def colorize_image_cmd(images_dir, is_person, target_size): + from imaginairy.training_tools.image_prep import prep_images + + prep_images(images_dir=images_dir, is_person=is_person, target_size=target_size) + + @click.argument("ckpt_paths", nargs=-1) @aimg.command("prune-ckpt") def prune_ckpt(ckpt_paths): diff --git a/imaginairy/colorize.py b/imaginairy/colorize.py new file mode 100644 index 0000000..6bf6569 --- /dev/null +++ b/imaginairy/colorize.py @@ -0,0 +1,46 @@ +from PIL import Image + +from imaginairy import ImaginePrompt, imagine +from imaginairy.enhancers.describe_image_blip import generate_caption + + +def colorize_img(img): + caption = generate_caption(img) + caption = caption.replace("black and white", "color") + + prompt = ImaginePrompt( + prompt=caption, + init_image=img, + init_image_strength=0.01, + control_image=img, + control_mode="hed", + negative_prompt="black and white", + # width=img.width, + # height=img.height, + ) + result = list(imagine(prompt))[0] + colorized_img = replace_color(img, result.images["generated"]) + + prompt = ImaginePrompt( + prompt=caption, + init_image=colorized_img, + init_image_strength=0.1, + control_image=img, + control_mode="hed", + negative_prompt="black and white", + width=min(img.width, 1024), + height=min(img.height, 1024), + steps=30 + ) + result = list(imagine(prompt))[0] + colorized_img = replace_color(img, result.images["generated"]) + return colorized_img + + +def replace_color(target_img, color_src_img): + color_src_img = color_src_img.resize(target_img.size) + + _, _, value = target_img.convert("HSV").split() + hue, saturation, _ = color_src_img.convert("HSV").split() + + return Image.merge("HSV", (hue, saturation, value)).convert("RGB")