app_inpaint2.py

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(
        "SG161222/RealVisXL_V5.0_Lightning",
        torch_dtype=torch.float16,
        vae=vae,
        controlnet=model,
    ).to("cuda")
    # 此函数原用于 Misc 选项卡，现在无需返回 gr.update
    print("Default pipeline loaded!")

@spaces.GPU(duration=7)
def fill_image(prompt, image, model_selection, paste_back):
    """
    处理 ImageMask（gr.ImageMask）输入的 fill/repair 流程。
    在这里对用户绘制的 mask 做默认 5% 膨胀。
    """
    print(f"Received image: {image}")
    if image is None:
        yield None, None
        return

    # 如果用户选择了不同的模型 key，则加载对应预训练仓库
    # 注意：此逻辑原 Outpaint 中有，Inpaint 中缺失，现在补充以支持模型切换
    global pipe
    if model_selection in MODELS and pipe.config.model_name != MODELS[model_selection]:
        # 释放旧模型显存
        del pipe
        torch.cuda.empty_cache()
        pipe = StableDiffusionXLFillPipeline.from_pretrained(
            MODELS[model_selection],
            torch_dtype=torch.float16,
            vae=vae,
            controlnet=model,
            variant="fp16", # 保持 variant 设置
        )
        pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
        pipe.to("cuda")
        print(f"Loaded new SDXL model: {pipe.config.model_name}")


    (
        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 final_output, cnet_image
    else:
        # 如果不 paste back，只返回生成的修复图像
        yield image_out, cnet_image

def clear_result():
    return gr.update(value=None)

def use_output_as_input(output_image):
    """
    接收 ImageSlider 的输出 (image_out, cnet_image)
    返回 cnet_image 作为新的输入。
    """
    # ImageSlider 的 value 是一个 tuple (image1, image2)
    # 这里的 image2 (即 cnet_image) 是包含修复结果的图像 (如果 paste_back 为 True)
    # 或者只是 ControlNet 输入（如果 fill_image 逻辑有变）
    # 假设我们想要将修复后的结果图作为新的输入图像（新的 source）
    # 在 fill_image 中，最终 yield 的是 (final_output, cnet_image)
    # 我们应该使用 final_output 作为新的背景图。
    return gr.update(value=output_image[0]) # output_image[0] 是最终修复图像

css = """
.nulgradio-container {
    width: 86vw !important;
}
.nulcontain {
    overflow-y: scroll !important;
    padding: 10px 40px !important;
}
div#component-17 { /* 这是一个动态 ID，可能需要调整或移除 */
    height: auto !important;
}

@media screen and (max-width: 600px) {
    .img-row{
        display: block !important;
        margin-bottom: 20px !important;
    }
    /* 移除掉 component-16 的引用，因为它是动态的 */
}

"""

title = """<h1 align="center">Diffusers Image Inpaint</h1>
<div align="center">Upload an image, draw a mask, and enter a prompt to repair/inpaint the masked area.</div>
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
    <p style="display: flex;gap: 6px;">
         <a href="https://huggingface.co/spaces/fffiloni/diffusers-image-outpout?duplicate=true">
            <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-md.svg" alt="Duplicate this Space">
        </a> to skip the queue and enjoy faster inference on the GPU of your choice
    </p>
</div>
"""

with gr.Blocks(css=css, fill_height=True) as demo:
    gr.Markdown(title) # 使用更简洁的 Markdown 标题
    
    # 只保留 Inpaint 选项卡的内容，移除 Tabs 结构，让 Inpaint 成为主界面
    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,
        )

# 将 queue 和 launch 保持不变
demo.queue(max_size=10).launch(show_error=True)