brycedrennan · brycedrennan · Sep 30, 2023 · Sep 29, 2023 · Sep 30, 2023 · Sep 30, 2023
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -5,44 +5,60 @@ on:
     branches:
       - master
   workflow_dispatch:
+env:
+  PIP_DISABLE_PIP_VERSION_CHECK: 1
+  CACHE_PATHS: |
+          ~/.cache/huggingface
+          ~/.cache/clip
+          ~/.cache/imaginairy
+          ~/.cache/torch
+
 jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v3
-    - uses: actions/setup-python@v4
-      with:
-        python-version: 3.9
-        cache: pip
-        cache-dependency-path: requirements-dev.txt
-    - name: Install dependencies
-      run: |
-        python -m pip install --disable-pip-version-check wheel pip-tools
-        pip-sync requirements-dev.txt
-        python -m pip install --disable-pip-version-check --no-deps .
-    - name: Lint
-      run: |
-        echo "::add-matcher::.github/pylama_matcher.json"
-        pylama --options tox.ini
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4.5.0
+        with:
+          python-version: 3.9
+      - name: Cache dependencies
+        uses: actions/cache@v3.2.4
+        id: cache
+        with:
+          path: ${{ env.pythonLocation }}
+          key: ${{ env.pythonLocation }}-${{ hashFiles('requirements-dev.txt') }}-lint
+      - name: Install Ruff
+        if: steps.cache.outputs.cache-hit != 'true'
+        run: grep -E 'ruff==' requirements-dev.txt | xargs pip install
+      - name: Lint
+        run: |
+          echo "::add-matcher::.github/pylama_matcher.json"
+          ruff --config tests/ruff.toml .
   autoformat:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v3
-    - uses: actions/setup-python@v4
-      with:
-        python-version: 3.9
-    - name: Install dependencies
-      run: |
-        python -m pip install --disable-pip-version-check black==23.1.0 isort==5.12.0
-    - name: Autoformatter
-      run: |
-        black --diff .
-        isort --atomic --profile black --check-only .
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4.5.0
+        with:
+          python-version: 3.9
+      - name: Cache dependencies
+        uses: actions/cache@v3.2.4
+        id: cache
+        with:
+          path: ${{ env.pythonLocation }}
+          key: ${{ env.pythonLocation }}-${{ hashFiles('requirements-dev.txt') }}-autoformat
+      - name: Install Black
+        if: steps.cache.outputs.cache-hit != 'true'
+        run: grep -E 'black==' requirements-dev.txt | xargs pip install
+      - name: Lint
+        run: |
+          black --diff --fast .
   test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version: ["3.8", "3.10"]
+        subset: ["1/10", "2/10", "3/10", "4/10", "5/10", "6/10", "7/10", "8/10", "9/10", "10/10"]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
@@ -53,22 +69,46 @@ jobs:
         cache-dependency-path: requirements-dev.txt
     - name: Install dependencies
       run: |
-        python -m pip install --disable-pip-version-check -r requirements-dev.txt
-        python -m pip install --disable-pip-version-check .
+        python -m pip install torch==1.13.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+        python -m pip install -r requirements-dev.txt .
     - name: Get current date
       id: date
       run: echo "::set-output name=curmonth::$(date +'%Y-%m')"
     - name: Cache Model Files
       id: cache-model-files
-      uses: actions/cache@v3
+      uses: actions/cache/restore@v3
       with:
-        path: |
-          ~/.cache/huggingface
-          ~/.cache/clip
-          ~/.cache/imaginairy
-          ~/.cache/torch
+        path: ${{ env.CACHE_PATHS }}
         key: ${{ steps.date.outputs.curmonth }}-b
+    # Generate initial file list for all directories
+    - name: Generate initial model file list
+      run: |
+        for dir in $CACHE_PATHS; do
+          if [ -d "$dir" ]; then
+            find $dir
+          fi
+        done > initial_file_list.txt
     - name: Test with pytest
       timeout-minutes: 20
       run: |
-        pytest --durations=50 -v
+        pytest --durations=50 -v --subset ${{ matrix.subset }}
+    # Generate final file list and check for new files
+    - name: Generate final model file list
+      run: |
+        for dir in CACHE_PATHS; do
+          if [ -d "$dir" ]; then
+            find $dir
+          fi
+        done > final_file_list.txt
+        if ! diff initial_file_list.txt final_file_list.txt > /dev/null; then
+          echo "New files detected."
+          echo "new_files=true" >> $GITHUB_ENV
+        else
+          echo "No new files detected."
+        fi
+    - uses: actions/cache/save@v3
+      id: cache
+      if: env.new_files == 'true'
+      with:
+        path: ${{ env.CACHE_PATHS }}
+        key: ${{ steps.date.outputs.curmonth }}-b
diff --git a/Makefile b/Makefile
@@ -28,18 +28,16 @@ init: require_pyenv  ## Setup a dev environment for local development.
 
 af: autoformat  ## Alias for `autoformat`
 autoformat:  ## Run the autoformatter.
-	@pycln . --all --quiet --extend-exclude __init__\.py
-	@# ERA,T201
-	@-ruff --extend-ignore ANN,ARG001,C90,DTZ,D100,D101,D102,D103,D202,D203,D212,D415,E501,RET504,S101,UP006,UP007 --extend-select C,D400,I,W --unfixable T,ERA --fix-only .
+	@-ruff check --config tests/ruff.toml . --fix-only
 	@black .
-	@isort --atomic --profile black --skip downloads/** .
+
 
 test:  ## Run the tests.
 	@pytest
 	@echo -e "The tests pass! ✨ 🍰 ✨"
 
 lint:  ## Run the code linter.
-	@pylama
+	@ruff check --config tests/ruff.toml .
 	@echo -e "No linting errors - well done! ✨ 🍰 ✨"
 
 deploy:  ## Deploy the package to pypi.org

diff --git a/docs/examples/immortal_pearl_earring.py b/docs/examples/immortal_pearl_earring.py
@@ -63,7 +63,7 @@ def generate_image_morph_video():
         if os.path.exists(filename):
             continue
 
-        result = list(imagine([prompt]))[0]
+        result = next(iter(imagine([prompt])))
         generated_image = result.images["generated"]
 
         draw = ImageDraw.Draw(generated_image)

diff --git a/imaginairy/animations.py b/imaginairy/animations.py
@@ -33,7 +33,7 @@ def make_bounce_animation(
 
     middle_imgs = shrink_list(middle_imgs, max_frames)
 
-    frames = [first_img] + middle_imgs + [last_img] + list(reversed(middle_imgs))
+    frames = [first_img, *middle_imgs, last_img, *list(reversed(middle_imgs))]
 
     # convert from latents
     converted_frames = []

diff --git a/imaginairy/api.py b/imaginairy/api.py
@@ -92,13 +92,13 @@ def _record_step(img, description, image_count, step_count, prompt):
             os.makedirs(subpath, exist_ok=True)
             filepath = os.path.join(subpath, f"{basefilename}.gif")
 
-            frames = result.progress_latents + [result.images["generated"]]
+            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 = [resized_init_image, *frames]
             frames.reverse()
             make_bounce_animation(
                 imgs=frames,
@@ -170,7 +170,7 @@ def imagine(
             logger.info(
                 f"🖼  Generating  {i + 1}/{num_prompts}: {prompt.prompt_description()}"
             )
-            for attempt in range(0, unsafe_retry_count + 1):
+            for attempt in range(unsafe_retry_count + 1):
                 if attempt > 0 and isinstance(prompt.seed, int):
                     prompt.seed += 100_000_000 + attempt
                 result = _generate_single_image(
@@ -238,7 +238,7 @@ def _generate_single_image(
     latent_channels = 4
     downsampling_factor = 8
     batch_size = 1
-    global _most_recent_result  # noqa
+    global _most_recent_result
     # handle prompt pulling in previous values
     # if isinstance(prompt.init_image, str) and prompt.init_image.startswith("*prev"):
     #     _, img_type = prompt.init_image.strip("*").split(".")
@@ -457,16 +457,17 @@ def latent_logger(latents):
                 if control_image_t.shape[1] != 3:
                     raise RuntimeError("Control image must have 3 channels")
 
-                if control_input.mode != "inpaint":
-                    if control_image_t.min() < 0 or control_image_t.max() > 1:
-                        raise RuntimeError(
-                            f"Control image must be in [0, 1] but we received {control_image_t.min()} and {control_image_t.max()}"
-                        )
+                if (
+                    control_input.mode != "inpaint"
+                    and control_image_t.min() < 0
+                    or control_image_t.max() > 1
+                ):
+                    msg = f"Control image must be in [0, 1] but we received {control_image_t.min()} and {control_image_t.max()}"
+                    raise RuntimeError(msg)
 
                 if control_image_t.max() == control_image_t.min():
-                    raise RuntimeError(
-                        f"No control signal found in control image {control_input.mode}."
-                    )
+                    msg = f"No control signal found in control image {control_input.mode}."
+                    raise RuntimeError(msg)
 
                 c_cat.append(control_image_t)
                 control_strengths.append(control_input.strength)
@@ -517,7 +518,7 @@ def latent_logger(latents):
         if (
             prompt.allow_compose_phase
             and not is_controlnet_model
-            and not model.cond_stage_key == "edit"
+            and model.cond_stage_key != "edit"
         ):
             if prompt.init_image:
                 comp_image = _generate_composition_image(

diff --git a/imaginairy/cli/clickshell_mod.py b/imaginairy/cli/clickshell_mod.py
@@ -4,6 +4,7 @@
 import shlex
 import traceback
 from functools import update_wrapper
+from typing import ClassVar
 
 import click
 from click_help_colors import HelpColorsCommand, HelpColorsMixin
@@ -43,27 +44,23 @@ def invoke_(self, arg):  # pylint: disable=unused-argument
             # and that's not ideal when running in a shell.
             pass
         except Exception as e:  # noqa
-            traceback.print_exception(e)  # noqa
+            traceback.print_exception(e)
             # logger.warning(traceback.format_exc())
 
         # Always return False so the shell doesn't exit
         return False
 
     invoke_ = update_wrapper(invoke_, command.callback)
-    invoke_.__name__ = "do_%s" % command.name  # noqa
+    invoke_.__name__ = "do_%s" % command.name
     return invoke_
 
 
 class ModClickShell(ClickShell):
     def add_command(self, cmd, name):
         # Use the MethodType to add these as bound methods to our current instance
-        setattr(
-            self, "do_%s" % name, get_method_type(mod_get_invoke(cmd), self)  # noqa
-        )
-        setattr(self, "help_%s" % name, get_method_type(get_help(cmd), self))  # noqa
-        setattr(
-            self, "complete_%s" % name, get_method_type(get_complete(cmd), self)  # noqa
-        )
+        setattr(self, "do_%s" % name, get_method_type(mod_get_invoke(cmd), self))
+        setattr(self, "help_%s" % name, get_method_type(get_help(cmd), self))
+        setattr(self, "complete_%s" % name, get_method_type(get_complete(cmd), self))
 
 
 class ModShell(Shell):
@@ -85,7 +82,7 @@ class ColorShell(HelpColorsMixin, ModShell):
 
 
 class ImagineColorsCommand(HelpColorsCommand):
-    _option_order = []
+    _option_order: ClassVar = []
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

diff --git a/imaginairy/cli/edit.py b/imaginairy/cli/edit.py
@@ -43,7 +43,7 @@
 )
 @add_options(edit_options)
 @click.pass_context
-def edit_cmd(  # noqa
+def edit_cmd(
     ctx,
     image_paths,
     image_strength,
@@ -77,7 +77,7 @@ def edit_cmd(  # noqa
     model_weights_path,
     model_config_path,
     prompt_library_path,
-    version,  # noqa
+    version,
     make_gif,
     make_compare_gif,
     arg_schedules,
@@ -130,7 +130,7 @@ def edit_cmd(  # noqa
         model_weights_path,
         model_config_path,
         prompt_library_path,
-        version,  # noqa
+        version,
         make_gif,
         make_compare_gif,
         arg_schedules,

diff --git a/imaginairy/cli/imagine.py b/imaginairy/cli/imagine.py
@@ -90,7 +90,7 @@ def imagine_cmd(
     model_weights_path,
     model_config_path,
     prompt_library_path,
-    version,  # noqa
+    version,
     make_gif,
     make_compare_gif,
     arg_schedules,
@@ -110,7 +110,7 @@ def imagine_cmd(
     # hacky method of getting order of control images (mixing raw and normal images)
     control_images = [
         (o, path)
-        for o, path in ImagineColorsCommand._option_order  # noqa
+        for o, path in ImagineColorsCommand._option_order
         if o.name in ("control_image", "control_image_raw")
     ]
     control_inputs = []
@@ -176,7 +176,7 @@ def imagine_cmd(
         model_weights_path,
         model_config_path,
         prompt_library_path,
-        version,  # noqa
+        version,
         make_gif,
         make_compare_gif,
         arg_schedules,
@@ -187,4 +187,4 @@ def imagine_cmd(
 
 
 if __name__ == "__main__":
-    imagine_cmd()  # noqa
+    imagine_cmd()
diff --git a/imaginairy/cli/main.py b/imaginairy/cli/main.py
@@ -92,4 +92,4 @@ def model_list_cmd():
 
 
 if __name__ == "__main__":
-    aimg()  # noqa
+    aimg()