fix crop in Cond2ImPipe

singnet · Jan 15, 2025 · 0565eb0 · 0565eb0
1 parent c6580b1
commit 0565eb0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/multigen/__init__.py b/multigen/__init__.py
@@ -1,4 +1,4 @@
-from .pipes import Prompt2ImPipe, Im2ImPipe, MaskedIm2ImPipe, CIm2ImPipe
+from .pipes import Prompt2ImPipe, Im2ImPipe, MaskedIm2ImPipe, CIm2ImPipe, Cond2ImPipe
 from .sessions import GenSession
 from .prompting import Cfgen
 from .loader import Loader

diff --git a/multigen/pipes.py b/multigen/pipes.py
@@ -862,6 +862,7 @@ def setup(self, fimage, width=None, height=None,
         self.fname = fimage
         image = Image.open(fimage).convert("RGB") if image is None else image
         self._original_size = image.size
+        self._use_input_size = width is None or height is None
         image = util.pad_image_to_multiple_of_8(image)
         self._condition_image = [image]
         self._input_image = [image]
@@ -914,7 +915,8 @@ def gen(self, inputs):
         inputs.update({"image": self._input_image,
                        "control_image": self._condition_image})
         image = self.pipe(**inputs).images[0]
-        result = image.crop((0, 0, self._original_size[0], self._original_size[1]))
+        result = image.crop((0, 0, self._original_size[0] if self._use_input_size else inputs.get('height'),
+                                   self._original_size[1] if self._use_input_size else inputs.get('width') ))
         return result