From 9c48b749d87cc99982bc0b90fd9618dd05913233 Mon Sep 17 00:00:00 2001 From: Bryce Date: Tue, 5 Mar 2024 17:38:52 -0800 Subject: [PATCH] feature: script for running imaginairy in the modal.com cloud --- imaginairy/api/generate.py | 165 +++++++++++++++----------- imaginairy/utils/gitignore.py | 117 +++++++++++++++++++ scripts/modal_entrypoint.py | 211 ++++++++++++++++++++++++++++++++++ tox.ini | 2 +- 4 files changed, 427 insertions(+), 68 deletions(-) create mode 100644 imaginairy/utils/gitignore.py create mode 100644 scripts/modal_entrypoint.py diff --git a/imaginairy/api/generate.py b/imaginairy/api/generate.py index 3806b34..4a2e1f5 100755 --- a/imaginairy/api/generate.py +++ b/imaginairy/api/generate.py @@ -2,6 +2,7 @@ import logging import os +from pathlib import Path from typing import TYPE_CHECKING, Callable if TYPE_CHECKING: @@ -59,8 +60,6 @@ def imagine_image_files( from imaginairy.api.video_sample import generate_video from imaginairy.utils import get_next_filenumber, prompt_normalized - from imaginairy.utils.animations import make_bounce_animation - from imaginairy.utils.img_utils import pillow_fit_image_within generated_imgs_path = os.path.join(outdir, "generated") os.makedirs(generated_imgs_path, exist_ok=True) @@ -93,73 +92,26 @@ def imagine_image_files( debug_img_callback=_record_step if record_step_images else None, add_caption=print_caption, ): - prompt = result.prompt - if prompt.is_intermediate: - # we don't save intermediate images - continue - img_str = "" - if prompt.init_image: - img_str = f"_img2img-{prompt.init_image_strength}" - - basefilename = ( - f"{base_count:06}_{prompt.seed}_{prompt.solver_type.replace('_', '')}{prompt.steps}_" - f"PS{prompt.prompt_strength}{img_str}_{prompt_normalized(prompt.prompt_text)}" + primary_filename = save_image_result( + result, + base_count, + outdir=outdir, + output_file_extension=output_file_extension, + primary_filename_type=return_filename_type, + make_gif=make_gif, + make_compare_gif=make_compare_gif, ) - for image_type in result.images: - subpath = os.path.join(outdir, image_type) - os.makedirs(subpath, exist_ok=True) - filepath = os.path.join( - subpath, f"{basefilename}_[{image_type}].{output_file_extension}" - ) - result.save(filepath, image_type=image_type) - logger.info(f" {image_type:<22} {filepath}") - if image_type == return_filename_type: - result_filenames.append(filepath) - if videogen: - try: - generate_video( - input_path=filepath, - ) - except FileNotFoundError as e: - logger.error(str(e)) - exit(1) - - if make_gif and result.progress_latents: - subpath = os.path.join(outdir, "gif") - os.makedirs(subpath, exist_ok=True) - filepath = os.path.join(subpath, f"{basefilename}.gif") - - frames = [*result.progress_latents, result.images["generated"]] - - if prompt.init_image: - resized_init_image = pillow_fit_image_within( - prompt.init_image, prompt.width, prompt.height + if not primary_filename: + continue + result_filenames.append(primary_filename) + if primary_filename and videogen: + try: + generate_video( + input_path=primary_filename, ) - frames = [resized_init_image, *frames] - frames.reverse() - make_bounce_animation( - imgs=frames, - outpath=filepath, - start_pause_duration_ms=1500, - end_pause_duration_ms=1000, - ) - image_type = "gif" - logger.info(f" {image_type:<22} {filepath}") - if make_compare_gif and prompt.init_image: - subpath = os.path.join(outdir, "gif") - os.makedirs(subpath, exist_ok=True) - filepath = os.path.join(subpath, f"{basefilename}_[compare].gif") - resized_init_image = pillow_fit_image_within( - prompt.init_image, prompt.width, prompt.height - ) - frames = [result.images["generated"], resized_init_image] - - make_bounce_animation( - imgs=frames, - outpath=filepath, - ) - image_type = "gif" - logger.info(f" {image_type:<22} {filepath}") + except FileNotFoundError as e: + logger.error(str(e)) + exit(1) base_count += 1 del result @@ -167,6 +119,85 @@ def imagine_image_files( return result_filenames +def save_image_result( + result, + base_count: int, + outdir: str | Path, + output_file_extension: str, + primary_filename_type, + make_gif=False, + make_compare_gif=False, +): + from imaginairy.utils import prompt_normalized + from imaginairy.utils.animations import make_bounce_animation + from imaginairy.utils.img_utils import pillow_fit_image_within + + prompt = result.prompt + if prompt.is_intermediate: + # we don't save intermediate images + return + + img_str = "" + if prompt.init_image: + img_str = f"_img2img-{prompt.init_image_strength}" + + basefilename = ( + f"{base_count:06}_{prompt.seed}_{prompt.solver_type.replace('_', '')}{prompt.steps}_" + f"PS{prompt.prompt_strength}{img_str}_{prompt_normalized(prompt.prompt_text)}" + ) + primary_filename = None + for image_type in result.images: + subpath = os.path.join(outdir, image_type) + os.makedirs(subpath, exist_ok=True) + filepath = os.path.join( + subpath, f"{basefilename}_[{image_type}].{output_file_extension}" + ) + result.save(filepath, image_type=image_type) + logger.info(f" {image_type:<22} {filepath}") + + if image_type == primary_filename_type: + primary_filename = filepath + + if make_gif and result.progress_latents: + subpath = os.path.join(outdir, "gif") + os.makedirs(subpath, exist_ok=True) + filepath = os.path.join(subpath, f"{basefilename}.gif") + + frames = [*result.progress_latents, result.images["generated"]] + + if prompt.init_image: + resized_init_image = pillow_fit_image_within( + prompt.init_image, prompt.width, prompt.height + ) + frames = [resized_init_image, *frames] + frames.reverse() + make_bounce_animation( + imgs=frames, + outpath=filepath, + start_pause_duration_ms=1500, + end_pause_duration_ms=1000, + ) + image_type = "gif" + logger.info(f" {image_type:<22} {filepath}") + if make_compare_gif and prompt.init_image: + subpath = os.path.join(outdir, "gif") + os.makedirs(subpath, exist_ok=True) + filepath = os.path.join(subpath, f"{basefilename}_[compare].gif") + resized_init_image = pillow_fit_image_within( + prompt.init_image, prompt.width, prompt.height + ) + frames = [result.images["generated"], resized_init_image] + + make_bounce_animation( + imgs=frames, + outpath=filepath, + ) + image_type = "gif" + logger.info(f" {image_type:<22} {filepath}") + + return primary_filename + + def imagine( prompts: "list[ImaginePrompt] | str | ImaginePrompt", precision: str = "autocast", diff --git a/imaginairy/utils/gitignore.py b/imaginairy/utils/gitignore.py new file mode 100644 index 0000000..1c75dc6 --- /dev/null +++ b/imaginairy/utils/gitignore.py @@ -0,0 +1,117 @@ +import os +from pathlib import Path + +import pathspec + + +def find_project_root(start_path): + """ + Traverse up from a starting path to find the project root. + + The project root is identified by the presence of a '.git' directory inside it. + """ + current_path = Path(start_path) + + while current_path != current_path.root: + if (current_path / ".git").is_dir(): + return str(current_path) + + if (current_path / ".hg").is_dir(): + return current_path + + if (current_path / "pyproject.toml").is_file(): + return current_path + + if (current_path / "setup.py").is_file(): + return current_path + + current_path = current_path.parent + + return None + + +ALWAYS_IGNORE = """ +.git +__pycache__ +.direnv +.eggs +.git +.hg +.mypy_cache +.nox +.tox +.venv +venv +.svn +.ipynb_checkpoints +_build +buck-out +build +dist +__pypackages__ +""" + + +def load_gitignore_spec_at_path(path): + gitignore_path = os.path.join(path, ".gitignore") + + if os.path.exists(gitignore_path): + with open(gitignore_path, encoding="utf-8") as f: + patterns = f.read().split("\n") + patterns.extend(ALWAYS_IGNORE.split("\n")) + ignore_spec = pathspec.PathSpec.from_lines("gitwildmatch", patterns) + else: + ignore_spec = pathspec.PathSpec.from_lines("gitwildmatch", []) + return ignore_spec + + +def get_nonignored_file_paths(directory, gitignore_dict=None, extensions=()): + return_relative = False + if gitignore_dict is None: + gitignore_dict = {} + return_relative = True + gitignore_dict = { + **gitignore_dict, + directory: load_gitignore_spec_at_path(directory), + } + + file_paths = [] + + for entry in os.scandir(directory): + if path_is_ignored(Path(entry.path), gitignore_dict): + continue + + if entry.is_file(): + if any(entry.path.endswith(ext) for ext in extensions): + continue + + file_paths.append(entry.path) + + elif entry.is_dir(): + subdir_file_paths = get_nonignored_file_paths( + entry.path, gitignore_dict=gitignore_dict + ) + file_paths.extend(subdir_file_paths) + if return_relative: + file_paths = [os.path.relpath(f, directory) for f in file_paths] + file_paths.sort(key=lambda p: ("/" in p, p)) + + return file_paths + + +def path_is_ignored(path: Path, gitignore_dict) -> bool: + for gitignore_path, pattern in gitignore_dict.items(): + try: + abspath = path if path.is_absolute() else Path.cwd() / path + normalized_path = abspath.resolve() + try: + relative_path = normalized_path.relative_to(gitignore_path).as_posix() + except ValueError: + return False + + except OSError: + return False + + if pattern.match_file(relative_path): + return True + return False diff --git a/scripts/modal_entrypoint.py b/scripts/modal_entrypoint.py new file mode 100644 index 0000000..9b5c374 --- /dev/null +++ b/scripts/modal_entrypoint.py @@ -0,0 +1,211 @@ +import os.path +from functools import lru_cache +from pathlib import Path + +from modal import ( + Image, + Mount, + Stub, + Volume, + gpu, +) + +os.environ["MODAL_AUTOMOUNT"] = "0" + +# find project root that is two levels up from __file__ +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +requirements_path = os.path.join(project_root, "requirements-dev.txt") + + +def file_filter(path: str): + # print(f"Checking {path}") + if "/tests/" in path: + return False + include = path in files_for_inclusion() + return include + + +@lru_cache +def files_for_inclusion(): + from imaginairy.utils.gitignore import get_nonignored_file_paths + + filepaths = get_nonignored_file_paths(project_root) + for f in filepaths: + print(f) + filepaths = [f"{project_root}/{f}" for f in filepaths if os.path.isfile(f)] + return set(filepaths) + + +local_mount = Mount.from_local_dir( + project_root, remote_path="/root/workspace/imaginairy", condition=file_filter +) + +imaginairy_image = ( + Image.debian_slim(python_version="3.10") + .apt_install( + "libglib2.0-0", "libsm6", "libxrender1", "libxext6", "ffmpeg", "libgl1" + ) + .pip_install_from_requirements(requirements_path) + .workdir("/root/workspace") + .env({"PYTHONPATH": "/root/workspace/imaginairy"}) +) +weights_cache_volume = Volume.from_name("ai-wights-cache", create_if_missing=True) +weights_cache_path = "/root/.cache/" + +stub = Stub( + "imaginairy", + mounts=[local_mount], + volumes={weights_cache_path: weights_cache_volume}, +) + + +@stub.function(gpu=gpu.A100(), container_idle_timeout=2, image=imaginairy_image) +def generate_image(imagine_prompt): + from imaginairy.api import imagine + from imaginairy.utils.log_utils import configure_logging + + configure_logging() + + results = list(imagine(imagine_prompt)) + result = results[-1] + + weights_cache_volume.commit() + return result + + +standard_prompts = [ + "a flower", + "portrait photo of a woman with a few freckles. red hair", + "photo of a bowl of fruit", + "gold coins", + "a scenic landscape", + "old growth redwood forest with a stream meandering through it. award-winning photo, diffuse lighting, beautiful, high-resolution, 4k", + "abstract desktop background", + "highway system seen from above", + "the galaxy", + "a photo of the rainforest", + "the starry night painting", + "photo of flowers", + "girl with a pearl earring", + "a painting of 1920s luncheon on a boat", + "napolean crossing the alps", + "american gothic painting", + "the tower of babel", + "god creating the {universe|earth}", + "a fishing boat in a storm on the sea of galilee. oil painting", + "the american flag", + "the tree of life. oil painting", + "the last supper. oil painting", + "the statue of liberty", + "a maze", + "HD desktop wallpaper", + "a beautiful garden", + "the garden of eden", + "the circus. {photo|oil painting}", + "a futuristic city", + "a retro spaceship", + "yosemite national park. {photo|oil painting}", + "a seacliff with a lighthouse. giant ocean waves, {photo|oil painting}", + "blueberries", + "strawberries", + "a single giant diamond", + "disneyland thomas kinkade painting", + "ancient books in an ancient library. cinematic", + "mormon missionaries", + "salt lake city", + "oil painting of heaven", + "oil painting of hell", + "an x-wing. digital art", + "star trek uss enterprise", + "a giant pile of treasure", + "the white house", + "a grizzly bear. nature photography", + "a unicorn. nature photography", + "elon musk. normal rockwell painting", + "a cybertruck", + "elon musk with a halo dressed in greek robes. oil painting", + "a crowd of people", + "a stadium full of people", + "macro photography of a drop of water", + "macro photography of a leaf", + "macro photography of a spider", + "flames", + "a robot", + "the stars at night", + "a lovely sunset", + "an oil painting", +] + + +@stub.local_entrypoint() +def main( + prompt: str, + size: str = "fhd", + upscale: bool = False, + model_weights: str = "opendalle", + n_images: int = 1, + n_steps: int = 50, + seed=None, +): + from imaginairy.enhancers.prompt_expansion import expand_prompts + from imaginairy.schema import ImaginePrompt + from imaginairy.utils import get_next_filenumber + from imaginairy.utils.log_utils import configure_logging + + configure_logging() + + prompt_texts = expand_prompts( + n=n_images, + prompt_text=prompt, + ) + # model_weights = ModelWeightsConfig( + # name="ProteusV0.4", + # aliases=["proteusv4"], + # architecture="sdxl", + # defaults={ + # "negative_prompt": DEFAULT_NEGATIVE_PROMPT, + # "composition_strength": 0.6, + # }, + # weights_location="https://huggingface.co/dataautogpt3/ProteusV0.4/tree/0dfa4101db540e7a4b2b6ba6f87d8d7219e84513", + # ) + + imagine_prompts = [ + ImaginePrompt( + prompt_text, + steps=n_steps, + size=size, + upscale=upscale, + model_weights=model_weights, + seed=seed, + ) + for prompt_text in prompt_texts + ] + # imagine_prompts = [] + # for sp in standard_prompts: + # imagine_prompts.append( + # ImaginePrompt( + # sp, + # steps=n_steps, + # size=size, + # upscale=upscale, + # model_weights=model_weights, + # seed=seed + # ) + # ) + # imagine_prompts = imagine_prompts[:n_images] + + outdir = Path("./outputs/modal-inference") + outdir.mkdir(exist_ok=True, parents=True) + file_num = get_next_filenumber(f"{outdir}/generated") + + for result in generate_image.map(imagine_prompts): + from imaginairy.api.generate import save_image_result + + save_image_result( + result=result, + base_count=file_num, + outdir=outdir, + output_file_extension="jpg", + primary_filename_type="generated", + ) + file_num += 1 diff --git a/tox.ini b/tox.ini index 80e180f..0b0dc0f 100644 --- a/tox.ini +++ b/tox.ini @@ -1,6 +1,6 @@ [pytest] addopts = --doctest-modules -s --tb=native -v --durations=10 -norecursedirs = build dist downloads other prolly_delete imaginairy/vendored +norecursedirs = build dist downloads other prolly_delete imaginairy/vendored scripts filterwarnings = ignore::DeprecationWarning ignore::UserWarning