refactor

app.py

@@ -7,7 +7,8 @@ import numpy as np
 import os
 import time
 
-from models import make_image_controlnet, make_inpainting, segment_image
+from models import make_image_controlnet, make_inpainting
+from segmentation import segment_image
 from config import HEIGHT, WIDTH, POS_PROMPT, NEG_PROMPT, COLOR_MAPPING, map_colors, map_colors_rgb
 from palette import COLOR_MAPPING_CATEGORY
 from preprocessing import preprocess_seg_mask, get_image, get_mask

helpers.py

@@ -0,0 +1,46 @@
+import gc
+import torch
+from scipy.signal import fftconvolve
+from PIL import Image
+
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+
+
+def convolution(mask: Image.Image, size=9) -> Image:
+    """Method to blur the mask
+    Args:
+        mask (Image): masking image
+        size (int, optional): size of the blur. Defaults to 9.
+    Returns:
+        Image: blurred mask
+    """
+    mask = np.array(mask.convert("L"))
+    conv = np.ones((size, size)) / size**2
+    mask_blended = fftconvolve(mask, conv, 'same')
+    mask_blended = mask_blended.astype(np.uint8).copy()
+
+    border = size
+
+    # replace borders with original values
+    mask_blended[:border, :] = mask[:border, :]
+    mask_blended[-border:, :] = mask[-border:, :]
+    mask_blended[:, :border] = mask[:, :border]
+    mask_blended[:, -border:] = mask[:, -border:]
+
+    return Image.fromarray(mask_blended).convert("L")
+
+
+def postprocess_image_masking(inpainted: Image, image: Image, mask: Image) -> Image:
+    """Method to postprocess the inpainted image
+    Args:
+        inpainted (Image): inpainted image
+        image (Image): original image
+        mask (Image): mask
+    Returns:
+        Image: inpainted image
+    """
+    final_inpainted = Image.composite(inpainted.convert("RGBA"), image.convert("RGBA"), mask)
+    return final_inpainted.convert("RGB")

models.py

@@ -8,176 +8,18 @@ import gc
 import time
 import numpy as np
 from PIL import Image
-from time import perf_counter
-from contextlib import contextmanager
-from scipy.signal import fftconvolve
 from PIL import ImageFilter
 
-from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
 from diffusers import ControlNetModel, UniPCMultistepScheduler
-from diffusers import StableDiffusionInpaintPipeline
 
 from config import WIDTH, HEIGHT
 from palette import ade_palette
 from stable_diffusion_controlnet_inpaint_img2img import StableDiffusionControlNetInpaintImg2ImgPipeline
+from helpers import flush, postprocess_image_masking, convolution
+from pipelines import ControlNetPipeline, SDPipeline, get_inpainting_pipeline, get_controlnet
 
 LOGGING = logging.getLogger(__name__)
 
-def flush():
-    gc.collect()
-    torch.cuda.empty_cache()
-
-class ControlNetPipeline:
-    def __init__(self):
-        self.in_use = False
-        self.controlnet = ControlNetModel.from_pretrained(
-        "BertChristiaens/controlnet-seg-room", torch_dtype=torch.float16)
-
-        self.pipe = StableDiffusionControlNetInpaintImg2ImgPipeline.from_pretrained(
-            "runwayml/stable-diffusion-inpainting",
-            controlnet=self.controlnet,
-            safety_checker=None,
-            torch_dtype=torch.float16
-        )
-
-        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
-        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
-
-        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
-    
-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
-
-        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:
-    """Method to blur the mask
-    Args:
-        mask (Image): masking image
-        size (int, optional): size of the blur. Defaults to 9.
-    Returns:
-        Image: blurred mask
-    """
-    mask = np.array(mask.convert("L"))
-    conv = np.ones((size, size)) / size**2
-    mask_blended = fftconvolve(mask, conv, 'same')
-    mask_blended = mask_blended.astype(np.uint8).copy()
-
-    border = size
-
-    # replace borders with original values
-    mask_blended[:border, :] = mask[:border, :]
-    mask_blended[-border:, :] = mask[-border:, :]
-    mask_blended[:, :border] = mask[:, :border]
-    mask_blended[:, -border:] = mask[:, -border:]
-
-    return Image.fromarray(mask_blended).convert("L")
-
-
-def postprocess_image_masking(inpainted: Image, image: Image, mask: Image) -> Image:
-    """Method to postprocess the inpainted image
-    Args:
-        inpainted (Image): inpainted image
-        image (Image): original image
-        mask (Image): mask
-    Returns:
-        Image: inpainted image
-    """
-    final_inpainted = Image.composite(inpainted.convert("RGBA"), image.convert("RGBA"), mask)
-    return final_inpainted.convert("RGB")
-
-
-@st.experimental_singleton(max_entries=5)
-def get_controlnet() -> ControlNetModel:
-    """Method to load the controlnet model
-    Returns:
-        ControlNetModel: controlnet model
-    """
-    pipe = ControlNetPipeline()
-    return pipe
-
-
-@st.experimental_singleton(max_entries=5)
-def get_segmentation_pipeline() -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
-    """Method to load the segmentation pipeline
-    Returns:
-        Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]: segmentation pipeline
-    """
-    image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
-    image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
-        "openmmlab/upernet-convnext-small")
-    return image_processor, image_segmentor
-
-
-@st.experimental_singleton(max_entries=5)
-def get_inpainting_pipeline() -> StableDiffusionInpaintPipeline:
-    """Method to load the inpainting pipeline
-    Returns:
-        StableDiffusionInpaintPipeline: inpainting pipeline
-    """
-    pipe = SDPipeline()
-    return pipe
-
 
 @torch.inference_mode()
 def make_image_controlnet(image: np.ndarray,
@@ -238,12 +80,13 @@ def make_inpainting(positive_prompt: str,
         List[Image.Image]: list of generated images
     """
     pipe = get_inpainting_pipeline()
+    mask_image = Image.fromarray((mask_image * 255).astype(np.uint8))
     mask_image_postproc = convolution(mask_image)
 
     flush()
     st.success(f"{pipe.queue_size} images in the queue, can take up to {(pipe.queue_size+1) * 10} seconds")
     generated_image = pipe(image=image,
-                    mask_image=Image.fromarray((mask_image * 255).astype(np.uint8)),
+                    mask_image=mask_image,
                     prompt=positive_prompt,
                     negative_prompt=negative_prompt,
                     num_inference_steps=20,
@@ -252,29 +95,4 @@ def make_inpainting(positive_prompt: str,
                     ).images[0]
     generated_image = postprocess_image_masking(generated_image, image, mask_image_postproc)
 
-    return image_
-
-
-@torch.inference_mode()
-@torch.autocast('cuda')
-def segment_image(image: Image) -> Image:
-    """Method to segment image
-    Args:
-        image (Image): input image
-    Returns:
-        Image: segmented image
-    """
-    image_processor, image_segmentor = get_segmentation_pipeline()
-    pixel_values = image_processor(image, return_tensors="pt").pixel_values
-    with torch.no_grad():
-        outputs = image_segmentor(pixel_values)
-
-    seg = image_processor.post_process_semantic_segmentation(
-        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
-    color_seg = color_seg.astype(np.uint8)
-    seg_image = Image.fromarray(color_seg).convert('RGB')
-    return seg_image
+    return generated_image

pipelines.py

@@ -0,0 +1,126 @@
+import logging
+from typing import List, Tuple, Dict
+
+import streamlit as st
+import torch
+import gc
+import time
+import numpy as np
+from PIL import Image
+from time import perf_counter
+from contextlib import contextmanager
+from scipy.signal import fftconvolve
+from PIL import ImageFilter
+
+from diffusers import ControlNetModel, UniPCMultistepScheduler
+from diffusers import StableDiffusionInpaintPipeline
+
+from config import WIDTH, HEIGHT
+from stable_diffusion_controlnet_inpaint_img2img import StableDiffusionControlNetInpaintImg2ImgPipeline
+from helpers import flush
+
+LOGGING = logging.getLogger(__name__)
+
+class ControlNetPipeline:
+    def __init__(self):
+        self.in_use = False
+        self.controlnet = ControlNetModel.from_pretrained(
+        "BertChristiaens/controlnet-seg-room", torch_dtype=torch.float16)
+
+        self.pipe = StableDiffusionControlNetInpaintImg2ImgPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
+            controlnet=self.controlnet,
+            safety_checker=None,
+            torch_dtype=torch.float16
+        )
+
+        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
+        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
+
+        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
+    
+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
+
+        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
+
+
+
+@st.experimental_singleton(max_entries=5)
+def get_controlnet():
+    """Method to load the controlnet model
+    Returns:
+        ControlNetModel: controlnet model
+    """
+    pipe = ControlNetPipeline()
+    return pipe
+
+
+
+@st.experimental_singleton(max_entries=5)
+def get_inpainting_pipeline():
+    """Method to load the inpainting pipeline
+    Returns:
+        StableDiffusionInpaintPipeline: inpainting pipeline
+    """
+    pipe = SDPipeline()
+    return pipe

segmentation.py

@@ -0,0 +1,55 @@
+import logging
+from typing import List, Tuple, Dict
+
+import streamlit as st
+import torch
+import gc
+import numpy as np
+from PIL import Image
+
+from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
+
+from palette import ade_palette
+
+LOGGING = logging.getLogger(__name__)
+
+
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+@st.experimental_singleton(max_entries=5)
+def get_segmentation_pipeline() -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
+    """Method to load the segmentation pipeline
+    Returns:
+        Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]: segmentation pipeline
+    """
+    image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
+    image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
+        "openmmlab/upernet-convnext-small")
+    return image_processor, image_segmentor
+
+
+@torch.inference_mode()
+@torch.autocast('cuda')
+def segment_image(image: Image) -> Image:
+    """Method to segment image
+    Args:
+        image (Image): input image
+    Returns:
+        Image: segmented image
+    """
+    image_processor, image_segmentor = get_segmentation_pipeline()
+    pixel_values = image_processor(image, return_tensors="pt").pixel_values
+    with torch.no_grad():
+        outputs = image_segmentor(pixel_values)
+
+    seg = image_processor.post_process_semantic_segmentation(
+        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
+    color_seg = color_seg.astype(np.uint8)
+    seg_image = Image.fromarray(color_seg).convert('RGB')
+    return seg_image