import gradio as gr import torch import cv2 from diffusers import ControlNetModel, StableDiffusionControlNetPipeline 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, ) stable_model_list = [ "runwayml/stable-diffusion-v1-5", "dreamlike-art/dreamlike-diffusion-1.0", "kadirnar/maturemalemix_v0", "kadirnar/DreamShaper_v6" ] stable_inpiant_model_list = [ "stabilityai/stable-diffusion-2-inpainting", "runwayml/stable-diffusion-inpainting", "saik0s/realistic_vision_inpainting", ] controlnet_model_list = [ "lllyasviel/control_v11p_sd15_canny", "lllyasviel/control_v11f1p_sd15_depth", "lllyasviel/control_v11p_sd15_openpose", "lllyasviel/control_v11p_sd15_scribble", "lllyasviel/control_v11p_sd15_mlsd", "lllyasviel/control_v11e_sd15_shuffle", "lllyasviel/control_v11e_sd15_ip2p", "lllyasviel/control_v11p_sd15_lineart", "lllyasviel/control_v11p_sd15s2_lineart_anime", "lllyasviel/control_v11p_sd15_softedge", ] class StableDiffusionControlNetGenerator(ControlnetPipeline): def __init__(self): self.pipe = None 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 = StableDiffusionControlNetPipeline.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.scheduler_name = scheduler self.pipe.to("cuda") self.pipe.enable_xformers_memory_efficient_attention() return self.pipe 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, height: int, width: int, guess_mode: bool, controlnet_conditioning_scale: int, prompt: str, negative_prompt: str, num_images_per_prompt: int, guidance_scale: int, num_inference_step: int, scheduler: str, seed_generator: int, preprocces_type: str, ): pipe = self.load_model( stable_model_path=stable_model_path, controlnet_model_path=controlnet_model_path, scheduler=scheduler, ) if preprocces_type== "ScribbleXDOG": read_image = cv2.imread(image_path) controlnet_image = self.controlnet_preprocces(read_image=read_image, preprocces_type=preprocces_type)[0] controlnet_image = Image.fromarray(controlnet_image) elif preprocces_type== "None": controlnet_image = self.controlnet_preprocces(read_image=image_path, preprocces_type=preprocces_type) else: read_image = Image.open(image_path) controlnet_image = self.controlnet_preprocces(read_image=read_image, preprocces_type=preprocces_type) 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, height=height, width=width, controlnet_conditioning_scale=float(controlnet_conditioning_scale), guess_mode=guess_mode, image=controlnet_image, negative_prompt=negative_prompt, num_images_per_prompt=num_images_per_prompt, num_inference_steps=num_inference_step, guidance_scale=guidance_scale, generator=generator, ).images return output def app(): with gr.Blocks(): with gr.Row(): with gr.Column(): controlnet_image_path = gr.Image( type="filepath", label="Image" ).style(height=260) controlnet_prompt = gr.Textbox( lines=1, placeholder="Prompt", show_label=False ) controlnet_negative_prompt = gr.Textbox( lines=1, placeholder="Negative Prompt", show_label=False ) with gr.Row(): with gr.Column(): controlnet_stable_model_path = gr.Dropdown( choices=stable_model_list, value=stable_model_list[0], label="Stable Model Path", ) controlnet_preprocces_type = gr.Dropdown( choices=list(PREPROCCES_DICT.keys()), value=list(PREPROCCES_DICT.keys())[0], label="Preprocess Type", ) controlnet_conditioning_scale = gr.Slider( minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="ControlNet Conditioning Scale", ) controlnet_guidance_scale = gr.Slider( minimum=0.1, maximum=15, step=0.1, value=7.5, label="Guidance Scale", ) controlnet_height = gr.Slider( minimum=128, maximum=1280, step=32, value=512, label="Height", ) controlnet_width = gr.Slider( minimum=128, maximum=1280, step=32, value=512, label="Width", ) with gr.Row(): with gr.Column(): controlnet_model_path = gr.Dropdown( choices=controlnet_model_list, value=controlnet_model_list[0], label="ControlNet Model Path", ) controlnet_scheduler = gr.Dropdown( choices=list(SCHEDULER_MAPPING.keys()), value=list(SCHEDULER_MAPPING.keys())[0], label="Scheduler", ) controlnet_num_inference_step = gr.Slider( minimum=1, maximum=150, step=1, value=30, label="Num Inference Step", ) controlnet_num_images_per_prompt = gr.Slider( minimum=1, maximum=4, step=1, value=1, label="Number Of Images", ) controlnet_seed_generator = gr.Slider( minimum=0, maximum=1000000, step=1, value=0, label="Seed(0 for random)", ) controlnet_guess_mode = gr.Checkbox( label="Guess Mode" ) # Button to generate the image predict_button = gr.Button(value="Generate Image") with gr.Column(): # Gallery to display the generated images output_image = gr.Gallery( label="Generated images", show_label=False, elem_id="gallery", ).style(grid=(1, 2)) predict_button.click( fn=StableDiffusionControlNetGenerator().generate_image, inputs=[ controlnet_image_path, controlnet_stable_model_path, controlnet_model_path, controlnet_height, controlnet_width, controlnet_guess_mode, controlnet_conditioning_scale, controlnet_prompt, controlnet_negative_prompt, controlnet_num_images_per_prompt, controlnet_guidance_scale, controlnet_num_inference_step, controlnet_scheduler, controlnet_seed_generator, controlnet_preprocces_type, ], outputs=[output_image], )