ml6team/controlnet-interior-design

1import logging2from typing import List, Tuple, Dict3 4import streamlit as st5import torch6import gc7import time8import numpy as np9from PIL import Image10from time import perf_counter11from contextlib import contextmanager12from scipy.signal import fftconvolve13from PIL import ImageFilter14 15from diffusers import ControlNetModel, UniPCMultistepScheduler16from diffusers import StableDiffusionInpaintPipeline17 18from config import WIDTH, HEIGHT19from stable_diffusion_controlnet_inpaint_img2img import StableDiffusionControlNetInpaintImg2ImgPipeline20from helpers import flush21 22LOGGING = logging.getLogger(__name__)23 24class ControlNetPipeline:25    def __init__(self):26        self.in_use = False27        self.controlnet = ControlNetModel.from_pretrained(28        "BertChristiaens/controlnet-seg-room", torch_dtype=torch.float16)29 30        self.pipe = StableDiffusionControlNetInpaintImg2ImgPipeline.from_pretrained(31            "runwayml/stable-diffusion-inpainting",32            controlnet=self.controlnet,33            safety_checker=None,34            torch_dtype=torch.float1635        )36 37        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)38        self.pipe.enable_xformers_memory_efficient_attention()39        self.pipe = self.pipe.to("cuda")40        41        self.waiting_queue = []42        self.count = 043    44    @property45    def queue_size(self):46        return len(self.waiting_queue)47    48    def __call__(self, **kwargs):49        self.count += 150        number = self.count51 52        self.waiting_queue.append(number)53        54        # wait until the next number in the queue is the current number55        while self.waiting_queue[0] != number:56            print(f"Wait for your turn {number} in queue {self.waiting_queue}")57            time.sleep(0.5)58            pass59 60        # it's your turn, so remove the number from the queue61        # and call the function62        print("It's the turn of", self.count)63        results = self.pipe(**kwargs)64        self.waiting_queue.pop(0)65        flush()66        return results67    68class SDPipeline:69    def __init__(self):70        self.pipe = StableDiffusionInpaintPipeline.from_pretrained(71            "stabilityai/stable-diffusion-2-inpainting",72            torch_dtype=torch.float16,73            safety_checker=None,74        )75 76        self.pipe.enable_xformers_memory_efficient_attention()77        self.pipe = self.pipe.to("cuda")78        79        self.waiting_queue = []80        self.count = 081    82    @property83    def queue_size(self):84        return len(self.waiting_queue)85    86    def __call__(self, **kwargs):87        self.count += 188        number = self.count89 90        self.waiting_queue.append(number)91        92        # wait until the next number in the queue is the current number93        while self.waiting_queue[0] != number:94            print(f"Wait for your turn {number} in queue {self.waiting_queue}")95            time.sleep(0.5)96            pass97 98        # it's your turn, so remove the number from the queue99        # and call the function100        print("It's the turn of", self.count)101        results = self.pipe(**kwargs)102        self.waiting_queue.pop(0)103        flush()104        return results105 106 107 108@st.cache_resource(max_entries=5)109def get_controlnet():110    """Method to load the controlnet model111    Returns:112        ControlNetModel: controlnet model113    """114    pipe = ControlNetPipeline()115    return pipe116 117 118 119@st.cache_resource(max_entries=5)120def get_inpainting_pipeline():121    """Method to load the inpainting pipeline122    Returns:123        StableDiffusionInpaintPipeline: inpainting pipeline124    """125    pipe = SDPipeline()126    return pipe127