Stable Diffusion中的各种插件功能是怎么实现的

起因

之前用gradio搭一些项目的时候一般只有1~2个tab页,官方的技术文档已能满足需求,因此对组件的拓展性没有太关注。最近在研究Stable Diffusion,发现绝大部分web ui的实现都非常全面,一般都有十几个tab页,而且把鼠标移动到组件上,会发现有文字提示,如下图所示,这肯定不是原生gradio的组件。这引起了我的好奇心,gradio应该是比较简单,主要用于demo展示的,那这些花里胡哨的组件是怎么实现的?

探究

通过搜索相关的资料,发现大多插件都是以github仓库的形式存在的,那这就好办了,通过一个简单的项目即可窥斑见豹。在这里我找到一个插件:github.com/butaixianra... 这个是一个实现prompt翻译的插件,目录结构如下:

lua 复制代码
--img # demo用到的图片
|-- xxx.jpg
--javascript # js函数定义
|-- prompt_translator.js
--scripts
|-- services # 谷歌翻译相关的模块,这里以百度翻译为例,因此无需细究
  |-- __init__.py
  |-- google.py
  |-- interface.py
  |-- schema.py
|-- lang_code.py # 不同的翻译api中,语言和对应缩写。如:中文 对应 zh
|-- prompt_translator.py # 核心功能模块
--.gitignore
--LICENSE
--README.cn.md
--README.md

各个模块的功能如上所述,这里只摘取prompt_translator.py进行说明

分析

以下分析针对prompt_translator.py中的关键变量、函数进行说明。

1 从原始web-ui模块中引入的4个组件

这4个组件不是在这个项目中搭建的,而是从原始web-ui项目中引入(几乎所有开源的sd-web-ui实现都是在原始项目的基础上开发的),而这4个组件则代表最原始的功能,即文生图的prompt、文生图的负向prompt、图生图的prompt、图生图的负向prompt。当然原始项目还有很多其他组件,主要是因为这个插件是实现翻译,因此只需要引入这几个即可。

python 复制代码
    # 从原始web-ui项目中引入的4个组件    
    txt2img_prompt = modules.ui.txt2img_paste_fields[0][0]
    txt2img_neg_prompt = modules.ui.txt2img_paste_fields[1][0]
    img2img_prompt = modules.ui.img2img_paste_fields[0][0]
    img2img_neg_prompt = modules.ui.img2img_paste_fields[1][0]

2 隐藏的按钮

trans_prompt_js_btntrans_neg_prompt_js_btn 这两个按钮隐藏起来了,可以看到其初始化的属性:visible=False。因此对于该项目搭建的gradio app来说,这两个按钮是没用的,对应的绑定的两个监听函数也是没用的。 用更简单的话说,如果用该项目启动gradio app,不会出现这两个按钮。如单独部署的界面。

python 复制代码
# gradio搭建 
    with gr.Blocks(analytics_enabled=False) as prompt_translator:
        # ====ui====        
        # Prompt Area
        with gr.Row():
            tar_lang_drop = gr.Dropdown(label="Target Language", choices=tar_langs, value=def_tar_lang, elem_id="pt_tar_lang")
        with gr.Row():
            prompt = gr.Textbox(label="Prompt", lines=3, value="", elem_id="pt_prompt")
            translated_prompt = gr.Textbox(label="Translated Prompt", lines=3, value="", elem_id="pt_translated_prompt")

        with gr.Row():
            trans_prompt_btn = gr.Button(value="Translate", elem_id="pt_trans_prompt_btn")
            # add a hidden button, used by fake click with javascript. To simulate msg between server and client side.
            # this is the only way.
            # 隐藏的按钮,用于js触发
            trans_prompt_js_btn = gr.Button(value="Trans Js", visible=False, elem_id="pt_trans_prompt_js_btn")
            send_prompt_btn = gr.Button(value="Send to txt2img and img2img", elem_id="pt_send_prompt_btn")


        with gr.Row():
            neg_prompt = gr.Textbox(label="Negative Prompt", lines=2, value="", elem_id="pt_neg_prompt")
            translated_neg_prompt = gr.Textbox(label="Translated Negative Prompt", lines=2, value=json.dumps(tar_langs), elem_id="pt_translated_neg_prompt")

        with gr.Row():
            trans_neg_prompt_btn = gr.Button(value="Translate", elem_id="pt_trans_neg_prompt_btn")
            # add a hidden button, used by fake click with javascript. To simulate msg between server and client side.
            # this is the only way.
             # 隐藏的按钮,用于js触发
            trans_neg_prompt_js_btn = gr.Button(value="Trans Js", visible=False, elem_id="pt_trans_neg_prompt_js_btn")
            send_neg_prompt_btn = gr.Button(value="Send to txt2img and img2img", elem_id="pt_send_neg_prompt_btn")

3 交互

上述代码再往后看,会发现该模块最后有这么一行代码:

bash 复制代码
script_callbacks.on_ui_tabs(on_ui_tabs)

经过对比,发现这是用了原始web-ui项目中提供的一个回调方法,实现插件功能,即把这个插件页面补充到原始web-ui项目中去。

4 原始项目中的方法

查看回原始项目(github.com/AUTOMATIC11...),发现2个模块比较关键。 对于modules/script_callbacks.py,定义了一些把需要新增的插件补充进原始项目中去的方法,这里仅列出2个跟本翻译插件项目关联最紧密的方法说明。

python 复制代码
# file: modules/script_callbacks.py

# 把当前的插件添加进callback_map字典中
def on_ui_tabs(callback):
    """register a function to be called when the UI is creating new tabs.
    The function must either return a None, which means no new tabs to be added, or a list, where
    each element is a tuple:
        (gradio_component, title, elem_id)

    gradio_component is a gradio component to be used for contents of the tab (usually gr.Blocks)
    title is tab text displayed to user in the UI
    elem_id is HTML id for the tab
    """
    add_callback(callback_map['callbacks_ui_tabs'], callback)
    
   
# 遍历callback_map字典,把新增的插件都添加进来
def ui_tabs_callback():
    res = []

    for c in callback_map['callbacks_ui_tabs']:
        try:
            res += c.callback() or []
        except Exception:
            report_exception(c, 'ui_tabs_callback')

    return res

对于modules/ui.py,即实现了原始项目的核心gradio界面功能,通常初始化6个tab页面,然后通过上述回调方法,把各种需要的插件功能添加进原始项目中去。

python 复制代码
# file: modules/ui.py
# ...
# 前面部分与通常的gradio搭建大同小异
# 初始的6个tab页
interfaces = [
    (txt2img_interface, "txt2img", "txt2img"),
    (img2img_interface, "img2img", "img2img"),
    (extras_interface, "Extras", "extras"),
    (pnginfo_interface, "PNG Info", "pnginfo"),
    (modelmerger_ui.blocks, "Checkpoint Merger", "modelmerger"),
    (train_interface, "Train", "train"),
]
# 通过ui_tabs_callback方法实现tab的拓展,这些拓展的tab即插件,如本文说的翻译插件
interfaces += script_callbacks.ui_tabs_callback()
interfaces += [(settings.interface, "Settings", "settings")]

总结

通过一个翻译插件项目,我了解到怎么在原始的的gradio项目中添加额外的插件功能,而无需进行太多改动。而且这些插件功能通过js可以实现各种额外的功能,大大弥补了gradio的弱点。

相关推荐
qzhqbb1 小时前
基于统计方法的语言模型
人工智能·语言模型·easyui
冷眼看人间恩怨2 小时前
【话题讨论】AI大模型重塑软件开发:定义、应用、优势与挑战
人工智能·ai编程·软件开发
2401_883041082 小时前
新锐品牌电商代运营公司都有哪些?
大数据·人工智能
AI极客菌3 小时前
Controlnet作者新作IC-light V2:基于FLUX训练,支持处理风格化图像,细节远高于SD1.5。
人工智能·计算机视觉·ai作画·stable diffusion·aigc·flux·人工智能作画
阿_旭3 小时前
一文读懂| 自注意力与交叉注意力机制在计算机视觉中作用与基本原理
人工智能·深度学习·计算机视觉·cross-attention·self-attention
王哈哈^_^3 小时前
【数据集】【YOLO】【目标检测】交通事故识别数据集 8939 张,YOLO道路事故目标检测实战训练教程!
前端·人工智能·深度学习·yolo·目标检测·计算机视觉·pyqt
Power20246664 小时前
NLP论文速读|LongReward:基于AI反馈来提升长上下文大语言模型
人工智能·深度学习·机器学习·自然语言处理·nlp
数据猎手小k4 小时前
AIDOVECL数据集:包含超过15000张AI生成的车辆图像数据集,目的解决旨在解决眼水平分类和定位问题。
人工智能·分类·数据挖掘
好奇龙猫4 小时前
【学习AI-相关路程-mnist手写数字分类-win-硬件:windows-自我学习AI-实验步骤-全连接神经网络(BPnetwork)-操作流程(3) 】
人工智能·算法
沉下心来学鲁班4 小时前
复现LLM:带你从零认识语言模型
人工智能·语言模型