fix: tile mode was broken since latest perf improvements (#220)

.github/workflows/ci.yaml

@@ -69,6 +69,6 @@ jobs:
           ~/.cache/torch
         key: ${{ steps.date.outputs.curmonth }}-b
     - name: Test with pytest
-      timeout-minutes: 10
+      timeout-minutes: 20
       run: |
         pytest --durations=50 -v

README.md

@@ -284,6 +284,7 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface -
 ## ChangeLog
 
 - feature: `openjourney-v1` and `openjourney-v2` models added. available via `--model openjourney-v2`
+- fix: tile mode was broken since latest perf improvements
 
 **8.2.0**
 - feature: added `aimg system-info` command to help debug issues

imaginairy/modules/diffusion/ddpm.py

@@ -715,9 +715,8 @@ def _TileModeConv2DConvForward(
     self, input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor  # noqa
 ):
     if self.padding_modeX == self.padding_modeY:
-        return F.conv2d(
-            input, weight, bias, self.stride, self.padding, self.dilation, self.groups
-        )
+        self.padding_mode = self.padding_modeX
+        return self._orig_conv_forward(input, weight, bias)  # noqa
 
     w1 = F.pad(input, self.paddingX, mode=self.padding_modeX)
     del input
@@ -790,6 +789,7 @@ class LatentDiffusion(DDPM):
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
                 m._initial_padding_mode = m.padding_mode
+                m._orig_conv_forward = m._conv_forward
                 m._conv_forward = _TileModeConv2DConvForward.__get__(  # noqa
                     m, nn.Conv2d
                 )

tests/test_api.py

@@ -258,3 +258,20 @@ def test_cliptext_inpainting_pearl_doctor(
     pillow_fit_image_within(img).save(f"{filename_base_for_orig_outputs}_orig.jpg")
     img_path = f"{filename_base_for_outputs}.png"
     assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=2800)
+
+
+@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
+def test_tile_mode(filename_base_for_outputs):
+    prompt_text = "gold coins"
+    prompt = ImaginePrompt(
+        prompt_text,
+        width=400,
+        height=400,
+        steps=5,
+        seed=1,
+        tile_mode="xy",
+    )
+    result = next(imagine(prompt))
+
+    img_path = f"{filename_base_for_outputs}.png"
+    assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=1000)