init

app.py

@@ -12,6 +12,7 @@ from colors import ade_palette
 from utils import map_colors_rgb
 from diffusers import StableDiffusionXLPipeline
 import gradio as gr
+import gc
 
 device = "cpu"#"cuda"
 dtype = torch.float32#torch.float16
@@ -167,6 +168,11 @@ def resize_dimensions(dimensions, target_size):
         return (int(target_size * aspect_ratio), target_size)
 
 
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    
 class ControlNetDepthDesignModelMulti:
     """ Produces random noise images """
     def __init__(self):
@@ -204,7 +210,7 @@ class ControlNetDepthDesignModelMulti:
         self.depth_feature_extractor, self.depth_estimator = get_depth_pipeline()
         self.depth_estimator = self.depth_estimator.to(device)
 
-    def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9) -> Image:
+    def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9, img_size: int = 640) -> Image:
         """
         Given an image of an empty room and a prompt
         generate the designed room according to the prompt
@@ -216,12 +222,13 @@ class ControlNetDepthDesignModelMulti:
                            If the size is not the same the submission will fail.
         """
         print(prompt)
+        flush()
         self.generator = torch.Generator(device=device).manual_seed(self.seed)
 
         pos_prompt = prompt + f', {self.additional_quality_suffix}'
 
         orig_w, orig_h = empty_room_image.size
-        new_width, new_height = resize_dimensions(empty_room_image.size, 768)
+        new_width, new_height = resize_dimensions(empty_room_image.size, img_size)
         input_image = empty_room_image.resize((new_width, new_height))
         real_seg = np.array(segment_image(input_image,
                                           self.seg_image_processor,
@@ -247,6 +254,7 @@ class ControlNetDepthDesignModelMulti:
         image_depth = get_depth_image(image, self.depth_feature_extractor, self.depth_estimator)
 
         # generate image that would be used as IP-adapter
+        flush()
         new_width_ip = int(new_width / 8) * 8
         new_height_ip = int(new_height / 8) * 8
         ip_image = self.guide_pipe(pos_prompt,
@@ -256,7 +264,7 @@ class ControlNetDepthDesignModelMulti:
                                    width=new_width_ip,
                                    generator=[self.generator]).images[0]
 
-
+        flush()
         generated_image = self.pipe(
             prompt=pos_prompt,
             negative_prompt=self.neg_prompt,
@@ -270,7 +278,8 @@ class ControlNetDepthDesignModelMulti:
             control_image=[image_depth, segmentation_cond_image],
             controlnet_conditioning_scale=[0.5, 0.5]
         ).images[0]
-
+        
+        flush()
         design_image = generated_image.resize(
             (orig_w, orig_h), Image.Resampling.LANCZOS
         )
@@ -278,7 +287,7 @@ class ControlNetDepthDesignModelMulti:
         return design_image
 
 
-def create_refseg_demo(model):
+def create_demo(model):
     gr.Markdown("### Stable Design demo")
     with gr.Row():
         with gr.Column():
@@ -290,6 +299,11 @@ def create_refseg_demo(model):
                                       maximum=50,
                                       value=50,
                                       step=1)
+                img_size = gr.Slider(label='Image size',
+                                      minimum=256,
+                                      maximum=768,
+                                      value=768,
+                                      step=64)
                 guidance_scale = gr.Slider(label='Guidance Scale',
                                            minimum=0.1,
                                            maximum=30.0,
@@ -318,17 +332,17 @@ def create_refseg_demo(model):
             design_image = gr.Image(label="Output Mask", elem_id='img-display-output')
     
     
-    def on_submit(image, text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt):
+    def on_submit(image, text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt, img_size):
         model.seed = seed
         model.neg_prompt = n_prompt
         model.additional_quality_suffix = a_prompt
         
         with torch.no_grad():
-            out_img = model.generate_design(image, text, guidance_scale=guidance_scale, num_steps=num_steps, strength=strength)
+            out_img = model.generate_design(image, text, guidance_scale=guidance_scale, num_steps=num_steps, strength=strength, img_size=img_size)
 
         return out_img
 
-    submit.click(on_submit, inputs=[input_image, input_text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt], outputs=design_image)
+    submit.click(on_submit, inputs=[input_image, input_text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt, img_size], outputs=design_image)
     examples = gr.Examples(examples=[["imgs/bedroom_1.jpg", "An elegantly appointed bedroom in the Art Deco style, featuring a grand king-size bed with geometric bedding, a luxurious velvet armchair, and a mirrored nightstand that reflects the room's opulence. Art Deco-inspired artwork adds a touch of glamour"], ["imgs/bedroom_2.jpg", "A bedroom that exudes French country charm with a soft upholstered bed, walls adorned with floral wallpaper, and a vintage wooden wardrobe. A crystal chandelier casts a warm, inviting glow over the space"], ["imgs/dinning_room_1.jpg", "A cozy dining room that captures the essence of rustic charm with a solid wooden farmhouse table at its core, surrounded by an eclectic mix of mismatched chairs. An antique sideboard serves as a statement piece, and the ambiance is warmly lit by a series of quaint Edison bulbs dangling from the ceiling"], ["imgs/dinning_room_3.jpg", "A dining room that epitomizes contemporary elegance, anchored by a sleek, minimalist dining table paired with stylish modern chairs. Artistic lighting fixtures create a focal point above, while the surrounding minimalist decor ensures the space feels open, airy, and utterly modern"], ["imgs/image_1.jpg", "A glamorous master bedroom in Hollywood Regency style, boasting a plush tufted headboard, mirrored furniture reflecting elegance, luxurious fabrics in rich textures, and opulent gold accents for a touch of luxury."], ["imgs/image_2.jpg", "A vibrant living room with a tropical theme, complete with comfortable rattan furniture, large leafy plants bringing the outdoors in, bright cushions adding pops of color, and bamboo blinds for natural light control."], ["imgs/living_room_1.jpg", "A stylish living room embracing mid-century modern aesthetics, featuring a vintage teak coffee table at its center, complemented by a classic sunburst clock on the wall and a cozy shag rug underfoot, creating a warm and inviting atmosphere"]],
                            inputs=[input_image, input_text])
 
@@ -346,7 +360,7 @@ def main():
         gr.Markdown(title)
         gr.Markdown(description)
         
-        create_refseg_demo(model)
+        create_demo(model)
         gr.HTML('''<br><br><br><center>You can duplicate this Space to skip the queue:<a href="https://huggingface.co/spaces/MykolaL/StableDesign?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a><br>
                 <p><img src="https://visitor-badge.glitch.me/badge?page_id=MykolaL/StableDesign" alt="visitors"></p></center>''')