flush

app.py

@@ -182,11 +182,6 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
 
 
     elif generation_mode == "Re-generate objects":
-        st.write("This mode allows you to choose which objects you want to re-generate in the image. "
-                 "Use the selection dropdown to add or remove objects. If you are ready, press the generate button"
-                 " to generate the image, which can take up to 30 seconds. If you want to improve the generated image, click"
-                 " the 'move image to input' button."
-                 )
         canvas = st_canvas(
             **canvas_dict,
         )
@@ -209,6 +204,12 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
             default=st.session_state['unique_colors'],
             format_func=map_colors_rgb,
         )
+        with st.expander("Explanation", expanded=False):
+            st.write("This mode allows you to choose which objects you want to re-generate in the image. "
+                 "Use the selection dropdown to add or remove objects. If you are ready, press the generate button"
+                 " to generate the image, which can take up to 30 seconds. If you want to improve the generated image, click"
+                 " the 'move image to input' button."
+                 )
 
         if st.button("generate image", key='generate_button'):
             image = get_image()

models.py

@@ -4,6 +4,7 @@ from typing import List, Tuple, Dict
 
 import streamlit as st
 import torch
+import gc
 import time
 import numpy as np
 from PIL import Image
@@ -23,22 +24,26 @@ from stable_diffusion_controlnet_inpaint_img2img import StableDiffusionControlNe
 
 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)
+        "BertChristiaens/controlnet-seg-room", torch_dtype=torch.float32)
 
         self.pipe = StableDiffusionControlNetInpaintImg2ImgPipeline.from_pretrained(
             "runwayml/stable-diffusion-inpainting",
             controlnet=self.controlnet,
             safety_checker=None,
-            torch_dtype=torch.float16
+            torch_dtype=torch.float32
         )
 
         self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
         self.pipe.enable_xformers_memory_efficient_attention()
+        self.pipe.enable_attention_slicing("max")
         self.pipe = self.pipe.to("cuda")
         
         self.waiting_queue = []
@@ -59,8 +64,10 @@ class ControlNetPipeline:
         # 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)
-        return self.pipe(**kwargs)
+        flush()
+        return results
 
 
 @contextmanager