feature: aimg colorize attempts to use controlnet to colorize images

This commit is contained in:
Bryce 2023-02-23 01:28:40 -08:00 committed by Bryce Drennan
parent 7899d08181
commit 9831d8a24a
3 changed files with 119 additions and 0 deletions

View File

@ -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

View File

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

46
imaginairy/colorize.py Normal file
View File

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