feature: `aimg colorize` attempts to use controlnet to colorize images

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

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):

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")