add tests

rudalle/__init__.py

@@ -4,7 +4,6 @@ from .dalle import get_rudalle_model
 from .tokenizer import get_tokenizer
 from .realesrgan import get_realesrgan
 from .ruclip import get_ruclip
-from .image_prompts import ImagePrompts
 from . import vae, dalle, tokenizer, realesrgan, pipelines, ruclip, image_prompts
 
 
@@ -20,7 +19,8 @@ __all__ = [
     'tokenizer',
     'realesrgan',
     'pipelines',
-    'ImagePrompts',
+    'image_prompts',
 ]
 
 __version__ = '0.0.1-rc1'
+

tests/test_image_prompts.py

@@ -0,0 +1,20 @@
+import pytest
+
+from rudalle.image_prompts import ImagePrompts
+
+@pytest.mark.parametrize('borders, crop_first', [
+    ({'up': 4, 'right': 0, 'left': 0, 'down': 0}, False),
+    ({'up': 4, 'right': 0, 'left': 0, 'down': 0}, True),
+    ({'up': 4, 'right': 3, 'left': 3, 'down': 3}, False)
+])
+def test_image_prompts(sample_image, vae, borders, crop_first):
+    img = sample_image.copy()
+    img = img.resize((256, 256))
+    image_prompt = ImagePrompts(img, borders, vae, crop_first=crop_first)
+    if crop_first:
+        assert image_prompt.image_prompts.shape[1] == borders['up'] * 32
+        assert len(image_prompt.image_prompts_idx) == borders['up'] * 32
+    else:
+        assert image_prompt.image_prompts.shape[1] == 32 * 32
+        assert len(image_prompt.image_prompts_idx) == (borders['up'] + borders['down']) * 32 \
+               + (borders['left'] + borders['right']) * (32 - borders['up'] - borders['down'])