MykolaL/StableDesign
120
1import spaces2from typing import Tuple, Union, List3import os4 5import numpy as np6from PIL import Image7 8import torch9from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline10from diffusers import ControlNetModel, UniPCMultistepScheduler, AutoPipelineForText2Image11from transformers import AutoImageProcessor, UperNetForSemanticSegmentation, AutoModelForDepthEstimation12from colors import ade_palette13from utils import map_colors_rgb14from diffusers import StableDiffusionXLPipeline15import gradio as gr16import gc17 18device = "cuda"19dtype = torch.float1620 21 22css = """23#img-display-container {24 max-height: 50vh;25 }26#img-display-input {27 max-height: 40vh;28 }29#img-display-output {30 max-height: 40vh;31 }32 33"""34 35 36def filter_items(37 colors_list: Union[List, np.ndarray],38 items_list: Union[List, np.ndarray],39 items_to_remove: Union[List, np.ndarray]40) -> Tuple[Union[List, np.ndarray], Union[List, np.ndarray]]:41 """42 Filters items and their corresponding colors from given lists, excluding43 specified items.44 45 Args:46 colors_list: A list or numpy array of colors corresponding to items.47 items_list: A list or numpy array of items.48 items_to_remove: A list or numpy array of items to be removed.49 50 Returns:51 A tuple of two lists or numpy arrays: filtered colors and filtered52 items.53 """54 filtered_colors = []55 filtered_items = []56 for color, item in zip(colors_list, items_list):57 if item not in items_to_remove:58 filtered_colors.append(color)59 filtered_items.append(item)60 return filtered_colors, filtered_items61 62def get_segmentation_pipeline(63) -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:64 """Method to load the segmentation pipeline65 Returns:66 Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]: segmentation pipeline67 """68 image_processor = AutoImageProcessor.from_pretrained(69 "openmmlab/upernet-convnext-small"70 )71 image_segmentor = UperNetForSemanticSegmentation.from_pretrained(72 "openmmlab/upernet-convnext-small"73 )74 return image_processor, image_segmentor75 76 77@torch.inference_mode()78@spaces.GPU79def segment_image(80 image: Image,81 image_processor: AutoImageProcessor,82 image_segmentor: UperNetForSemanticSegmentation83) -> Image:84 """85 Segments an image using a semantic segmentation model.86 87 Args:88 image (Image): The input image to be segmented.89 image_processor (AutoImageProcessor): The processor to prepare the90 image for segmentation.91 image_segmentor (UperNetForSemanticSegmentation): The semantic92 segmentation model used to identify different segments in the image.93 94 Returns:95 Image: The segmented image with each segment colored differently based96 on its identified class.97 """98 # image_processor, image_segmentor = get_segmentation_pipeline()99 pixel_values = image_processor(image, return_tensors="pt").pixel_values100 with torch.no_grad():101 outputs = image_segmentor(pixel_values)102 103 seg = image_processor.post_process_semantic_segmentation(104 outputs, target_sizes=[image.size[::-1]])[0]105 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)106 palette = np.array(ade_palette())107 for label, color in enumerate(palette):108 color_seg[seg == label, :] = color109 color_seg = color_seg.astype(np.uint8)110 seg_image = Image.fromarray(color_seg).convert('RGB')111 return seg_image112 113 114def get_depth_pipeline():115 feature_extractor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf",116 torch_dtype=dtype)117 depth_estimator = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf",118 torch_dtype=dtype)119 return feature_extractor, depth_estimator120 121 122@torch.inference_mode()123@spaces.GPU124def get_depth_image(125 image: Image,126 feature_extractor: AutoImageProcessor,127 depth_estimator: AutoModelForDepthEstimation128) -> Image:129 image_to_depth = feature_extractor(images=image, return_tensors="pt").to(device)130 with torch.no_grad():131 depth_map = depth_estimator(**image_to_depth).predicted_depth132 133 width, height = image.size134 depth_map = torch.nn.functional.interpolate(135 depth_map.unsqueeze(1).float(),136 size=(height, width),137 mode="bicubic",138 align_corners=False,139 )140 depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)141 depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)142 depth_map = (depth_map - depth_min) / (depth_max - depth_min)143 image = torch.cat([depth_map] * 3, dim=1)144 145 image = image.permute(0, 2, 3, 1).cpu().numpy()[0]146 image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))147 return image148 149 150def resize_dimensions(dimensions, target_size):151 """ 152 Resize PIL to target size while maintaining aspect ratio 153 If smaller than target size leave it as is154 """155 width, height = dimensions156 157 # Check if both dimensions are smaller than the target size158 if width < target_size and height < target_size:159 return dimensions160 161 # Determine the larger side162 if width > height:163 # Calculate the aspect ratio164 aspect_ratio = height / width165 # Resize dimensions166 return (target_size, int(target_size * aspect_ratio))167 else:168 # Calculate the aspect ratio169 aspect_ratio = width / height170 # Resize dimensions171 return (int(target_size * aspect_ratio), target_size)172 173 174def flush():175 gc.collect()176 torch.cuda.empty_cache()177 178 179class ControlNetDepthDesignModelMulti:180 """ Produces random noise images """181 182 def __init__(self):183 """ Initialize your model(s) here """184 #os.environ['HF_HUB_OFFLINE'] = "True"185 186 self.seed = 323*111187 self.neg_prompt = "window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner"188 self.control_items = ["windowpane;window", "door;double;door"]189 self.additional_quality_suffix = "interior design, 4K, high resolution, photorealistic"190 191 @spaces.GPU192 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:193 """194 Given an image of an empty room and a prompt195 generate the designed room according to the prompt196 Inputs - 197 empty_room_image - An RGB PIL Image of the empty room198 prompt - Text describing the target design elements of the room199 Returns - 200 design_image - PIL Image of the same size as the empty room image201 If the size is not the same the submission will fail.202 """203 print(prompt)204 flush()205 self.generator = torch.Generator(device=device).manual_seed(self.seed)206 207 pos_prompt = prompt + f', {self.additional_quality_suffix}'208 209 orig_w, orig_h = empty_room_image.size210 new_width, new_height = resize_dimensions(empty_room_image.size, img_size)211 input_image = empty_room_image.resize((new_width, new_height))212 real_seg = np.array(segment_image(input_image,213 seg_image_processor,214 image_segmentor))215 unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)216 unique_colors = [tuple(color) for color in unique_colors]217 segment_items = [map_colors_rgb(i) for i in unique_colors]218 chosen_colors, segment_items = filter_items(219 colors_list=unique_colors,220 items_list=segment_items,221 items_to_remove=self.control_items222 )223 mask = np.zeros_like(real_seg)224 for color in chosen_colors:225 color_matches = (real_seg == color).all(axis=2)226 mask[color_matches] = 1227 228 image_np = np.array(input_image)229 image = Image.fromarray(image_np).convert("RGB")230 mask_image = Image.fromarray((mask * 255).astype(np.uint8)).convert("RGB")231 segmentation_cond_image = Image.fromarray(real_seg).convert("RGB")232 233 image_depth = get_depth_image(image, depth_feature_extractor, depth_estimator)234 235 # generate image that would be used as IP-adapter236 flush()237 new_width_ip = int(new_width / 8) * 8238 new_height_ip = int(new_height / 8) * 8239 ip_image = guide_pipe(pos_prompt,240 num_inference_steps=num_steps,241 negative_prompt=self.neg_prompt,242 height=new_height_ip,243 width=new_width_ip,244 generator=[self.generator]).images[0]245 246 flush()247 generated_image = pipe(248 prompt=pos_prompt,249 negative_prompt=self.neg_prompt,250 num_inference_steps=num_steps,251 strength=strength,252 guidance_scale=guidance_scale,253 generator=[self.generator],254 image=image,255 mask_image=mask_image,256 ip_adapter_image=ip_image,257 control_image=[image_depth, segmentation_cond_image],258 controlnet_conditioning_scale=[0.5, 0.5]259 ).images[0]260 261 flush()262 design_image = generated_image.resize(263 (orig_w, orig_h), Image.Resampling.LANCZOS264 )265 266 return design_image267 268 269def create_demo(model):270 gr.Markdown("### Stable Design demo")271 with gr.Row():272 with gr.Column():273 input_image = gr.Image(label="Input Image", type='pil', elem_id='img-display-input')274 input_text = gr.Textbox(label='Prompt', placeholder='Please upload your image first', lines=2)275 with gr.Accordion('Advanced options', open=False):276 num_steps = gr.Slider(label='Steps',277 minimum=1,278 maximum=50,279 value=50,280 step=1)281 img_size = gr.Slider(label='Image size',282 minimum=256,283 maximum=768,284 value=768,285 step=64)286 guidance_scale = gr.Slider(label='Guidance Scale',287 minimum=0.1,288 maximum=30.0,289 value=10.0,290 step=0.1)291 seed = gr.Slider(label='Seed',292 minimum=-1,293 maximum=2147483647,294 value=323*111,295 step=1,296 randomize=True)297 strength = gr.Slider(label='Strength',298 minimum=0.1,299 maximum=1.0,300 value=0.9,301 step=0.1)302 a_prompt = gr.Textbox(303 label='Added Prompt',304 value="interior design, 4K, high resolution, photorealistic")305 n_prompt = gr.Textbox(306 label='Negative Prompt',307 value="window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner")308 submit = gr.Button("Submit")309 310 with gr.Column():311 design_image = gr.Image(label="Output Mask", elem_id='img-display-output')312 313 314 def on_submit(image, text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt, img_size):315 model.seed = seed316 model.neg_prompt = n_prompt317 model.additional_quality_suffix = a_prompt318 319 with torch.no_grad():320 out_img = model.generate_design(image, text, guidance_scale=guidance_scale, num_steps=num_steps, strength=strength, img_size=img_size)321 322 return out_img323 324 submit.click(on_submit, inputs=[input_image, input_text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt, img_size], outputs=design_image)325 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"]],326 inputs=[input_image, input_text], cache_examples=False)327 328 329controlnet_depth= ControlNetModel.from_pretrained(330 "controlnet_depth", torch_dtype=dtype, use_safetensors=True)331controlnet_seg = ControlNetModel.from_pretrained(332 "own_controlnet", torch_dtype=dtype, use_safetensors=True)333 334pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(335 "SG161222/Realistic_Vision_V5.1_noVAE",336 #"models/runwayml--stable-diffusion-inpainting",337 controlnet=[controlnet_depth, controlnet_seg],338 safety_checker=None,339 torch_dtype=dtype340)341 342pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models",343 weight_name="ip-adapter_sd15.bin")344pipe.set_ip_adapter_scale(0.4)345pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)346pipe = pipe.to(device)347guide_pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B",348 torch_dtype=dtype, use_safetensors=True, variant="fp16")349guide_pipe = guide_pipe.to(device)350 351seg_image_processor, image_segmentor = get_segmentation_pipeline()352depth_feature_extractor, depth_estimator = get_depth_pipeline()353depth_estimator = depth_estimator.to(device)354 355 356def main():357 model = ControlNetDepthDesignModelMulti()358 print('Models uploaded successfully')359 360 title = "# StableDesign"361 description = """362 <p style='font-size: 14px; margin-bottom: 10px;'><a href='https://www.linkedin.com/in/mykola-lavreniuk/'>Mykola Lavreniuk</a>, <a href='https://www.linkedin.com/in/bartosz-ludwiczuk-a677a760/'>Bartosz Ludwiczuk</a></p>363 <p style='font-size: 16px; margin-bottom: 0px; margin-top=0px;'>Official demo for <strong>StableDesign:</strong> 2nd place solution for the Generative Interior Design 2024 <a href='https://www.aicrowd.com/challenges/generative-interior-design-challenge-2024/leaderboards?challenge_round_id=1314'>competition</a>. StableDesign is a deep learning model designed to harness the power of AI, providing innovative and creative tools for designers. Using our algorithms, images of empty rooms can be transformed into fully furnished spaces based on text descriptions. Please refer to our <a href='https://github.com/Lavreniuk/generative-interior-design'>GitHub</a> for more details.</p>364 """365 with gr.Blocks() as demo:366 gr.Markdown(title)367 gr.Markdown(description)368 369 create_demo(model)370 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>371 <p><img src="https://visitor-badge.glitch.me/badge?page_id=MykolaL/StableDesign" alt="visitors"></p></center>''')372 373 demo.queue().launch(share=False)374 375 376if __name__ == '__main__':377 main()378