stable-diffusion-webui之extension

https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Developing-extensionshttps://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Developing-extensionshttps://github.com/udon-universe/stable-diffusion-webui-extension-templateshttps://github.com/udon-universe/stable-diffusion-webui-extension-templateDeveloping custom scripts · AUTOMATIC1111/stable-diffusion-webui Wiki · GitHubStable Diffusion web UI. Contribute to AUTOMATIC1111/stable-diffusion-webui development by creating an account on GitHub.https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Developing-custom-scripts主要是对scripts的原理进行分析。

extension中：

1.install.py一定会被执行。环境安装都在launch.py中，在launch.py/run_extension_installer=launch_utils.run_extension_installer中执行。

import launch
 
if not launch.is_installed("aitextgen"):
    launch.run_pip("install aitextgen==0.6.0", "requirements for MagicPrompt")

2.extension中一定要包括scripts目录下脚本，extension目录已被包含在sys.path中。

Script类在modules/scripts.py中，

import modules.scripts as scripts
import gradio as gr
import os
 
from modules import images
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state
 
 
class Script(scripts.Script):  
    # The title of the script. This is what will be displayed in the dropdown menu.
    def title(self):
        return "Flip/Rotate Output"
 
    # Determines when the script should be shown in the dropdown menu via the 
    # returned value. As an example:
    # is_img2img is True if the current tab is img2img, and False if it is txt2img.
    # Thus, return is_img2img to only show the script on the img2img tab.
    # 在文生图、图生图页面是否出现
    def show(self, is_img2img):
        return is_img2img
 
    # How the script's is displayed in the UI. See https://gradio.app/docs/#components
    # for the different UI components you can use and how to create them.
    # Most UI components can return a value, such as a boolean for a checkbox.
    # The returned values are passed to the run method as parameters.
    # 具体展示的页面，这块也可以用on_ui_tabs()实现
    def ui(self, is_img2img):
        angle = gr.Slider(minimum=0.0, maximum=360.0, step=1, value=0,
        label="Angle")
        hflip = gr.Checkbox(False, label="Horizontal flip")
        vflip = gr.Checkbox(False, label="Vertical flip")
        overwrite = gr.Checkbox(False, label="Overwrite existing files")
        return [angle, hflip, vflip, overwrite] 
 
    # This is where the additional processing is implemented. The parameters include
    # self, the model object "p" (a StableDiffusionProcessing class, see
    # processing.py), and the parameters returned by the ui method.
    # Custom functions can be defined here, and additional libraries can be imported 
    # to be used in processing. The return value should be a Processed object, which is
    # what is returned by the process_images method.
    # 一般process函数，主函数
    def run(self, p, angle, hflip, vflip, overwrite):
        # function which takes an image from the Processed object, 
        # and the angle and two booleans indicating horizontal and
        # vertical flips from the UI, then returns the 
        # image rotated and flipped accordingly
        def rotate_and_flip(im, angle, hflip, vflip):
            from PIL import Image            
            raf = im            
            if angle != 0:
                raf = raf.rotate(angle, expand=True)
            if hflip:
                raf = raf.transpose(Image.FLIP_LEFT_RIGHT)
            if vflip:
                raf = raf.transpose(Image.FLIP_TOP_BOTTOM)
            return raf  
 
        # If overwrite is false, append the rotation information to the filename
        # using the "basename" parameter and save it in the same directory.
        # If overwrite is true, stop the model from saving its outputs and
        # save the rotated and flipped images instead.
        basename = ""
        if(not overwrite):
            if angle != 0:
                basename += "rotated_" + str(angle)
            if hflip:
                basename += "_hflip"
            if vflip:
                basename += "_vflip"
        else:
            p.do_not_save_samples = True
 
        proc = process_images(p)
 
        # rotate and flip each image in the processed images
        # use the save_images method from images.py to save
        # them.
        for i in range(len(proc.images)):
            proc.images[i] = rotate_and_flip(proc.images[i], angle, hflip, vflip)
            images.save_image(proc.images[i], p.outpath_samples, basename,
            proc.seed + i, proc.prompt, opts.samples_format, info= proc.info, p=p)
 
        return proc

3.extension中的JavaScript和css都将添加到页面。

4.extension中的localizations也会添加，名称相同会替换。

5.extension中的preload.py中的参数，在解析命令行参数之前加载该文件。一般含有preload函数。

一般放置到extensions目录下。

下面是一个例子，我把代码都列出来方便参考：

template.py

import modules.scripts as scripts
import gradio as gr
import os
 
from modules import images, script_callbacks
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state
 
class ExtensionTemplateScript(scripts.Script):
        # Extension title in menu UI
        def title(self):
                return "Extension Template"
 
        # Decide to show menu in txt2img or img2img
        # - in "txt2img" -> is_img2img is `False`
        # - in "img2img" -> is_img2img is `True`
        #
        # below code always show extension menu
        def show(self, is_img2img):
                return scripts.AlwaysVisible
 
        # Setup menu ui detail
        def ui(self, is_img2img):
                with gr.Accordion('Extension Template', open=False):
                        with gr.Row():
                                angle = gr.Slider(
                                        minimum=0.0,
                                        maximum=360.0,
                                        step=1,
                                        value=0,
                                        label="Angle"
                                )
                                checkbox = gr.Checkbox(
                                        False,
                                        label="Checkbox"
                                )
                # TODO: add more UI components (cf. https://gradio.app/docs/#components)
                return [angle, checkbox]
 
        # Extension main process
        # Type: (StableDiffusionProcessing, List) -> (Processed)
        # args is [StableDiffusionProcessing, UI1, UI2, ...]
        def run(self, p, angle, checkbox):
                # TODO: get UI info through UI object angle, checkbox
                proc = process_images(p)
                # TODO: add image edit process via Processed object proc
                return proc

template_on_settings.py

import modules.scripts as scripts
import gradio as gr
import os
 
from modules import shared
from modules import script_callbacks
 
def on_ui_settings():
    section = ('template', "Template")
    shared.opts.add_option(
        "option1",
        shared.OptionInfo(
            False,
            "option1 description",
            gr.Checkbox,
            {"interactive": True},
            section=section)
    )
 
script_callbacks.on_ui_settings(on_ui_settings)

template_on_tab.py

import modules.scripts as scripts
import gradio as gr
import os
 
from modules import script_callbacks
 
 
def on_ui_tabs():
    with gr.Blocks(analytics_enabled=False) as ui_component:
        with gr.Row():
            angle = gr.Slider(
                minimum=0.0,
                maximum=360.0,
                step=1,
                value=0,
                label="Angle"
            )
            checkbox = gr.Checkbox(
                False,
                label="Checkbox"
            )
            # TODO: add more UI components (cf. https://gradio.app/docs/#components)
        return [(ui_component, "Extension Template", "extension_template_tab")]
 
script_callbacks.on_ui_tabs(on_ui_tabs)