import gradio as gr from diffusers import DiffusionPipeline,StableDiffusionInpaintPipeline import torch from .utils.prompt2prompt import generate from .utils.device import get_device from .utils.schedulers import SCHEDULER_LIST, get_scheduler_list from .download import get_share_js, CSS, get_community_loading_icon INPAINT_MODEL_LIST = { "Stable Diffusion 2" : "stabilityai/stable-diffusion-2-inpainting", "Stable Diffusion 1" : "runwayml/stable-diffusion-inpainting", } class StableDiffusionInpaintGenerator: def __init__(self): self.pipe = None def load_model(self, model_path, scheduler): model_path = INPAINT_MODEL_LIST[model_path] if self.pipe is None: self.pipe = StableDiffusionInpaintPipeline.from_pretrained( model_path, torch_dtype=torch.float32 ) device = get_device() self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler) self.pipe.to(device) self.pipe.enable_attention_slicing() return self.pipe def generate_image( self, pil_image: str, model_path: str, prompt: str, negative_prompt: str, scheduler: str, guidance_scale: int, num_inference_step: int, height: int, width: int, seed_generator=0, ): image = pil_image["image"].convert("RGB").resize((width, height)) mask_image = pil_image["mask"].convert("RGB").resize((width, height)) pipe = self.load_model(model_path,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=image, mask_image=mask_image, negative_prompt=negative_prompt, num_images_per_prompt=1, num_inference_steps=num_inference_step, guidance_scale=guidance_scale, generator=generator, ).images return output def app(): demo = gr.Blocks(css=CSS) with demo: with gr.Row(): with gr.Column(): stable_diffusion_inpaint_image_file = gr.Image( source="upload", tool="sketch", elem_id="image-upload-inpainting", type="pil", label="Upload", ).style(height=260) stable_diffusion_inpaint_prompt = gr.Textbox( lines=1, placeholder="Prompt, keywords that explains how you want to modify the image.", show_label=False, elem_id="prompt-text-input-inpainting", value='' ) stable_diffusion_inpaint_negative_prompt = gr.Textbox( lines=1, placeholder="Negative Prompt, keywords that describe what you don't want in your image", show_label=False, elem_id = "negative-prompt-text-input-inpainting", value='' ) # add button for generating a prompt from the prompt stable_diffusion_inpaint_generate = gr.Button( label="Generate Prompt", type="primary", align="center", value = "Generate Prompt" ) # show a text box with the generated prompt stable_diffusion_inpaint_generated_prompt = gr.Textbox( lines=1, placeholder="Generated Prompt", show_label=False, info="Auto generated prompts for inspiration.", ) stable_diffusion_inpaint_model_id = gr.Dropdown( choices=list(INPAINT_MODEL_LIST.keys()), value=list(INPAINT_MODEL_LIST.keys())[0], label="Inpaint Model Selection", elem_id="model-dropdown-inpainting", info="Select the model you want to use for inpainting." ) stable_diffusion_inpaint_scheduler = gr.Dropdown( choices=SCHEDULER_LIST, value=SCHEDULER_LIST[0], label="Scheduler", elem_id="scheduler-dropdown-inpainting", info="Scheduler list for models. Different schdulers result in different outputs." ) stable_diffusion_inpaint_guidance_scale = gr.Slider( minimum=0.1, maximum=15, step=0.1, value=7.5, label="Guidance Scale", elem_id = "guidance-scale-slider-inpainting", info = "Guidance scale determines how much the prompt will affect the image. Higher the value, more the effect." ) stable_diffusion_inpaint_num_inference_step = gr.Slider( minimum=1, maximum=100, step=1, value=50, label="Num Inference Step", elem_id = "num-inference-step-slider-inpainting", info = "Number of inference step determines the quality of the image. Higher the number, better the quality." ) stable_diffusion_inpaint_size = gr.Slider( minimum=128, maximum=1280, step=32, value=512, label="Image Size", elem_id="image-size-slider-inpainting", info = "Image size determines the height and width of the generated image. Higher the value, better the quality however slower the computation." ) stable_diffusion_inpaint_seed_generator = gr.Slider( label="Seed(0 for random)", minimum=0, maximum=1000000, value=0, elem_id="seed-slider-inpainting", info="Set the seed to a specific value to reproduce the results." ) stable_diffusion_inpaint_predict = gr.Button( value="Generate image" ) with gr.Column(): output_image = gr.Gallery( label="Generated images", show_label=False, elem_id="gallery-inpainting", ).style(grid=(1, 2)) with gr.Group(elem_id="container-advanced-btns"): with gr.Group(elem_id="share-btn-container"): community_icon_html, loading_icon_html = get_community_loading_icon("inpainting") community_icon = gr.HTML(community_icon_html) loading_icon = gr.HTML(loading_icon_html) share_button = gr.Button("Save artwork", elem_id="share-btn-inpainting") gr.HTML( """

Inpainting Models

Inpainting models will take a masked image and modify the masked image with the given prompt.

Prompt should describe how you want to modify the image. For example, if you want to modify the image to have a blue sky, you can use the prompt "sky is blue".

Negative prompt should describe what you don't want in your image. For example, if you don't want the image to have a red sky, you can use the negative prompt "sky is red".

Stable Diffusion 1 & 2: Default model for many tasks.

""" ) stable_diffusion_inpaint_predict.click( fn=StableDiffusionInpaintGenerator().generate_image, inputs=[ stable_diffusion_inpaint_image_file, stable_diffusion_inpaint_model_id, stable_diffusion_inpaint_prompt, stable_diffusion_inpaint_negative_prompt, stable_diffusion_inpaint_scheduler, stable_diffusion_inpaint_guidance_scale, stable_diffusion_inpaint_num_inference_step, stable_diffusion_inpaint_size, stable_diffusion_inpaint_size, stable_diffusion_inpaint_seed_generator, ], outputs=[output_image], ) stable_diffusion_inpaint_generate.click( fn=generate, inputs=[stable_diffusion_inpaint_prompt], outputs=[stable_diffusion_inpaint_generated_prompt], ) return demo