Update pipelines.py

pipelines.py

@@ -16,46 +16,52 @@ from diffusers import ControlNetModel, UniPCMultistepScheduler
 from diffusers import StableDiffusionInpaintPipeline
 
 from config import WIDTH, HEIGHT
-from stable_diffusion_controlnet_inpaint_img2img import StableDiffusionControlNetInpaintImg2ImgPipeline
+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)
+            "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
+            torch_dtype=torch.float16,
         )
 
-        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
+        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
+        # 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
@@ -64,7 +70,8 @@ class ControlNetPipeline:
         self.waiting_queue.pop(0)
         flush()
         return results
-    
+
+
 class SDPipeline:
     def __init__(self):
         self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
@@ -75,25 +82,25 @@ class SDPipeline:
 
         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
+        # 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
@@ -104,8 +111,7 @@ class SDPipeline:
         return results
 
 
-
-@st.experimental_singleton(max_entries=5)
+@st.experimental_singleton(max_entries=5,show_spinner=False)
 def get_controlnet():
     """Method to load the controlnet model
     Returns:
@@ -115,8 +121,7 @@ def get_controlnet():
     return pipe
 
 
-
-@st.experimental_singleton(max_entries=5)
+@st.experimental_singleton(max_entries=5,show_spinner=False)
 def get_inpainting_pipeline():
     """Method to load the inpainting pipeline
     Returns: