tests: fix tests

- disable details mode. needs more work done to support

.github/workflows/ci.yaml

@@ -54,7 +54,7 @@ jobs:
         run: |
           black --diff --fast .
   test:
-    runs-on: macos-13-xlarge
+    runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:

docs/todo.md

@@ -1,5 +1,20 @@
 ## Todo
 
+### v14 todo
+
+ - configurable composition cutoff
+ - rename model parameter weights 
+ - rename model_config parameter to architecture and make it case insensitive
+ - add --size parameter that accepts strings (e.g. 256x256, 4k, uhd, 8k, etc)
+ - detect if cuda torch missing and give better error message
+ - add method to install correct torch version
+ - add composition cutoff parameter
+ - allow selection of output video format
+ - chain multiple operations together imggen => videogen
+ - make sure terminal output on windows doesn't suck
+
+### Old Todo
+
  - Inference Performance Optimizations
    - ✅ fp16
    - ✅ [Doggettx Sliced attention](https://github.com/CompVis/stable-diffusion/compare/main...Doggettx:stable-diffusion:autocast-improvements#)

imaginairy/api_refiners.py

@@ -124,9 +124,11 @@ def _generate_single_image(
         ]
 
         init_latent = None
+        noise_step = None
         if prompt.init_image:
             starting_image = prompt.init_image
-            first_step = int((prompt.steps - 1) * prompt.init_image_strength)
+            first_step = int((prompt.steps) * prompt.init_image_strength)
+            # noise_step = int((prompt.steps - 1) * prompt.init_image_strength)
 
             if prompt.mask_prompt:
                 mask_image, mask_grayscale = get_img_mask(
@@ -201,6 +203,9 @@ def _generate_single_image(
                     max_height=prompt.height,
                     max_width=prompt.width,
                 )
+                if control_input.mode == "inpaint":
+                    control_image_input = ImageOps.invert(control_image_input)
+
                 control_image_input_t = pillow_img_to_torch_image(control_image_input)
                 control_image_input_t = control_image_input_t.to(get_device())
 
@@ -251,7 +256,6 @@ def _generate_single_image(
                 )
                 controlnets.append((controlnet, control_image_t))
 
-        noise_step = None
         if prompt.allow_compose_phase:
             if prompt.init_image:
                 comp_image, comp_img_orig = _generate_composition_image(
@@ -272,7 +276,8 @@ def _generate_single_image(
                 result_images["composition-upscaled"] = comp_image
                 # noise = noise[:, :, : comp_image.height, : comp_image.shape[3]]
                 comp_cutoff = 0.60
-                first_step = int((prompt.steps - 1) * comp_cutoff)
+                first_step = int((prompt.steps) * comp_cutoff)
+                noise_step = int((prompt.steps - 1) * comp_cutoff)
                 # noise_step = int(prompt.steps * max(comp_cutoff - 0.05, 0))
                 # noise_step = max(noise_step, 0)
                 # noise_step = min(noise_step, prompt.steps - 1)
@@ -303,9 +308,12 @@ def _generate_single_image(
 
         if init_latent is not None:
             noise_step = noise_step if noise_step is not None else first_step
-            noised_latent = sd.scheduler.add_noise(
-                x=init_latent, noise=noise, step=sd.steps[noise_step]
-            )
+            if first_step >= len(sd.steps):
+                noised_latent = init_latent
+            else:
+                noised_latent = sd.scheduler.add_noise(
+                    x=init_latent, noise=noise, step=sd.steps[noise_step]
+                )
 
         x = noised_latent
         x = x.to(device=sd.device, dtype=sd.dtype)

imaginairy/img_processors/control_modes.py

@@ -243,6 +243,6 @@ CONTROL_MODES = {
     "shuffle": shuffle_map_torch,
     "edit": noop,
     "inpaint": inpaint_prep,
-    "details": noop,
+    # "details": noop,
     "colorize": to_grayscale,
 }

imaginairy/samplers/__init__.py

@@ -1,21 +1,20 @@
 from imaginairy.samplers import kdiff
 from imaginairy.samplers.base import SamplerName  # noqa
 from imaginairy.samplers.ddim import DDIMSampler
-from imaginairy.samplers.plms import PLMSSampler
 
 SAMPLERS = [
-    PLMSSampler,
+    # PLMSSampler,
     DDIMSampler,
-    kdiff.DPMFastSampler,
-    kdiff.DPMAdaptiveSampler,
-    kdiff.LMSSampler,
-    kdiff.DPM2Sampler,
-    kdiff.DPM2AncestralSampler,
+    # kdiff.DPMFastSampler,
+    # kdiff.DPMAdaptiveSampler,
+    # kdiff.LMSSampler,
+    # kdiff.DPM2Sampler,
+    # kdiff.DPM2AncestralSampler,
     kdiff.DPMPP2MSampler,
-    kdiff.DPMPP2SAncestralSampler,
-    kdiff.EulerSampler,
-    kdiff.EulerAncestralSampler,
-    kdiff.HeunSampler,
+    # kdiff.DPMPP2SAncestralSampler,
+    # kdiff.EulerSampler,
+    # kdiff.EulerAncestralSampler,
+    # kdiff.HeunSampler,
 ]
 
 SAMPLER_LOOKUP = {sampler.short_name: sampler for sampler in SAMPLERS}

imaginairy/samplers/kdiff.py

@@ -151,32 +151,34 @@ class KDiffusionSampler(ImageSampler, ABC):
         return samples
 
 
-class DPMFastSampler(KDiffusionSampler):
-    short_name = SamplerName.K_DPM_FAST
-    name = "Diffusion probabilistic models - fast"
-    default_steps = 15
-    sampler_func = staticmethod(sample_dpm_fast)
-
-
-class DPMAdaptiveSampler(KDiffusionSampler):
-    short_name = SamplerName.K_DPM_ADAPTIVE
-    name = "Diffusion probabilistic models - adaptive"
-    default_steps = 40
-    sampler_func = staticmethod(sample_dpm_adaptive)
-
-
-class DPM2Sampler(KDiffusionSampler):
-    short_name = SamplerName.K_DPM_2
-    name = "Diffusion probabilistic models - 2"
-    default_steps = 40
-    sampler_func = staticmethod(k_sampling.sample_dpm_2)
-
-
-class DPM2AncestralSampler(KDiffusionSampler):
-    short_name = SamplerName.K_DPM_2_ANCESTRAL
-    name = "Diffusion probabilistic models - 2 ancestral"
-    default_steps = 40
-    sampler_func = staticmethod(k_sampling.sample_dpm_2_ancestral)
+#
+# class DPMFastSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_DPM_FAST
+#     name = "Diffusion probabilistic models - fast"
+#     default_steps = 15
+#     sampler_func = staticmethod(sample_dpm_fast)
+#
+#
+# class DPMAdaptiveSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_DPM_ADAPTIVE
+#     name = "Diffusion probabilistic models - adaptive"
+#     default_steps = 40
+#     sampler_func = staticmethod(sample_dpm_adaptive)
+#
+#
+# class DPM2Sampler(KDiffusionSampler):
+#     short_name = SamplerName.K_DPM_2
+#     name = "Diffusion probabilistic models - 2"
+#     default_steps = 40
+#     sampler_func = staticmethod(k_sampling.sample_dpm_2)
+#
+#
+# class DPM2AncestralSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_DPM_2_ANCESTRAL
+#     name = "Diffusion probabilistic models - 2 ancestral"
+#     default_steps = 40
+#     sampler_func = staticmethod(k_sampling.sample_dpm_2_ancestral)
+#
 
 
 class DPMPP2MSampler(KDiffusionSampler):
@@ -186,39 +188,40 @@ class DPMPP2MSampler(KDiffusionSampler):
     sampler_func = staticmethod(k_sampling.sample_dpmpp_2m)
 
 
-class DPMPP2SAncestralSampler(KDiffusionSampler):
-    short_name = SamplerName.K_DPMPP_2S_ANCESTRAL
-    name = "Ancestral sampling with DPM-Solver++(2S) second-order steps."
-    default_steps = 15
-    sampler_func = staticmethod(k_sampling.sample_dpmpp_2s_ancestral)
-
-
-class EulerSampler(KDiffusionSampler):
-    short_name = SamplerName.K_EULER
-    name = "Algorithm 2 (Euler steps) from Karras et al. (2022)"
-    default_steps = 40
-    sampler_func = staticmethod(k_sampling.sample_euler)
-
-
-class EulerAncestralSampler(KDiffusionSampler):
-    short_name = SamplerName.K_EULER_ANCESTRAL
-    name = "Euler ancestral"
-    default_steps = 40
-    sampler_func = staticmethod(k_sampling.sample_euler_ancestral)
-
-
-class HeunSampler(KDiffusionSampler):
-    short_name = SamplerName.K_HEUN
-    name = "Algorithm 2 (Heun steps) from Karras et al. (2022)."
-    default_steps = 40
-    sampler_func = staticmethod(k_sampling.sample_heun)
-
-
-class LMSSampler(KDiffusionSampler):
-    short_name = SamplerName.K_LMS
-    name = "LMS"
-    default_steps = 40
-    sampler_func = staticmethod(k_sampling.sample_lms)
+#
+# class DPMPP2SAncestralSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_DPMPP_2S_ANCESTRAL
+#     name = "Ancestral sampling with DPM-Solver++(2S) second-order steps."
+#     default_steps = 15
+#     sampler_func = staticmethod(k_sampling.sample_dpmpp_2s_ancestral)
+#
+#
+# class EulerSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_EULER
+#     name = "Algorithm 2 (Euler steps) from Karras et al. (2022)"
+#     default_steps = 40
+#     sampler_func = staticmethod(k_sampling.sample_euler)
+#
+#
+# class EulerAncestralSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_EULER_ANCESTRAL
+#     name = "Euler ancestral"
+#     default_steps = 40
+#     sampler_func = staticmethod(k_sampling.sample_euler_ancestral)
+#
+#
+# class HeunSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_HEUN
+#     name = "Algorithm 2 (Heun steps) from Karras et al. (2022)."
+#     default_steps = 40
+#     sampler_func = staticmethod(k_sampling.sample_heun)
+#
+#
+# class LMSSampler(KDiffusionSampler):
+#     short_name = SamplerName.K_LMS
+#     name = "LMS"
+#     default_steps = 40
+#     sampler_func = staticmethod(k_sampling.sample_lms)
 
 
 class CFGDenoiser(nn.Module):

requirements-dev.txt

@@ -4,7 +4,7 @@
 #
 #    pip-compile --output-file=requirements-dev.txt requirements-dev.in setup.py
 #
-aiohttp==3.9.0
+aiohttp==3.9.1
     # via fsspec
 aiosignal==1.3.1
     # via aiohttp
@@ -65,7 +65,7 @@ filelock==3.13.1
     #   transformers
 filterpy==1.4.5
     # via facexlib
-fonttools==4.45.0
+fonttools==4.45.1
     # via matplotlib
 frozenlist==1.4.0
     # via
@@ -89,7 +89,7 @@ huggingface-hub==0.19.4
     #   timm
     #   tokenizers
     #   transformers
-idna==3.4
+idna==3.6
     # via
     #   anyio
     #   requests
@@ -100,7 +100,7 @@ importlib-metadata==6.8.0
     # via diffusers
 iniconfig==2.0.0
     # via pytest
-jaxtyping==0.2.23
+jaxtyping==0.2.24
     # via refiners
 jinja2==3.1.2
     # via torch
@@ -266,7 +266,7 @@ sympy==1.12
     # via torch
 termcolor==2.3.0
     # via pytest-sugar
-timm==0.9.11
+timm==0.9.12
     # via
     #   imaginAIry (setup.py)
     #   open-clip-torch
@@ -333,7 +333,7 @@ uvicorn==0.24.0.post1
     # via imaginAIry (setup.py)
 wcwidth==0.2.12
     # via ftfy
-wheel==0.41.3
+wheel==0.42.0
     # via -r requirements-dev.in
 yarl==1.9.3
     # via aiohttp

setup.py

@@ -111,7 +111,7 @@ setup(
         "triton>=2.0.0; sys_platform!='darwin' and platform_machine!='aarch64'",
         "kornia>=0.6",
         "uvicorn>=0.16.0",
-        "xformers>=0.0.16; sys_platform!='darwin' and platform_machine!='aarch64'",
+        # "xformers>=0.0.22; sys_platform!='darwin' and platform_machine!='aarch64'",
     ],
     # don't specify maximum python versions as it can cause very long dependency resolution issues as the resolver
     # goes back to older versions of packages that didn't specify a maximum

tests/test_api.py

@@ -315,19 +315,33 @@ control_modes = list(CONTROL_MODES.keys())
 @pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
 def test_controlnet(filename_base_for_outputs, control_mode):
     prompt_text = "a photo of a woman sitting on a bench"
+    img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2.png")
     control_input = ControlNetInput(
         mode=control_mode,
-        image=LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2.png"),
+        image=img,
     )
 
+    seed = 0
+    if control_mode == "inpaint":
+        prompt_text = "a wise old man"
+        seed = 1
+        mask_image = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2_mask.png")
+        control_input = ControlNetInput(
+            mode=control_mode,
+            image=mask_image,
+        )
+
     prompt = ImaginePrompt(
         prompt_text,
         width=512,
         height=512,
-        steps=15,
-        seed=0,
+        steps=45,
+        seed=seed,
+        init_image=img,
+        init_image_strength=0,
         control_inputs=[control_input],
         fix_faces=True,
+        sampler="ddim",
     )
     prompt.steps = 1
     prompt.width = 256

tests/test_cmds.py

@@ -31,7 +31,7 @@ def test_imagine_cmd(monkeypatch):
             f"{TESTS_FOLDER}/test_output",
         ],
     )
-    assert result.exit_code == 0
+    assert result.exit_code == 0, (result.stdout, result.stderr)
 
 
 def test_edit_cmd(monkeypatch):
@@ -45,13 +45,13 @@ def test_edit_cmd(monkeypatch):
             "1",
             "-p",
             "turn the dog into a cat",
-            "--model",
-            "empty",
+            # "--model",
+            # "empty",
             "--outdir",
             f"{TESTS_FOLDER}/test_output",
         ],
     )
-    assert result.exit_code == 0
+    assert result.exit_code == 0, (result.stdout, result.stderr)
 
 
 def test_aimg_shell():
@@ -74,7 +74,7 @@ def test_edit_demo(monkeypatch):
                 steps=1,
                 width=256,
                 height=256,
-                model="empty",
+                # model="empty",
             )
         ]
 

tests/test_schema/test_controlnetinput.py

@@ -35,5 +35,6 @@ def test_controlnetinput_filepath_input(lazy_img):
 
 
 def test_controlnetinput_big(lazy_img):
-    with pytest.raises(ValidationError, match=r".*less than or.*"):
+    ControlNetInput(mode="canny", strength=2)
+    with pytest.raises(ValidationError, match=r".*float_type.*"):
         ControlNetInput(mode="canny", strength=2**2048)