add adaptive model

README.md

@@ -24,4 +24,5 @@ a handsome man relaxing in a chair, shirt widely unbuttoned, eyes closed, close
 - [ ] добавить caption.csv в data/
 - [ ] добавить adaptive resize
 - [ ] прокинуть параметры в демке (seed и т д), + adaptive resize размеры
+- [ ] разобраться с выставлением device в пайплайн
 - [ ] настроить запуск в hugging space

src/pipeline.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Tuple
 
 from PIL import Image
 import numpy as np
@@ -49,8 +50,11 @@ class FashionPipeline:
 
         # extract segmentation mask
         segm_mask = generate_mask(control_image, self.segmentation_model, device=self.device)
+        control_mask = self.create_control_image(segm_mask)
+
         control_mask = self.adaptive_resize(
-            self.create_control_image(segm_mask),
+            image=control_mask,
+            initial_shape=(control_image.shape[1], control_image.shape[0]),
             target_image_size=resolution,
         )
 
@@ -84,13 +88,20 @@ class FashionPipeline:
         ch3 = (segm_mask == 3) * 255  # Full body(blue).
         return Image.fromarray(np.stack([ch1, ch2, ch3], axis=-1).astype('uint8'), 'RGB')
 
-    def adaptive_resize(self, image, target_image_size=512, max_image_size=768, divisible=64):
+    def adaptive_resize(
+        self,
+        image: Image,
+        initial_shape: Tuple[int, int],
+        target_image_size: int = 512,
+        max_image_size: int = 768,
+        divisible: int = 64,
+    ):
 
         assert target_image_size % divisible == 0
         assert max_image_size % divisible == 0
         assert max_image_size >= target_image_size
 
-        width, height = image.size
+        width, height = initial_shape
         aspect_ratio = width / height
 
         if height > width:
@@ -116,14 +127,14 @@ class FashionPipeline:
         self.controlnet = ControlNetModel.from_pretrained(
             self.config.controlnet_path,
             torch_dtype=torch.float16,
-            device_map="auto",
+            # device_map="auto",
         )
 
         self.pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
             self.config.base_model_path,
             controlnet=self.controlnet,
             torch_dtype=torch.float16,
-            device_map="auto",
+            # device_map="auto",
         )
 
         self.pipeline.scheduler = UniPCMultistepScheduler.from_config(self.pipeline.scheduler.config)