mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
feature: aimg colorize
attempts to use controlnet to colorize images
This commit is contained in:
parent
7899d08181
commit
9831d8a24a
@ -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
|
||||
|
@ -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
46
imaginairy/colorize.py
Normal 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")
|
Loading…
Reference in New Issue
Block a user