AI一键消除：手搓一个消除效果达国内Top3的消除工具

chuck
1,110 阅读7分钟

AI一键消除：手搓一个消除效果达国内Top3的消除工具

效果展示

在图像修复领域，以下场景一直是技术难点：

  • 大面积区域的自然重建
  • 复杂纹理的精确还原
  • 多物体交互区域的上下文连贯处理

看看下面的效果对比：

效果对比1左：原图 右：修复效果

传统的修复方法面临诸多挑战：

  • 大面积区域重建需要强大的场景理解能力
  • 精确的纹理还原依赖复杂的特征提取
  • 边缘处理要求细致的上下文融合

而本方案通过创新的多模型协同策略， 上面面临的这些问题无需额外训练即可解决。这种工程化的优化思路，为如何在AGI时代快速提升已有模型效果提供了新的思路。

技术测评

在AGI时代即将到来之际，AI图像处理技术也在飞速发展。 那么如果放眼国内的众多 AI 消除产品，这个效果是否依然能打呢？

下面是三种方案的效果对比：

product-comparison.jpg从左到右：美图秀秀、字节即梦、本方案

从上图可以直观地看出:

  • 美图秀秀的效果较为基础，边缘处理不够自然
  • 字节跳动的即梦整体效果不错，但在细节还原上略有不足
  • 本方案在保持图像整体风格的同时，实现了更自然的修复效果和更好的细节保留
产品效果自然度细节还原速度特点
美图秀秀⭐⭐⭐⭐⭐⭐⭐⭐⭐便捷但效果一般
字节即梦⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐效果稳定均衡
本方案⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐专注高质量修复

正如上面的效果展示，我们的方案虽然在处理速度上稍逊一筹，但在修复质量上却实现了突破性的提升。 值得一提的是，我们的方案，并不需要复杂的模型训练或者进一步finetune。 这对于工程师如何在 AGI 时代通过工程手段提升模型效果，具有一定的参考意义。

核心思想

在图像修复领域，LAMA模型是一个优秀的基础模型。但在实际应用中,我们发现单独使用LAMA存在一些局限性:

  • 当需要修复的区域较大时,修复质量会明显下降
  • 修复区域与原图的边缘衔接往往不够自然
  • 在细节还原方面还有提升空间

为了突破这些限制, 一个很直观的工程侧的解决方案是：多模型协同方案,将几个强大的AI模型优势互补:

  1. 模型串联策略

    • 使用 LAMA 模型生成基础修复结果
    • 将 LAMA 的输出作为 FLUX inpaint 模型的初始化输入
    • 通过 Controlnet 的 upscaler 模型实现两个模型的自然过渡

  2. 精确的局部重绘

    • 采用 alimama 的 Controlnet inpaint 模型进行局部区域的精细重绘
    • 智能识别和处理需要重点修复的区域
    • 确保修复区域与周围环境的无缝融合

核心实现

1. LAMA基础修复

首先使用LAMA模型生成初步的修复结果：

from PIL import Image, ImageFilter, ImageChops
import cv2
import numpy as np
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline as erase_pipe
from modelscope.utils.constant import Tasks

def lama_inpaint(image, mask):
    """
    使用LAMA模型进行图像修复
    
    Args:
        image (PIL.Image): 输入图像
        mask (PIL.Image): 蒙版图像
    
    Returns:
        PIL.Image: 修复后的图像
    """
    # 初始化LAMA模型
    pipe = erase_pipe(
        Tasks.image_inpainting,
        model="damo/cv_fft_inpainting_lama",
        device="cuda"
    )
    
    # 预处理:缩放到最大边896像素
    max_size = 896
    w, h = image.size
    scale = min(max_size / max(w, h), 1.0)
    if scale < 1.0:
        resize_size = (int(w * scale), int(h * scale))
        image = image.resize(resize_size, Image.Resampling.LANCZOS)
        mask = mask.resize(resize_size, Image.Resampling.LANCZOS)
    
    # 执行修复
    result = pipe({"img": image, "mask": mask})
    output_img = Image.fromarray(
        cv2.cvtColor(result[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGB)
    )
    
    # 还原原始尺寸
    if scale < 1.0:
        output_img = output_img.resize((w, h), Image.Resampling.LANCZOS)
        
    return output_img

2. Flux controlnet 模型的准备

虽然我们的目标是在不重新训练模型的前提下，使用工程侧的手段优化模型效果，但是不可否认的一点是，基础模型的好坏将很大程度的影响我们的工程化最终能优化得到的效果。所以，我们优先会选择更好，更强大，更通用的 controlnet 模型。经过测试，最终选择了下面两个 controlnet 模型：

  1. Upscaler Controlnet: 来自 jasperai 的扩图控制模型，专注于提升图像细节质量
  2. Inpainting Controlnet: 阿里妈妈开源的修复控制模型，擅长自然的图像修复

这两个模型的组合使用可以使得无需过多的提示词描述，Flux 即可理解消除任务，并生成高质量的修复结果。

from diffusers import FluxControlNetModel
from diffusers.pipelines import FluxControlNetPipeline

def prepare_controlnet_models(dtype=torch.float16):
    """
    准备和初始化 Controlnet 模型
    
    Args:
        dtype: 模型数据类型，默认为 float16 以优化性能
        
    Returns:
        FluxMultiControlNetModel: 组合后的多控制网络模型
    """
    # 加载上采样控制模型
    upscale_controlnet = FluxControlNetModel.from_pretrained(
        "jasperai/Flux.1-dev-Controlnet-Upscaler",
        torch_dtype=dtype,
    )
    
    # 加载修复控制模型
    inpaint_controlnet = FluxControlNetModel.from_pretrained(
        "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta",
        torch_dtype=dtype,
    )
    
    # 组合两个控制模型
    return FluxMultiControlNetModel([inpaint_controlnet, upscale_controlnet])

3. ControlNet模型的组合与参数调整

为了同时兼顾修复的整体效果和确保可以利用到上面的 lama 算法的结果，我们需要将Flux inpaint controlnet和upscaler controlnet这两个模型进行组合。

不过在实现过程中，我们遇到了一个技术挑战：Flux controlnet inpaint原生并不支持multiple controlnet的组合使用（参考issue: github.com/alimama-cre… 解决，感兴趣的读者可以尝试解决这个开放性问题。

此外，在实际应用中，模型权重的配置对最终效果有着重要影响。基于实践，我们总结出以下优化建议：

  1. Inpainting ControlNet权重配置

    • 建议设置在0.8-1.0之间
    • 这个权重决定了整体结构的重建质量
    • 权重过低可能导致结构不完整，过高则可能产生过度修复

  2. Upscaler ControlNet权重配置

    • 建议控制在0.2-0.4之间
    • 主要用于优化LAMA算法修复后的细节
    • lama 算法的修复结果通常并不理想， 过高会直接影响最终的消除质量

下面是实现这种多模型协同的完整代码示例：

from diffusers import FluxControlNetModel
from diffusers.pipelines import FluxControlNetPipeline
from diffusers import (
    FluxControlNetPipeline,
    FluxControlNetModel,
    FluxMultiControlNetModel,
)

cpu_device = torch.device("cpu")
gpu_device = torch.device("cuda")


def inpaint_with_upscale(
    image,
    control_images,
    mask_image,
    prompt="empty scene",
    num_inference_steps=30,
    guidance_scale=3.5,
    inpaint_strength=1.0,
    upscale_strength=0.3,
):
    """
    使用controlnet inpaint和upscale进行图像修复
    
    Args:
        image: PIL.Image, 输入图像
        control_images: List[PIL.Image], 控制图像列表
        mask_image: PIL.Image, 蒙版图像
        prompt: str, 正向提示
        num_inference_steps: int, 推理步数
        guidance_scale: float, 引导比例
        inpaint_strength: float, inpaint controlnet的权重
        upscale_strength: float, upscale controlnet的权重
    
    Returns:
        PIL.Image: 修复后的图像
    """
    # 组合两个controlnet
    controlnet = FluxMultiControlNetModel([inpaint_controlnet, upscale_controlnet])
    
    # 创建pipeline
    pipe = FluxControlNetPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev",
        torch_dtype=dtype
    )
    pipe.controlnet = controlnet
    
    # 将模型移到GPU
    pipe.to(gpu_device)
    controlnet.to(gpu_device)
    
    try:
        # 执行推理
        output = pipe(
            prompt=prompt,
            image=image,
            control_mask=mask_image,
            control_image=control_images,
            controlnet_conditioning_scale=[inpaint_strength, upscale_strength],
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            control_mode=[None, None],
            true_guidance_scale=1.0
        ).images[0]
        
        return output
        
    finally:
        # 清理显存
        pipe.to(cpu_device)
        controlnet.to(cpu_device)
        torch.cuda.empty_cache()

In the end

通过上述技术方案的详细阐述，我们看到了如何通过工程化手段将多个强大的AI模型进行优势互补，在不额外训练的情况下显著提升图像修复效果。我已经将这一方案实现为在线工具：

dy.dayanweilai.com/#/airemove/…

这个工具不仅完整实现了上述的技术方案，还 增加了多项工程优化的 tricky，进一步提升了修复效果。工具使用方式非常直观 - 只需上传图片并标记需要修复的区域，AI就会自动完成高质量的图像修复。

进一步展望，图像修复技术仍有很大的优化空间。 比如现在图像修复的单图耗时需要 10s 左右，有没有办法进一步优化性能？比如 differential-diffusion 项目提出了一种创新的扩散模型推理方案，可以让重绘区域与周围环境实现更自然的过渡。如果能将这种差分扩散方案与我们的多模型协同方案相结合，是否能够进一步提升修复效果的自然度？这些开放性问题就留给聪明的读者进一步探索。

本文使用 markdown.com.cn 排版

avatar
文章
阅读
粉丝