import gradio as gr import numpy as np import torch from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline from PIL import Image from diffusion_webui.diffusion_models.base_controlnet_pipeline import ( ControlnetPipeline, ) from diffusion_webui.utils.model_list import ( controlnet_model_list, stable_model_list, ) from diffusion_webui.utils.preprocces_utils import PREPROCCES_DICT from diffusion_webui.utils.scheduler_list import ( SCHEDULER_MAPPING, get_scheduler, ) class StableDiffusionControlNetInpaintGenerator(ControlnetPipeline): def __init__(self): super().__init__() def load_model(self, stable_model_path, controlnet_model_path, scheduler): if self.pipe is None or self.pipe.model_name != stable_model_path or self.pipe.scheduler_name != scheduler: controlnet = ControlNetModel.from_pretrained( controlnet_model_path, torch_dtype=torch.float16 ) self.pipe = ( StableDiffusionControlNetInpaintPipeline.from_pretrained( pretrained_model_name_or_path=stable_model_path, controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16, ) ) self.pipe.model_name = stable_model_path self.pipe.scheduler_name = scheduler self.pipe = get_scheduler(pipe=self.pipe, scheduler=scheduler) self.pipe.to("cuda") self.pipe.enable_xformers_memory_efficient_attention() return self.pipe def load_image(self, image): image = np.array(image) image = Image.fromarray(image) return image def controlnet_preprocces( self, read_image: str, preprocces_type: str, ): processed_image = PREPROCCES_DICT[preprocces_type](read_image) return processed_image def generate_image( self, image_path: str, stable_model_path: str, controlnet_model_path: str, prompt: str, negative_prompt: str, num_images_per_prompt: int, height: int, width: int, strength: int, guess_mode: bool, guidance_scale: int, num_inference_step: int, controlnet_conditioning_scale: int, scheduler: str, seed_generator: int, preprocces_type: str, ): normal_image = image_path["image"].convert("RGB").resize((512, 512)) mask_image = image_path["mask"].convert("RGB").resize((512, 512)) normal_image = self.load_image(image=normal_image) mask_image = self.load_image(image=mask_image) control_image = self.controlnet_preprocces( read_image=normal_image, preprocces_type=preprocces_type ) pipe = self.load_model( stable_model_path=stable_model_path, controlnet_model_path=controlnet_model_path, scheduler=scheduler, ) if seed_generator == 0: random_seed = torch.randint(0, 1000000, (1,)) generator = torch.manual_seed(random_seed) else: generator = torch.manual_seed(seed_generator) output = pipe( prompt=prompt, image=normal_image, height=height, width=width, mask_image=mask_image, strength=strength, guess_mode=guess_mode, control_image=control_image, negative_prompt=negative_prompt, num_images_per_prompt=num_images_per_prompt, num_inference_steps=num_inference_step, guidance_scale=guidance_scale, controlnet_conditioning_scale=float(controlnet_conditioning_scale), generator=generator, ).images return output def app(): with gr.Blocks(): with gr.Row(): with gr.Column(): controlnet_inpaint_image_path = gr.Image( source="upload", tool="sketch", elem_id="image_upload", type="pil", label="Upload", ).style(height=260) controlnet_inpaint_prompt = gr.Textbox( lines=1, placeholder="Prompt", show_label=False ) controlnet_inpaint_negative_prompt = gr.Textbox( lines=1, placeholder="Negative Prompt", show_label=False ) with gr.Row(): with gr.Column(): controlnet_inpaint_stable_model_path = gr.Dropdown( choices=stable_model_list, value=stable_model_list[0], label="Stable Model Path", ) controlnet_inpaint_preprocces_type = gr.Dropdown( choices=list(PREPROCCES_DICT.keys()), value=list(PREPROCCES_DICT.keys())[0], label="Preprocess Type", ) controlnet_inpaint_conditioning_scale = gr.Slider( minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="ControlNet Conditioning Scale", ) controlnet_inpaint_guidance_scale = gr.Slider( minimum=0.1, maximum=15, step=0.1, value=7.5, label="Guidance Scale", ) controlnet_inpaint_height = gr.Slider( minimum=128, maximum=1280, step=32, value=512, label="Height", ) controlnet_inpaint_width = gr.Slider( minimum=128, maximum=1280, step=32, value=512, label="Width", ) controlnet_inpaint_guess_mode = gr.Checkbox( label="Guess Mode" ) with gr.Column(): controlnet_inpaint_model_path = gr.Dropdown( choices=controlnet_model_list, value=controlnet_model_list[0], label="ControlNet Model Path", ) controlnet_inpaint_scheduler = gr.Dropdown( choices=list(SCHEDULER_MAPPING.keys()), value=list(SCHEDULER_MAPPING.keys())[0], label="Scheduler", ) controlnet_inpaint_strength = gr.Slider( minimum=0.1, maximum=15, step=0.1, value=7.5, label="Strength", ) controlnet_inpaint_num_inference_step = gr.Slider( minimum=1, maximum=150, step=1, value=30, label="Num Inference Step", ) controlnet_inpaint_num_images_per_prompt = ( gr.Slider( minimum=1, maximum=4, step=1, value=1, label="Number Of Images", ) ) controlnet_inpaint_seed_generator = gr.Slider( minimum=0, maximum=1000000, step=1, value=0, label="Seed(0 for random)", ) # Button to generate the image controlnet_inpaint_predict_button = gr.Button( value="Generate Image" ) with gr.Column(): # Gallery to display the generated images controlnet_inpaint_output_image = gr.Gallery( label="Generated images", show_label=False, elem_id="gallery", ).style(grid=(1, 2)) controlnet_inpaint_predict_button.click( fn=StableDiffusionControlNetInpaintGenerator().generate_image, inputs=[ controlnet_inpaint_image_path, controlnet_inpaint_stable_model_path, controlnet_inpaint_model_path, controlnet_inpaint_prompt, controlnet_inpaint_negative_prompt, controlnet_inpaint_num_images_per_prompt, controlnet_inpaint_height, controlnet_inpaint_width, controlnet_inpaint_strength, controlnet_inpaint_guess_mode, controlnet_inpaint_guidance_scale, controlnet_inpaint_num_inference_step, controlnet_inpaint_conditioning_scale, controlnet_inpaint_scheduler, controlnet_inpaint_seed_generator, controlnet_inpaint_preprocces_type, ], outputs=[controlnet_inpaint_output_image], )