import spaces import gradio as gr import torch from diffusers import AutoencoderKL, TCDScheduler from diffusers.models.model_loading_utils import load_state_dict from gradio_imageslider import ImageSlider from huggingface_hub import hf_hub_download from controlnet_union import ControlNetModel_Union from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline from PIL import Image, ImageFilter import numpy as np # from gradio.sketch.run import create MODELS = { "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning", "Lustify Lightning": "GraydientPlatformAPI/lustify-lightning", "Juggernaut XL Lightning": "RunDiffusion/Juggernaut-XL-Lightning", "Juggernaut-XL-V9-GE-RDPhoto2": "AiWise/Juggernaut-XL-V9-GE-RDPhoto2-Lightning_4S", "SatPony-Lightning": "John6666/satpony-lightning-v2-sdxl" } # --- ControlNet and Pipeline Setup (Retained) --- config_file = hf_hub_download( "xinsir/controlnet-union-sdxl-1.0", filename="config_promax.json", ) config = ControlNetModel_Union.load_config(config_file) controlnet_model = ControlNetModel_Union.from_config(config) model_file = hf_hub_download( "xinsir/controlnet-union-sdxl-1.0", filename="diffusion_pytorch_model_promax.safetensors", ) state_dict = load_state_dict(model_file) model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model( controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0" ) model.to(device="cuda", dtype=torch.float16) vae = AutoencoderKL.from_pretrained( "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16 ).to("cuda") pipe = StableDiffusionXLFillPipeline.from_pretrained( "SG161222/RealVisXL_V5.0_Lightning", torch_dtype=torch.float16, vae=vae, controlnet=model, variant="fp16", ) pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config) pipe.to("cuda") print(pipe) def load_default_pipeline(): """仅保留，但当前 Inpaint 逻辑未直接使用，可以删除，但保留以防将来扩展。""" global pipe pipe = StableDiffusionXLFillPipeline.from_pretrained( "GraydientPlatformAPI/lustify-lightning", torch_dtype=torch.float16, vae=vae, controlnet=model, ).to("cuda") print("Default pipeline loaded!") @spaces.GPU(duration=15) def fill_image(prompt, image, model_selection, paste_back): """ Handles the fill/repair process for inputs from ImageMask (gr. ImageMask). Applies a default 5% expansion to user-drawn masks here. """ global pipe print(f"Received image: {image}") if image is None: yield None, None return if model_selection in MODELS: current_model = pipe.config.get("_name_or_path", "") target_model = MODELS[model_selection] if current_model != target_model: # 释放旧模型显存 del pipe torch.cuda.empty_cache() pipe = StableDiffusionXLFillPipeline.from_pretrained( target_model, torch_dtype=torch.float16, vae=vae, controlnet=model ) pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config) pipe.to("cuda") print(f"Loaded new SDXL model: {target_model}") ( prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, ) = pipe.encode_prompt(prompt, "cuda", True) source = image["background"] # 用户绘制的 mask layer（通常是 RGBA） mask = image["layers"][0] # 取 alpha 通道并转为二值 mask（255 表示 mask 区域） alpha_channel = mask.split()[3] binary_mask = alpha_channel.point(lambda p: 255 if p > 0 else 0).convert("L") # ==== 扩大 5%（针对 fill_image 的二值 mask） ==== expand_px = max(1, int(min(binary_mask.width, binary_mask.height) * 0.05)) kernel_size = expand_px * 2 + 1 binary_mask = binary_mask.filter(ImageFilter.MaxFilter(kernel_size)) # ==== END 扩大 ==== cnet_image = source.copy() # 在控制网络输入图上把 mask 区域填黑（以便 ControlNet/pipe 根据此区域生成） cnet_image.paste(0, (0, 0), binary_mask) # 调用管线（通常是生成若干中间结果，这里按原逻辑 yield） for image_out in pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, image=cnet_image, # Inpaint 流程使用 image=cnet_image（原图 masked with black）， # 管道内部应该处理了 mask，但如果 StableDiffusionXLFillPipeline # 需要显式 mask，这里可能需要调整。根据原代码的命名和逻辑， # 假定 pipe(image=cnet_image) 适用于此填充流程。 ): yield image_out, cnet_image # 这里的 yield 是为了流式输出 print(f"{model_selection=}") print(f"{paste_back=}") # 最后 paste 回原图（如用户选择） if paste_back: # image_out 是生成的修复部分 # cnet_image 在循环中已被用作 ControlNet 输入图（黑块版） # 这里的 cnet_image 应该更新为 source.copy() 以避免和输入混淆， # 但遵循原代码逻辑，使用 image_out + source/binary_mask # 最终结果是 image_out（修复结果），我们将其粘贴回原图 source # 的非 mask 区域（即只替换 mask 区域） final_output = source.copy() image_out_rgba = image_out.convert("RGBA") # 使用二值 mask 的反转作为 paste 的 mask inverted_mask = binary_mask.point(lambda p: 255 if p == 0 else 0).convert("L") # 将 image_out 粘贴到 final_output 中，仅在 binary_mask 为 255 的区域（即修复区域） final_output.paste(image_out_rgba, (0, 0), binary_mask) yield cnet_image, final_output else: # 如果不 paste back，只返回生成的修复图像 yield cnet_image, image_out def clear_result(): return gr.update(value=None) def use_output_as_input(output_image): """ Receives the output of ImageSlider (image_out, cnet_image) and returns cnet_image as the new input. """ return gr.update(value=output_image[0]) css = """ .nulgradio-container { width: 86vw !important; } .nulcontain { overflow-y: scroll !important; padding: 10px 40px !important; } div#component-17 { height: auto !important; } @media screen and (max-width: 600px) { .img-row{ display: block !important; margin-bottom: 20px !important; } } """ title = """

Diffusers Image Inpaint

Upload an image, draw a mask, and enter a prompt to repair/inpaint the masked area.

Duplicate this Space to skip the queue and enjoy faster inference on the GPU of your choice

""" with gr.Blocks(css=css, fill_height=True) as demo: gr.Markdown(title) with gr.Column(): with gr.Row(): with gr.Column(): prompt = gr.Textbox( label="Prompt", info="Describe what to inpaint the mask with", lines=3, ) with gr.Column(): model_selection = gr.Dropdown( choices=list(MODELS.keys()), value="RealVisXL V5.0 Lightning", label="Model", ) with gr.Row(): run_button = gr.Button("Generate") paste_back = gr.Checkbox(True, label="Paste back original") with gr.Row(equal_height=False): input_image = gr.ImageMask( type="pil", label="Input Image", layers=True, elem_classes="img-row" ) result = ImageSlider( interactive=False, label="Generated Image", elem_classes="img-row" ) use_as_input_button = gr.Button("Use as Input Image", visible=False) # --- Event Handlers for Inpaint --- use_as_input_button.click( fn=use_output_as_input, inputs=[result], outputs=[input_image], queue=False ) # Generates image on button click run_button.click( fn=clear_result, inputs=None, outputs=result, queue=False, ).then( fn=lambda: gr.update(visible=False), inputs=None, outputs=use_as_input_button, queue=False, ).then( fn=fill_image, inputs=[prompt, input_image, model_selection, paste_back], outputs=[result], ).then( fn=lambda: gr.update(visible=True), inputs=None, outputs=use_as_input_button, queue=False, ) # Generates image on prompt submit prompt.submit( fn=clear_result, inputs=None, outputs=result, queue=False, ).then( fn=lambda: gr.update(visible=False), inputs=None, outputs=use_as_input_button, queue=False, ).then( fn=fill_image, inputs=[prompt, input_image, model_selection, paste_back], outputs=[result], ).then( fn=lambda: gr.update(visible=True), inputs=None, outputs=use_as_input_button, queue=False, ) demo.queue(max_size=10).launch(show_error=True)