cache

app.py

@@ -7,15 +7,7 @@ import numpy as np
 import os
 import time
 
-# get OS environment variables of EMULATED
-EMULATED = os.environ.get('EMULATED', False)
-print(EMULATED)
-
-if not EMULATED:
-    from models import make_image_controlnet, make_inpainting, segment_image
-else:
-    from models_stub import make_image_controlnet, make_inpainting, segment_image
-
+from models import make_image_controlnet, make_inpainting, 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
@@ -35,7 +27,6 @@ def on_upload() -> None:
             del st.session_state['unique_colors']
 
 
-
 def check_reset_state() -> bool:
     """Check whether the UI elements need to be reset
     Returns:

models.py

@@ -75,8 +75,7 @@ def postprocess_image_masking(inpainted: Image, image: Image, mask: Image) -> Im
     return final_inpainted.convert("RGB")
 
 
-@st.experimental_singleton(max_entries=1)
-@st.cache_resource
+@st.experimental_singleton(max_entries=5)
 def get_controlnet() -> ControlNetModel:
     """Method to load the controlnet model
     Returns:
@@ -100,8 +99,7 @@ def get_controlnet() -> ControlNetModel:
     return pipe, compel_proc
 
 
-@st.experimental_singleton(max_entries=1)
-@st.cache_resource
+@st.experimental_singleton(max_entries=5)
 def get_segmentation_pipeline() -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
     """Method to load the segmentation pipeline
     Returns:
@@ -113,8 +111,7 @@ def get_segmentation_pipeline() -> Tuple[AutoImageProcessor, UperNetForSemanticS
     return image_processor, image_segmentor
 
 
-@st.experimental_singleton(max_entries=1)
-@st.cache_resource
+@st.experimental_singleton(max_entries=5)
 def get_inpainting_pipeline() -> StableDiffusionInpaintPipeline:
     """Method to load the inpainting pipeline
     Returns:

models_stub.py

@@ -1,25 +0,0 @@
-import numpy as np
-
-def make_image_controlnet(image,
-                        mask_image,
-                        controlnet_conditioning_image,
-                        positive_prompt,
-                        negative_prompt,
-                        seed,
-                        ):
-    print("EMULATED CONTROLNET")
-    return [image]
-
-def make_inpainting(positive_prompt,
-                    image,
-                    mask_image,
-                    negative_prompt,
-                    ):
-    print("EMULATED INPAINTING")
-    return [image]
-
-def segment_image(image):
-    # numpy array of shape (width, height, 3) with ones
-    print("EMULATED SEGMENTATION")
-    return np.ones((image.width, image.height, 3))
-