Update app.py

app.py

@@ -18,6 +18,7 @@ import gc
 device = "cuda"
 dtype = torch.float16
 
+        
 css = """
 #img-display-container {
     max-height: 50vh;
@@ -30,11 +31,23 @@ css = """
     }
 """
 
+
 def filter_items(
     colors_list: Union[List, np.ndarray],
     items_list: Union[List, np.ndarray],
     items_to_remove: Union[List, np.ndarray]
 ) -> Tuple[Union[List, np.ndarray], Union[List, np.ndarray]]:
+    """
+    Filters items and their corresponding colors from given lists, excluding
+    specified items.
+    Args:
+        colors_list: A list or numpy array of colors corresponding to items.
+        items_list: A list or numpy array of items.
+        items_to_remove: A list or numpy array of items to be removed.
+    Returns:
+        A tuple of two lists or numpy arrays: filtered colors and filtered
+        items.
+    """
     filtered_colors = []
     filtered_items = []
     for color, item in zip(colors_list, items_list):
@@ -43,11 +56,21 @@ def filter_items(
             filtered_items.append(item)
     return filtered_colors, filtered_items
 
-def get_segmentation_pipeline() -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
-    image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
-    image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
+def get_segmentation_pipeline(
+) -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
+    """Method to load the segmentation pipeline
+    Returns:
+        Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]: segmentation pipeline
+    """
+    image_processor = AutoImageProcessor.from_pretrained(
+        "openmmlab/upernet-convnext-small"
+    )
+    image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
+        "openmmlab/upernet-convnext-small"
+    )
     return image_processor, image_segmentor
 
+
 @torch.inference_mode()
 @spaces.GPU
 def segment_image(
@@ -55,6 +78,19 @@ def segment_image(
         image_processor: AutoImageProcessor,
         image_segmentor: UperNetForSemanticSegmentation
 ) -> Image:
+    """
+    Segments an image using a semantic segmentation model.
+    Args:
+        image (Image): The input image to be segmented.
+        image_processor (AutoImageProcessor): The processor to prepare the
+            image for segmentation.
+        image_segmentor (UperNetForSemanticSegmentation): The semantic
+            segmentation model used to identify different segments in the image.
+    Returns:
+        Image: The segmented image with each segment colored differently based
+            on its identified class.
+    """
+    # image_processor, image_segmentor = get_segmentation_pipeline()
     pixel_values = image_processor(image, return_tensors="pt").pixel_values
     with torch.no_grad():
         outputs = image_segmentor(pixel_values)
@@ -69,11 +105,15 @@ def segment_image(
     seg_image = Image.fromarray(color_seg).convert('RGB')
     return seg_image
 
+
 def get_depth_pipeline():
-    feature_extractor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf", torch_dtype=dtype)
-    depth_estimator = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf", torch_dtype=dtype)
+    feature_extractor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf",
+                                                           torch_dtype=dtype)
+    depth_estimator = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf",
+                                                                  torch_dtype=dtype)
     return feature_extractor, depth_estimator
 
+
 @torch.inference_mode()
 @spaces.GPU
 def get_depth_image(
@@ -101,31 +141,53 @@ def get_depth_image(
     image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
     return image
 
+
 def resize_dimensions(dimensions, target_size):
+    """ 
+    Resize PIL to target size while maintaining aspect ratio 
+    If smaller than target size leave it as is
+    """
     width, height = dimensions
 
-    if width < target_size và height < target_size:
+    # Check if both dimensions are smaller than the target size
+    if width < target_size and height < target_size:
         return dimensions
 
+    # Determine the larger side
     if width > height:
+        # Calculate the aspect ratio
         aspect_ratio = height / width
+        # Resize dimensions
         return (target_size, int(target_size * aspect_ratio))
     else:
+        # Calculate the aspect ratio
         aspect_ratio = width / height
+        # Resize dimensions
         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):
+        """ Initialize your model(s) here """
+        #os.environ['HF_HUB_OFFLINE'] = "True"
+        
         self.seed = 323*111
         self.neg_prompt = "window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner"
-        self.additional_quality_suffix = "4K, high resolution"
+        self.control_items = ["windowpane;window", "door;double;door"]
+        self.additional_quality_suffix = "interior design, 4K, high resolution, photorealistic"
         
     @spaces.GPU
-    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:
+    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.
+        """
         print(prompt)
         flush()
         self.generator = torch.Generator(device=device).manual_seed(self.seed)
@@ -158,6 +220,7 @@ class ControlNetDepthDesignModelMulti:
 
         image_depth = get_depth_image(image, depth_feature_extractor, 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
@@ -190,12 +253,13 @@ class ControlNetDepthDesignModelMulti:
         
         return design_image
 
+
 def create_demo(model):
     gr.Markdown("### Stable Design demo")
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type='pil', elem_id='img-display-input')
-            input_text = gr.Textbox(label='Prompt', placeholder='clay render style, grayscale, photorealistic', lines=2)
+            input_text = gr.Textbox(label='Prompt', placeholder='Please upload your image first', lines=2)
             with gr.Accordion('Advanced options', open=False):
                 num_steps = gr.Slider(label='Steps',
                                       minimum=1,
@@ -223,50 +287,60 @@ def create_demo(model):
                                            maximum=1.0,
                                            value=0.9,
                                            step=0.1)
-                a_prompt = gr.Textbox(                    label='Added Prompt',
-                    value="4K, high resolution")
+                a_prompt = gr.Textbox(
+                    label='Added Prompt',
+                    value="4K, high resolution, photorealistic")
                 n_prompt = gr.Textbox(
                     label='Negative Prompt',
-                    value="low resolution, banner, logo, watermark, deformed, blurry, out of focus, surreal, ugly, beginner")
+                    value=" low resolution, banner, logo, watermark, deformed, blurry, out of focus, surreal, ugly, beginner")
             submit = gr.Button("Submit")
         
         with gr.Column():
             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, 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, 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, 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], cache_examples=False)
 
-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, img_size=img_size)
-    return out_img
 
-controlnet_depth = ControlNetModel.from_pretrained(
+controlnet_depth= ControlNetModel.from_pretrained(
     "controlnet_depth", torch_dtype=dtype, use_safetensors=True)
 controlnet_seg = ControlNetModel.from_pretrained(
     "own_controlnet", torch_dtype=dtype, use_safetensors=True)
 
 pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     "SG161222/Realistic_Vision_V5.1_noVAE",
+    #"models/runwayml--stable-diffusion-inpainting",
     controlnet=[controlnet_depth, controlnet_seg],
     safety_checker=None,
     torch_dtype=dtype
 )
 
-pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
+pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models",
+                     weight_name="ip-adapter_sd15.bin")
 pipe.set_ip_adapter_scale(0.4)
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 pipe = pipe.to(device)
-
-guide_pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", torch_dtype=dtype, use_safetensors=True, variant="fp16")
+guide_pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B",
+                                                       torch_dtype=dtype, use_safetensors=True, variant="fp16")
 guide_pipe = guide_pipe.to(device)
-
+   
 seg_image_processor, image_segmentor = get_segmentation_pipeline()
 depth_feature_extractor, depth_estimator = get_depth_pipeline()
 depth_estimator = depth_estimator.to(device)
 
+
 def main():
     model = ControlNetDepthDesignModelMulti()
     print('Models uploaded successfully')
@@ -275,7 +349,7 @@ def main():
     description = """
     WELCOME
     """
-    with gr.Blocks(css=css) as demo:
+    with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
         
@@ -283,6 +357,6 @@ def main():
 
     demo.queue().launch(share=False)
 
+
 if __name__ == '__main__':
-    main()
-                
+    main()