refactor

models.py

@@ -29,7 +29,6 @@ def flush():
 
 class ControlNetPipeline:
     def __init__(self):
-        print(torch.__version__)
         self.in_use = False
         self.controlnet = ControlNetModel.from_pretrained(
         "BertChristiaens/controlnet-seg-room", torch_dtype=torch.float16)
@@ -43,7 +42,6 @@ class ControlNetPipeline:
 
         self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
         self.pipe.enable_xformers_memory_efficient_attention()
-        # self.pipe.enable_attention_slicing("max")
         self.pipe = self.pipe.to("cuda")
         
         self.waiting_queue = []
@@ -72,21 +70,44 @@ class ControlNetPipeline:
         self.waiting_queue.pop(0)
         flush()
         return results
+    
+class SDPipeline:
+    def __init__(self):
+        self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2-inpainting",
+            torch_dtype=torch.float16,
+            safety_checker=None,
+        )
 
+        self.pipe.enable_xformers_memory_efficient_attention()
+        self.pipe = self.pipe.to("cuda")
+        
+        self.waiting_queue = []
+        self.count = 0
+    
+    @property
+    def queue_size(self):
+        return len(self.waiting_queue)
+    
+    def __call__(self, **kwargs):
+        self.count += 1
+        number = self.count
 
-@contextmanager
-def catchtime(message: str) -> float:
-    """Context manager to measure time
-    Args:
-        message (str): message to log
-    Returns:
-        float: time in seconds
-    Yields:
-        Iterator[float]: time in seconds
-    """
-    start = perf_counter()
-    yield lambda: perf_counter() - start
-    LOGGING.info('%s: %.3f seconds', message, perf_counter() - start)
+        self.waiting_queue.append(number)
+        
+        # wait until the next number in the queue is the current number
+        while self.waiting_queue[0] != number:
+            print(f"Wait for your turn {number} in queue {self.waiting_queue}")
+            time.sleep(0.5)
+            pass
+
+        # it's your turn, so remove the number from the queue
+        # and call the function
+        print("It's the turn of", self.count)
+        results = self.pipe(**kwargs)
+        self.waiting_queue.pop(0)
+        flush()
+        return results
 
 
 def convolution(mask: Image.Image, size=9) -> Image:
@@ -154,53 +175,10 @@ def get_inpainting_pipeline() -> StableDiffusionInpaintPipeline:
     Returns:
         StableDiffusionInpaintPipeline: inpainting pipeline
     """
-    pipe = StableDiffusionInpaintPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-inpainting",
-        torch_dtype=torch.float16,
-        safety_checker=None,
-    )
-
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to("cuda")
-
+    pipe = SDPipeline()
     return pipe
 
 
-def make_grid_parameters(grid_search: Dict, params: Dict) -> List[Dict]:
-    """Method to make grid parameters
-    Args:
-        grid_search (Dict): grid search parameters
-        params (Dict): fixed parameters
-    Returns:
-        List[Dict]: grid parameters
-    """
-    options = []
-
-    for k in range(len(grid_search['generator'])):
-        for i in range(len(grid_search['strength'])):
-            for j in range(len(grid_search['guidance_scale'])):
-                options.append({'strength': grid_search['strength'][i],
-                                'guidance_scale': grid_search['guidance_scale'][j],
-                                'generator': grid_search['generator'][k],
-                                **params
-                                })
-    return options
-
-
-def make_captions(options: List[Dict]) -> List[str]:
-    """Method to make captions
-    Args:
-        options (List[Dict]): grid parameters
-    Returns:
-        List[str]: captions
-    """
-    captions = []
-    for option in options:
-        captions.append(
-            f"strength {option['strength']}, guidance {option['guidance_scale']}, steps {option['num_inference_steps']}")
-    return captions
-
-
 @torch.inference_mode()
 def make_image_controlnet(image: np.ndarray,
                           mask_image: np.ndarray,
@@ -219,49 +197,30 @@ def make_image_controlnet(image: np.ndarray,
         List[Image.Image]: list of generated images
     """
 
-    with catchtime("get controlnet"):
-        pipe = get_controlnet()
-
-    torch.cuda.empty_cache()
-    images = []
-
-    common_parameters = {'prompt': positive_prompt,
-                        'negative_prompt': negative_prompt,
-                        'num_inference_steps': 30,
-                            'controlnet_conditioning_scale': 1.1,
-                            'controlnet_conditioning_scale_decay': 0.96,
-                            'controlnet_steps': 28,
-                        }
-
-    grid_search = {'strength': [1.00, ],
-                   'guidance_scale': [7.0],
-                   'generator': [[torch.Generator(device="cuda").manual_seed(seed+i)] for i in range(1)],
-                   }
-
-    prompt_settings = make_grid_parameters(grid_search, common_parameters)
-
+    pipe = get_controlnet()
+    flush()
 
-    mask_image = Image.fromarray((mask_image * 255).astype(np.uint8)).convert("RGB")
     image = Image.fromarray(image).convert("RGB")
     controlnet_conditioning_image = Image.fromarray(controlnet_conditioning_image).convert("RGB").filter(ImageFilter.GaussianBlur(radius = 9))
-
+    mask_image = Image.fromarray((mask_image * 255).astype(np.uint8)).convert("RGB")
     mask_image_postproc = convolution(mask_image)
 
-    with catchtime("Controlnet generation total"):
-        for _, setting in enumerate(prompt_settings):
-            st.success(f"{pipe.queue_size} images in the queue, can take up to {(pipe.queue_size)+1 * 20} seconds")
-            with catchtime("Controlnet generation"):
-                generated_image = pipe(
-                    **setting,
-                    image=image,
-                    mask_image=mask_image,
-                    controlnet_conditioning_image=controlnet_conditioning_image,
-                ).images[0]
-                generated_image = postprocess_image_masking(
-                    generated_image, image, mask_image_postproc)
-            images.append(generated_image)
 
-    return images
+    st.success(f"{pipe.queue_size} images in the queue, can take up to {(pipe.queue_size)+1 * 10} seconds")
+    generated_image = pipe(
+        prompt=positive_prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=20,
+        strength=[1.00, ],
+        guidance_scale=[7.0],
+        generator=[torch.Generator(device="cuda").manual_seed(seed)],
+        image=image,
+        mask_image=mask_image,
+        controlnet_conditioning_image=controlnet_conditioning_image,
+    ).images[0]
+    generated_image = postprocess_image_masking(generated_image, image, mask_image_postproc)
+
+    return generated_image
 
 
 @torch.inference_mode()
@@ -278,27 +237,19 @@ def make_inpainting(positive_prompt: str,
     Returns:
         List[Image.Image]: list of generated images
     """
+    pipe = get_inpainting_pipeline()
 
-    with catchtime("Get inpainting pipeline"):
-        pipe = get_inpainting_pipeline()
-
-    common_parameters = {'prompt': positive_prompt,
-                        'negative_prompt': negative_prompt,
-                        'num_inference_steps': 20,
-                        }
-
-    torch.cuda.empty_cache()
-    images = []
-    for _ in range(1):
-        with catchtime("Inpainting generation"):
-            image_ = pipe(image=image,
-                          mask_image=Image.fromarray((mask_image * 255).astype(np.uint8)),
-                          height=HEIGHT,
-                          width=WIDTH,
-                          **common_parameters
-                          ).images[0]
-        images.append(image_)
-    return images
+    flush()
+    image_ = pipe(image=image,
+                    mask_image=Image.fromarray((mask_image * 255).astype(np.uint8)),
+                    prompt=positive_prompt,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=20,
+                    height=HEIGHT,
+                    width=WIDTH,
+                    **common_parameters
+                    ).images[0]
+    return image_
 
 
 @torch.inference_mode()
@@ -316,9 +267,8 @@ def segment_image(image: Image) -> Image:
         outputs = image_segmentor(pixel_values)
 
     seg = image_processor.post_process_semantic_segmentation(
-        outputs, target_sizes=[image.size[::-1]])
-    seg = seg[0]
-    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)  # height, width, 3
+        outputs, target_sizes=[image.size[::-1]])[0]
+    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
     palette = np.array(ade_palette())
     for label, color in enumerate(palette):
         color_seg[seg == label, :] = color