add queue

models.py

@@ -4,6 +4,7 @@ from typing import List, Tuple, Dict
 
 import streamlit as st
 import torch
+import time
 import numpy as np
 from PIL import Image
 from time import perf_counter
@@ -23,6 +24,45 @@ from stable_diffusion_controlnet_inpaint_img2img import StableDiffusionControlNe
 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
+    
+    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
+        self.waiting_queue.pop(0)
+        print("It's the turn of", self.count)
+        return self.pipe(**kwargs)
+
+
 @contextmanager
 def catchtime(message: str) -> float:
     """Context manager to measure time
@@ -81,22 +121,8 @@ def get_controlnet() -> ControlNetModel:
     Returns:
         ControlNetModel: controlnet model
     """
-    controlnet = ControlNetModel.from_pretrained(
-        "BertChristiaens/controlnet-seg-room", torch_dtype=torch.float16)
-
-    pipe = StableDiffusionControlNetInpaintImg2ImgPipeline.from_pretrained(
-        "runwayml/stable-diffusion-inpainting",
-        controlnet=controlnet,
-        safety_checker=None,
-        torch_dtype=torch.float16
-    )
-
-    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to("cuda")
-
-    compel_proc = Compel(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder)
-    return pipe, compel_proc
+    pipe = ControlNetPipeline()
+    return pipe
 
 
 @st.experimental_singleton(max_entries=5)
@@ -126,9 +152,7 @@ def get_inpainting_pipeline() -> StableDiffusionInpaintPipeline:
     pipe.enable_xformers_memory_efficient_attention()
     pipe = pipe.to("cuda")
 
-    compel_proc = Compel(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder)
-
-    return pipe, compel_proc
+    return pipe
 
 
 def make_grid_parameters(grid_search: Dict, params: Dict) -> List[Dict]:
@@ -185,27 +209,18 @@ def make_image_controlnet(image: np.ndarray,
     """
 
     with catchtime("get controlnet"):
-        pipe, proc = get_controlnet()
+        pipe = get_controlnet()
 
     torch.cuda.empty_cache()
     images = []
 
-    if '+' in positive_prompt or '-' in positive_prompt or '+' in negative_prompt or '-' in negative_prompt:
-        common_parameters = {'prompt_embeds': proc(positive_prompt),
-                            'negative_prompt_embeds': proc(negative_prompt),
-                            'num_inference_steps': 30,
-                             'controlnet_conditioning_scale': 1.1,
-                             'controlnet_conditioning_scale_decay': 0.96,
-                             'controlnet_steps': 28,
-                            }
-    else:
-        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,
-                            }
+    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],
@@ -253,18 +268,12 @@ def make_inpainting(positive_prompt: str,
     """
 
     with catchtime("Get inpainting pipeline"):
-        pipe, proc = get_inpainting_pipeline()
-
-    if '+' in positive_prompt or '-' in positive_prompt or '+' in negative_prompt or '-' in negative_prompt:
-        common_parameters = {'prompt_embeds': proc(positive_prompt),
-                            'negative_prompt_embeds': proc(negative_prompt),
-                            'num_inference_steps': 20,
-                            }
-    else:
-        common_parameters = {'prompt': positive_prompt,
-                            'negative_prompt': negative_prompt,
-                            'num_inference_steps': 20,
-                            }
+        pipe = get_inpainting_pipeline()
+
+    common_parameters = {'prompt': positive_prompt,
+                        'negative_prompt': negative_prompt,
+                        'num_inference_steps': 20,
+                        }
 
     torch.cuda.empty_cache()
     images = []