rapid12k4/diffusers-fast-inpaint

1import gradio as gr2import spaces3import torch4from diffusers import AutoencoderKL, TCDScheduler5from diffusers.models.model_loading_utils import load_state_dict6from gradio_imageslider import ImageSlider7from huggingface_hub import hf_hub_download8 9from controlnet_union import ControlNetModel_Union10from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline11 12MODELS = {13    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",14}15 16config_file = hf_hub_download(17    "xinsir/controlnet-union-sdxl-1.0",18    filename="config_promax.json",19)20 21config = ControlNetModel_Union.load_config(config_file)22controlnet_model = ControlNetModel_Union.from_config(config)23model_file = hf_hub_download(24    "xinsir/controlnet-union-sdxl-1.0",25    filename="diffusion_pytorch_model_promax.safetensors",26)27state_dict = load_state_dict(model_file)28model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(29    controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"30)31model.to(device="cuda", dtype=torch.float16)32 33vae = AutoencoderKL.from_pretrained(34    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float1635).to("cuda")36 37pipe = StableDiffusionXLFillPipeline.from_pretrained(38    "SG161222/RealVisXL_V5.0_Lightning",39    torch_dtype=torch.float16,40    vae=vae,41    controlnet=model,42    variant="fp16",43).to("cuda")44 45pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)46 47 48@spaces.GPU(duration=24)49def fill_image(prompt, image, model_selection, paste_back):50    (51        prompt_embeds,52        negative_prompt_embeds,53        pooled_prompt_embeds,54        negative_pooled_prompt_embeds,55    ) = pipe.encode_prompt(prompt, "cuda", True)56 57    source = image["background"]58    mask = image["layers"][0]59 60    alpha_channel = mask.split()[3]61    binary_mask = alpha_channel.point(lambda p: p > 0 and 255)62    cnet_image = source.copy()63    cnet_image.paste(0, (0, 0), binary_mask)64 65    for image in pipe(66        prompt_embeds=prompt_embeds,67        negative_prompt_embeds=negative_prompt_embeds,68        pooled_prompt_embeds=pooled_prompt_embeds,69        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,70        image=cnet_image,71    ):72        yield image, cnet_image73 74    print(f"{model_selection=}")75    print(f"{paste_back=}")76 77    if paste_back:78        image = image.convert("RGBA")79        cnet_image.paste(image, (0, 0), binary_mask)80    else:81        cnet_image = image82 83    yield source, cnet_image84 85 86def clear_result():87    return gr.update(value=None)88 89 90title = """<h1 align="center">Diffusers Fast Inpaint</h1>91<div align="center">Draw the mask over the subject you want to erase or change and write what you want to inpaint it with.</div>92<div align="center">This is a lighting model with almost no CFG and 12 steps, so don't expect high quality generations.</div>93<div align="center">This space is a PoC made for the guide <a href='https://huggingface.co/blog/OzzyGT/diffusers-image-fill'>Diffusers Image Fill</a>.</div>94"""95 96with gr.Blocks() as demo:97    gr.HTML(title)98    with gr.Row():99        with gr.Column():100            prompt = gr.Textbox(101                label="Prompt",102                info="Describe what to inpaint the mask with",103                lines=3,104            )105        with gr.Column():106            model_selection = gr.Dropdown(107                choices=list(MODELS.keys()),108                value="RealVisXL V5.0 Lightning",109                label="Model",110            )111 112            with gr.Row():113                with gr.Column():114                    run_button = gr.Button("Generate")115 116                with gr.Column():117                    paste_back = gr.Checkbox(True, label="Paste back original")118 119    with gr.Row():120        input_image = gr.ImageMask(121            type="pil", label="Input Image", crop_size=(1024, 1024), layers=False122        )123 124        result = ImageSlider(125            interactive=False,126            label="Generated Image",127        )128 129    use_as_input_button = gr.Button("Use as Input Image", visible=False)130 131    def use_output_as_input(output_image):132        return gr.update(value=output_image[1])133 134    use_as_input_button.click(135        fn=use_output_as_input, inputs=[result], outputs=[input_image]136    )137 138    run_button.click(139        fn=clear_result,140        inputs=None,141        outputs=result,142    ).then(143        fn=lambda: gr.update(visible=False),144        inputs=None,145        outputs=use_as_input_button,146    ).then(147        fn=fill_image,148        inputs=[prompt, input_image, model_selection, paste_back],149        outputs=result,150    ).then(151        fn=lambda: gr.update(visible=True),152        inputs=None,153        outputs=use_as_input_button,154    )155 156    prompt.submit(157        fn=clear_result,158        inputs=None,159        outputs=result,160    ).then(161        fn=lambda: gr.update(visible=False),162        inputs=None,163        outputs=use_as_input_button,164    ).then(165        fn=fill_image,166        inputs=[prompt, input_image, model_selection, paste_back],167        outputs=result,168    ).then(169        fn=lambda: gr.update(visible=True),170        inputs=None,171        outputs=use_as_input_button,172    )173 174 175demo.queue(max_size=12).launch(share=False)176