feature: script for running imaginairy in the modal.com cloud

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,78 +92,110 @@ def _record_step(img, description, image_count, step_count, prompt):
         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
+        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,
+        )
+        if not primary_filename:
             continue
-        img_str = ""
-        if prompt.init_image:
-            img_str = f"_img2img-{prompt.init_image_strength}"
+        result_filenames.append(primary_filename)
+        if primary_filename and videogen:
+            try:
+                generate_video(
+                    input_path=primary_filename,
+                )
+            except FileNotFoundError as e:
+                logger.error(str(e))
+                exit(1)
+
+        base_count += 1
+        del result
+
+    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)}"
+    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}"
         )
-        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
-                )
-                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")
+        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 = [result.images["generated"], resized_init_image]
-
-            make_bounce_animation(
-                imgs=frames,
-                outpath=filepath,
-            )
-            image_type = "gif"
-            logger.info(f"        {image_type:<22} {filepath}")
+            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]
 
-        base_count += 1
-        del result
+        make_bounce_animation(
+            imgs=frames,
+            outpath=filepath,
+        )
+        image_type = "gif"
+        logger.info(f"        {image_type:<22} {filepath}")
 
-    return result_filenames
+    return primary_filename
 
 
 def imagine(

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