自定义脚本编写与集成

摘要

Stable Diffusion WebUI 的脚本系统是其核心功能之一，允许用户通过编写自定义脚本来扩展和修改图像生成流程。本文将深入探讨 WebUI 脚本系统的架构、工作机制以及如何开发自定义脚本。我们将分析脚本生命周期、各类回调函数的作用，以及如何通过脚本与 UI 进行交互。此外，还会介绍如何利用脚本回调系统来扩展 WebUI 功能，为开发者提供全面的开发指导。

关键词： Stable Diffusion WebUI, 脚本系统, 自定义脚本, Python, 回调函数

1. 引言

Stable Diffusion WebUI 不仅是一个图像生成工具，更是一个高度可扩展的平台。其脚本系统是实现这种可扩展性的核心机制之一。通过脚本，用户可以：

1. 修改图像生成流程的各个阶段
2. 添加自定义 UI 控件
3. 实现复杂的图像后处理功能
4. 与其他系统集成

脚本系统的设计理念是提供一种非侵入式的方式来扩展 WebUI 功能，让用户能够在不修改核心代码的前提下，实现个性化的功能需求。

本文将从以下几个方面详细介绍脚本系统的开发：

1. 脚本系统架构概览
2. 脚本生命周期和核心方法
3. UI 交互和控件创建
4. 回调系统详解
5. 实际开发示例

2. 脚本系统架构概览

2.1 核心组件

WebUI 的脚本系统主要由以下几个核心组件构成：

1. Script 基类：所有自定义脚本都需要继承的基类
2. ScriptRunner：负责管理和执行脚本的运行时环境
3. 回调系统：提供扩展点的回调机制
4. UI 系统：与 Gradio 集成的用户界面系统

2.2 脚本分类

WebUI 中的脚本分为两种主要类型：

1. Selectable Scripts（可选脚本）：用户可以在下拉菜单中选择并配置的脚本
2. AlwaysVisible Scripts（常显脚本）：始终处于激活状态的脚本

这两种脚本的主要区别在于它们的显示方式和激活机制。

2.3 脚本执行环境

脚本在 WebUI 中有两种执行环境：

1. txt2img 环境：文本到图像生成环境
2. img2img 环境：图像到图像生成环境

每种环境下，脚本可能有不同的行为和配置选项。

3. 脚本生命周期和核心方法

3.1 Script 基类详解

所有的自定义脚本都必须继承 [Script](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L34-L337) 基类。这个基类定义了一系列的方法，这些方法会在图像生成的不同阶段被调用。

class Script:
    name = None
    section = None
    filename = None
    args_from = None
    args_to = None
    alwayson = False
    
    is_txt2img = False
    is_img2img = False
    tabname = None
    
    group = None
    create_group = True
    
    infotext_fields = None
    paste_field_names = None
    api_info = None
    
    on_before_component_elem_id = None
    on_after_component_elem_id = None
    setup_for_ui_only = False
    controls = None

3.2 核心方法详解

必须实现的方法

def title(self):
    """返回脚本的标题，会显示在下拉菜单中"""
    raise NotImplementedError()

def ui(self, is_img2img):
    """创建 Gradio UI 控件，返回控件列表"""
    pass

def show(self, is_img2img):
    """决定脚本在 UI 中的显示方式"""
    return True

def run(self, p, *args):
    """当脚本被选中执行时调用此方法"""
    pass

可选的生命周期方法

def setup(self, p, *args):
    """在处理对象设置完成后但在任何处理开始之前调用"""
    pass

def before_process(self, p, *args):
    """在处理开始前很早的时候调用"""
    pass

def process(self, p, *args):
    """在处理开始前调用"""
    pass

def before_process_batch(self, p, *args, **kwargs):
    """在处理每个批次前调用"""
    pass

def process_batch(self, p, *args, **kwargs):
    """处理每个批次时调用"""
    pass

def postprocess_batch(self, p, *args, **kwargs):
    """在每个批次生成后调用"""
    pass

def postprocess(self, p, processed, *args):
    """在整个处理结束后调用"""
    pass

def before_component(self, component, **kwargs):
    """在创建组件前调用"""
    pass

def after_component(self, component, **kwargs):
    """在创建组件后调用"""
    pass

3.3 脚本执行流程

脚本的执行流程大致如下：

1. 用户在 UI 中选择脚本或激活 AlwaysVisible 脚本
2. 调用 [setup()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L133-L138) 方法进行初始化
3. 调用 [before_process()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L140-L149) 进行预处理
4. 调用 [process()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L151-L160) 处理生成参数
5. 对每个批次调用 [before_process_batch()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L162-L173) 和 [process_batch()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L188-L199)
6. 生成图像
7. 对每个批次调用 [postprocess_batch()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L201-L213)
8. 调用 [postprocess()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L236-L244) 进行后处理

4. UI 交互和控件创建

4.1 创建 UI 控件

在 [ui()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L125-L131) 方法中，你可以创建任意的 Gradio 控件。这些控件的值将在脚本执行时传递给其他方法。

def ui(self, is_img2img):
    """创建 UI 控件"""
    # 创建滑块控件
    strength = gr.Slider(
        minimum=0.0,
        maximum=1.0,
        step=0.01,
        label="处理强度",
        value=0.5
    )
    
    # 创建复选框
    enable = gr.Checkbox(
        label="启用功能",
        value=False
    )
    
    # 返回控件列表，顺序很重要
    return [strength, enable]

4.2 控件值传递

UI 控件的值会按照它们在 [ui()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L125-L131) 方法返回列表中的顺序，作为参数传递给脚本的各个方法：

def process(self, p, strength, enable):
    """处理方法接收 UI 控件的值"""
    if not enable:
        return
        
    # 使用 strength 参数进行处理
    p.cfg_scale *= (1 + strength)

4.3 自定义组件注入

脚本还可以通过 [before_component()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L266-L278) 和 [after_component()](file:///e:/project/stable-diffusion-webui/modules/scripts.py#L280-L288) 方法向现有 UI 中注入自定义组件：

def before_component(self, component, **kwargs):
    """在组件创建前调用"""
    # 检查是否是我们感兴趣的组件
    if kwargs.get("elem_id") == "txt2img_generate":
        # 在生成按钮前添加我们的组件
        with gr.Accordion("高级选项"):
            gr.Markdown("这里是自定义内容")

def after_component(self, component, **kwargs):
    """在组件创建后调用"""
    pass

5. 回调系统详解

5.1 回调系统架构

WebUI 提供了一个强大的回调系统，允许脚本注册函数以响应系统中的各种事件。这些回调函数定义在 [modules/script_callbacks.py](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py) 文件中。

主要的回调类型包括：

1. UI 相关回调：如 [on_ui_tabs()](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py#L545-L555)、[on_ui_settings()](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py#L575-L580)
2. 处理流程回调：如 [on_before_image_saved()](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py#L585-L592)、[on_image_saved()](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py#L594-L600)
3. 模型相关回调：如 [on_model_loaded()](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py#L535-L540)
4. 组件相关回调：如 [on_before_component()](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py#L607-L613)、[on_after_component()](file:///e:/project/stable-diffusion-webui/modules/script_callbacks.py#L615-L618)

5.2 常用回调函数

UI 标签页回调

from modules import script_callbacks
import gradio as gr

def on_ui_tabs():
    with gr.Blocks(analytics_enabled=False) as demo:
        gr.Markdown("# 自定义标签页")
        # 添加自定义 UI 组件
        
    return [(demo, "自定义标签", "custom_tab")]

# 注册回调
script_callbacks.on_ui_tabs(on_ui_tabs)

图像保存回调

from modules import script_callbacks

def before_image_saved(params: script_callbacks.ImageSaveParams):
    """在图像保存前调用"""
    # 可以修改图像或其他参数
    params.pnginfo["CustomParam"] = "CustomValue"

def after_image_saved(params: script_callbacks.ImageSaveParams):
    """在图像保存后调用"""
    print(f"图像已保存至: {params.filename}")

# 注册回调
script_callbacks.on_before_image_saved(before_image_saved)
script_callbacks.on_image_saved(after_image_saved)

模型加载回调

from modules import script_callbacks

def on_model_loaded(sd_model):
    """在模型加载完成后调用"""
    print(f"模型已加载: {sd_model}")

# 注册回调
script_callbacks.on_model_loaded(on_model_loaded)

6. 实际开发示例

6.1 简单的图像后处理脚本

下面是一个简单的图像后处理脚本示例，它会在图像生成后应用亮度调整：

import modules.scripts as scripts
import gradio as gr
import numpy as np
from PIL import Image

class Script(scripts.Script):
    def title(self):
        return "亮度调整"

    def show(self, is_img2img):
        # 在两个环境中都显示
        return scripts.AlwaysVisible

    def ui(self, is_img2img):
        # 创建亮度调整滑块
        brightness = gr.Slider(
            minimum=-0.5,
            maximum=0.5,
            step=0.01,
            label="亮度调整",
            value=0.0
        )
        return [brightness]

    def postprocess(self, p, processed, brightness):
        # 后处理方法
        if brightness == 0:
            return
            
        # 调整每张图像的亮度
        for i in range(len(processed.images)):
            img = processed.images[i]
            
            # 将 PIL 图像转换为 numpy 数组
            if isinstance(img, Image.Image):
                img_array = np.array(img).astype(np.float32)
            else:
                img_array = img.astype(np.float32)
                
            # 调整亮度
            img_array += brightness * 255
            img_array = np.clip(img_array, 0, 255).astype(np.uint8)
            
            # 转换回 PIL 图像
            if isinstance(img, Image.Image):
                processed.images[i] = Image.fromarray(img_array)
            else:
                processed.images[i] = img_array

6.2 复杂的自定义脚本

下面是一个更复杂的脚本示例，它实现了图像水印功能：

import modules.scripts as scripts
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont

class WatermarkScript(scripts.Script):
    def title(self):
        return "图像水印"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def ui(self, is_img2img):
        with gr.Group():
            with gr.Row():
                enable = gr.Checkbox(label="启用水印", value=False)
                watermark_text = gr.Textbox(label="水印文字", value="AI Generated")
            with gr.Row():
                font_size = gr.Slider(minimum=10, maximum=100, step=1, label="字体大小", value=30)
                opacity = gr.Slider(minimum=0, maximum=1, step=0.05, label="透明度", value=0.5)
            with gr.Row():
                position = gr.Dropdown(
                    choices=["左上", "右上", "左下", "右下", "中央"],
                    label="位置",
                    value="右下"
                )
                
        return [enable, watermark_text, font_size, opacity, position]

    def postprocess(self, p, processed, enable, watermark_text, font_size, opacity, position):
        if not enable or not watermark_text:
            return
            
        # 为每张图像添加水印
        for i in range(len(processed.images)):
            if isinstance(processed.images[i], Image.Image):
                watermarked = self.add_watermark(
                    processed.images[i], 
                    watermark_text, 
                    font_size, 
                    opacity, 
                    position
                )
                processed.images[i] = watermarked

    def add_watermark(self, image, text, font_size, opacity, position):
        # 创建水印图像
        watermark = Image.new('RGBA', image.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(watermark)
        
        try:
            # 尝试使用系统字体
            font = ImageFont.truetype("arial.ttf", font_size)
        except:
            # 如果没有找到字体，使用默认字体
            font = ImageFont.load_default()
            
        # 计算文本尺寸
        bbox = draw.textbbox((0, 0), text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        
        # 根据位置确定坐标
        if position == "左上":
            x, y = 10, 10
        elif position == "右上":
            x, y = image.size[0] - text_width - 10, 10
        elif position == "左下":
            x, y = 10, image.size[1] - text_height - 10
        elif position == "右下":
            x, y = image.size[0] - text_width - 10, image.size[1] - text_height - 10
        else:  # 中央
            x, y = (image.size[0] - text_width) // 2, (image.size[1] - text_height) // 2
            
        # 绘制半透明文本
        draw.text((x, y), text, font=font, fill=(255, 255, 255, int(255 * opacity)))
        
        # 将水印合并到原图
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
            
        watermarked = Image.alpha_composite(image, watermark)
        return watermarked.convert('RGB')

6.3 利用回调系统扩展功能

下面是一个使用回调系统的示例，它会在 WebUI 启动时添加自定义设置：

import modules.scripts as scripts
from modules import script_callbacks, shared

def on_ui_settings():
    # 添加自定义设置项
    shared.opts.add_option(
        "custom_watermark_default_text",
        shared.OptionInfo(
            "AI Generated",
            "默认水印文字",
            section=("custom", "自定义设置")
        )
    )

def on_app_started(demo, app):
    print("WebUI 已启动，自定义脚本已加载")

# 注册回调
script_callbacks.on_ui_settings(on_ui_settings)
script_callbacks.on_app_started(on_app_started)

7. 最佳实践和注意事项

7.1 性能优化

1. 避免重复计算：在处理大量图像时，尽量缓存计算结果
2. 使用适当的数据类型：在图像处理中，合理使用 numpy 数组可以提高性能
3. 异步处理：对于耗时操作，考虑使用异步方式处理

7.2 错误处理

1. 异常捕获：在关键方法中添加适当的异常处理
2. 日志记录：使用适当的日志记录错误信息
3. 优雅降级：在出现错误时提供合理的默认行为

def process(self, p, *args):
    try:
        # 主要处理逻辑
        pass
    except Exception as e:
        # 记录错误但不中断流程
        print(f"脚本处理出错: {e}")
        # 可以选择继续处理或返回

7.3 兼容性考虑

1. 版本兼容：检查 WebUI 版本并在不兼容时给出提示
2. 环境检测：区分 txt2img 和 img2img 环境
3. 依赖管理：清楚地声明外部依赖

7.4 用户体验

1. 清晰的 UI：提供直观易懂的用户界面
2. 合理的默认值：设置合适的默认参数
3. 详细的文档：编写清晰的使用说明

8. 总结

WebUI 的脚本系统为用户提供了强大的扩展能力，使其成为一个真正可定制的平台。通过理解脚本的生命周期、掌握 UI 交互方法、熟练运用回调系统，开发者可以创建出功能丰富的自定义脚本。

本文介绍了脚本系统的核心概念和实际应用方法，通过具体示例展示了如何开发不同类型的脚本。随着对系统理解的加深，开发者可以创造出更加复杂和有用的功能，进一步丰富 WebUI 的生态系统。

脚本系统不仅是 WebUI 的重要特性，也是其社区活跃发展的基础。通过这个系统，无数优秀的功能得以实现和分享，形成了一个庞大的插件生态，这也是 Stable Diffusion WebUI 能够持续发展和壮大的重要原因。

参考资料

1. Stable Diffusion WebUI 官方文档
2. Gradio 框架文档
3. Python 标准库文档
4. Pillow 图像处理库文档
5. NumPy 科学计算库文档