Spaces:
Sleeping
Sleeping
| 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, | |
| num_images_per_prompt: int, | |
| 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=num_images_per_prompt, | |
| 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", | |
| show_label=False, | |
| elem_id="prompt-text-input-inpainting", | |
| value='' | |
| ) | |
| stable_diffusion_inpaint_negative_prompt = gr.Textbox( | |
| lines=1, | |
| placeholder="Negative Prompt", | |
| 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, | |
| ) | |
| 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", | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| 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" | |
| ) | |
| 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" | |
| ) | |
| stable_diffusion_inpiant_num_images_per_prompt = gr.Slider( | |
| minimum=1, | |
| maximum=10, | |
| step=1, | |
| value=1, | |
| label="Number Of Images", | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| stable_diffusion_inpaint_scheduler = gr.Dropdown( | |
| choices=SCHEDULER_LIST, | |
| value=SCHEDULER_LIST[0], | |
| label="Scheduler", | |
| elem_id="scheduler-dropdown-inpainting", | |
| ) | |
| stable_diffusion_inpaint_size = gr.Slider( | |
| minimum=128, | |
| maximum=1280, | |
| step=32, | |
| value=512, | |
| label="Image Size", | |
| elem_id="image-size-slider-inpainting", | |
| ) | |
| stable_diffusion_inpaint_seed_generator = gr.Slider( | |
| label="Seed(0 for random)", | |
| minimum=0, | |
| maximum=1000000, | |
| value=0, | |
| elem_id="seed-slider-inpainting", | |
| ) | |
| stable_diffusion_inpaint_predict = gr.Button( | |
| value="Generator" | |
| ) | |
| 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") | |
| 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_inpiant_num_images_per_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 | |