import io import requests import numpy as np import torch from PIL import Image from skimage.measure import block_reduce from typing import List, Optional from functools import reduce import gradio as gr from transformers import DetrFeatureExtractor, DetrForSegmentation, DetrConfig from transformers.models.detr.feature_extraction_detr import rgb_to_id from diffusers import StableDiffusionInpaintPipeline torch.inference_mode() torch.no_grad() # Load segmentation models def load_segmentation_models(model_name: str = 'facebook/detr-resnet-50-panoptic'): feature_extractor = DetrFeatureExtractor.from_pretrained(model_name) model = DetrForSegmentation.from_pretrained(model_name) cfg = DetrConfig.from_pretrained(model_name) return feature_extractor, model, cfg # Load diffusion pipeline def load_diffusion_pipeline(model_name: str = 'runwayml/stable-diffusion-inpainting'): return StableDiffusionInpaintPipeline.from_pretrained( model_name, revision='fp16', torch_dtype=torch.float16 ) # Device helper def get_device(try_cuda=True): return torch.device('cuda' if try_cuda and torch.cuda.is_available() else 'cpu') def min_pool(x: torch.Tensor, kernel_size: int): pad_size = (kernel_size - 1) // 2 return -torch.nn.functional.max_pool2d(-x, kernel_size, (1, 1), padding=pad_size) def max_pool(x: torch.Tensor, kernel_size: int): pad_size = (kernel_size - 1) // 2 return torch.nn.functional.max_pool2d(x, kernel_size, (1, 1), padding=pad_size) # Apply min-max pooling to clean up mask def clean_mask(mask, max_kernel: int = 23, min_kernel: int = 5): mask = torch.Tensor(mask[None, None]).float() mask = min_pool(mask, min_kernel) mask = max_pool(mask, max_kernel) mask = mask.bool().squeeze().numpy() return mask feature_extractor, segmentation_model, segmentation_cfg = load_segmentation_models() pipe = load_diffusion_pipeline() device = get_device() pipe = pipe.to(device) # Callback function that runs segmentation and updates CheckboxGroup def fn_segmentation(image, max_kernel, min_kernel): inputs = feature_extractor(images=image, return_tensors="pt") outputs = segmentation_model(**inputs) processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0) result = feature_extractor.post_process_panoptic(outputs, processed_sizes)[0] panoptic_seg = Image.open(io.BytesIO(result["png_string"])).resize((image.width, image.height)) panoptic_seg = np.array(panoptic_seg, dtype=np.uint8) panoptic_seg_id = rgb_to_id(panoptic_seg) raw_masks = [] for s in result['segments_info']: m = panoptic_seg_id == s['id'] raw_masks.append(m.astype(np.uint8) * 255) checkbox_choices = [f"{s['id']}:{segmentation_cfg.id2label[s['category_id']]}" for s in result['segments_info']] checkbox_group = gr.CheckboxGroup.update( choices=checkbox_choices ) return raw_masks, checkbox_group, gr.Image.update(value=np.zeros((image.height, image.width))), gr.Image.update(value=image) # Callback function that updates the displayed mask based on selected checkboxes def fn_update_mask( image: Image, masks: List[np.array], masks_enabled: List[int], max_kernel: int, min_kernel: int, ): masks_enabled = [int(m.split(':')[0]) for m in masks_enabled] combined_mask = reduce(lambda x, y: x | y, [masks[i] for i in masks_enabled], np.zeros_like(masks[0], dtype=bool)) combined_mask = clean_mask(combined_mask, max_kernel, min_kernel) masked_image = np.array(image).copy() masked_image[combined_mask] = 0.0 return combined_mask.astype(np.uint8) * 255, Image.fromarray(masked_image) # Callback function that runs diffusion given the current image, mask and prompt. def fn_diffusion( prompt: str, masked_image: Image, mask: Image, num_diffusion_steps: int, guidance_scale: float, negative_prompt: Optional[str] = None, ): if len(negative_prompt) == 0: negative_prompt = None # Resize image to a more stable diffusion friendly format. # TODO: remove magic number STABLE_DIFFUSION_SMALL_EDGE = 512 w, h = masked_image.size is_width_larger = w > h resize_ratio = STABLE_DIFFUSION_SMALL_EDGE / (h if is_width_larger else w) new_width = int(w * resize_ratio) if is_width_larger else STABLE_DIFFUSION_SMALL_EDGE new_height = STABLE_DIFFUSION_SMALL_EDGE if is_width_larger else int(h * resize_ratio) new_width += 8 - (new_width % 8) if is_width_larger else 0 new_height += 0 if is_width_larger else 8 - (new_height % 8) mask = Image.fromarray(mask).convert("RGB").resize((new_width, new_height)) masked_image = masked_image.convert("RGB").resize((new_width, new_height)) # Run diffusion inpainted_image = pipe( height=new_height, width=new_width, prompt=prompt, image=masked_image, mask_image=mask, num_inference_steps=num_diffusion_steps, guidance_scale=guidance_scale, negative_prompt=negative_prompt ).images[0] # Resize back to the original size inpainted_image = inpainted_image.resize((w, h)) return inpainted_image demo = gr.Blocks() with demo: # Input image control input_image = gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil', label="Input Image") # Combined mask controls bt_masks = gr.Button("Compute Masks") with gr.Row(): mask_image = gr.Image(type='numpy', label="Diffusion Mask") masked_image = gr.Image(type='pil', label="Masked Image") mask_storage = gr.State() # Mask editing controls with gr.Row(): max_slider = gr.Slider(minimum=1, maximum=99, value=23, step=2, label="Mask Overflow") min_slider = gr.Slider(minimum=1, maximum=99, value=5, step=2, label="Mask Denoising") mask_checkboxes = gr.CheckboxGroup(interactive=True, label="Mask Selection") # Diffusion controls and output with gr.Row(): with gr.Column(): prompt = gr.Textbox("Two ginger cats lying together on a pink sofa. There are two TV remotes. High definition.", label="Prompt") negative_prompt = gr.Textbox(label="Negative Prompt") with gr.Column(): steps_slider = gr.Slider(minimum=1, maximum=100, value=50, label="Inference Steps") guidance_slider = gr.Slider(minimum=0.0, maximum=50.0, value=7.5, step=0.1, label="Guidance Scale") bt_diffusion = gr.Button("Run Diffusion") inpainted_image = gr.Image(type='pil', label="Inpainted Image") update_mask_inputs = [input_image, mask_storage, mask_checkboxes, max_slider, min_slider] update_mask_outputs = [mask_image, masked_image] # Clear checkbox group on input image change input_image.change(lambda: gr.CheckboxGroup.update(choices=[], value=[]), outputs=mask_checkboxes) # Segmentation button callback bt_masks.click(fn_segmentation, inputs=[input_image, max_slider, min_slider], outputs=[mask_storage, mask_checkboxes, mask_image, masked_image]) # Update mask callbacks # TODO: can we replace this with `mask_image.change`? Not sure if it will actively update. max_slider.change(fn_update_mask, inputs=update_mask_inputs, outputs=update_mask_outputs, show_progress=False) min_slider.change(fn_update_mask, inputs=update_mask_inputs, outputs=update_mask_outputs, show_progress=False) mask_checkboxes.change(fn_update_mask, inputs=update_mask_inputs, outputs=update_mask_outputs, show_progress=False) # Diffusion button callback bt_diffusion.click(fn_diffusion, inputs=[ prompt, masked_image, mask_image, steps_slider, guidance_slider, negative_prompt ], outputs=inpainted_image) demo.launch()