This is an extension predicated upon SD-WebUI, Designed to save and restore Webui workspace. It can quickly save your current webui interface and restore it the next time you use it. 这是一个基于 sd-webui 的扩展，旨在保存和恢复Webui工作区.它可以快速保存当前您的webui界面,并在下次使用的时候恢复.
from workspace.hooks import *

scripts/external_code.py
Scholar01/sd-webui-workspace
import gradio as gr
from fastapi import FastAPI
from modules import script_callbacks
from .workspace import StableDiffusionWorkSpace

sd_workflow = StableDiffusionWorkSpace()


def on_app_started(blocks: gr.Blocks, app: FastAPI):
    sd_workflow.setup_components(blocks)
    sd_workflow.setup_ui()


script_callbacks.on_app_started(on_app_started)

workspace/hooks.py
Scholar01/sd-webui-workspace
from collections import defaultdict


def patch(key, obj, field, replacement):
    """Replaces a function in a module or a class.

    Also stores the original function in this module, possible to be retrieved via original(key, obj, field).
    If the function is already replaced by this caller (key), an exception is raised -- use undo() before that.

    Arguments:
        key: identifying information for who is doing the replacement. You can use __name__.
        obj: the module or the class
        field: name of the function as a string
        replacement: the new function

    Returns:
        the original function
    """

    patch_key = (obj, field)
    if patch_key in originals[key]:
        raise RuntimeError(f"patch for {field} is already applied")

    original_func = getattr(obj, field)
    originals[key][patch_key] = original_func

    setattr(obj, field, replacement)

    return original_func


def undo(key, obj, field):
    """Undoes the peplacement by the patch().

    If the function is not replaced, raises an exception.

    Arguments:
        key: identifying information for who is doing the replacement. You can use __name__.
        obj: the module or the class
        field: name of the function as a string

    Returns:
        Always None
    """

    patch_key = (obj, field)

    if patch_key not in originals[key]:
        raise RuntimeError(f"there is no patch for {field} to undo")

    original_func = originals[key].pop(patch_key)
    setattr(obj, field, original_func)

    return None


def original(key, obj, field):
    """Returns the original function for the patch created by the patch() function"""
    patch_key = (obj, field)

    return originals[key].get(patch_key, None)


originals = defaultdict(dict)

workspace/patches.py
Scholar01/sd-webui-workspace
from itertools import chain

import PIL.Image
import base64
import gradio as gr
import numpy as np
from modules import ui_components
from pandas import DataFrame

exclude_tabs = [
    'tab_settings',
    'tab_extensions',
]

exclude_components = [
    gr.Tab,
    gr.Row,
    gr.Column,
    gr.Group,
    gr.Blocks,
    gr.Accordion,
    gr.Tabs,
    gr.TabItem,
    gr.Box,
    gr.Button,
    ui_components.FormRow,
    ui_components.FormColumn,
    ui_components.FormGroup,
    ui_components.ToolButton,
]


def get_tabs(blocks):
    tabs = []
    for key, block in blocks.blocks.items():
        if hasattr(block, "elem_id") \
                and block.elem_id \
                and block.elem_id.startswith("tab_") \
                and block.elem_id not in exclude_tabs:
            tabs.append(block)
    return tabs


def get_all_elem_id(component):
    elem_ids = []
    if getattr(component, 'elem_id', None) is not None:
        # 判断是否在exclude_components中
        if not any([isinstance(component, c) for c in exclude_components]):
            elem_ids.append(component.elem_id)
    if hasattr(component, 'children'):
        elem_ids.extend(chain(*[get_all_elem_id(c) for c in component.children]))
    return elem_ids


def parse_image(component, value):
    if not value:
        return value
    if component.tool == 'editor':
        image = value
    else:
        image = value.get('image', None)
    mode = 'RGB'
    if component.type == 'numpy':
        image = PIL.Image.fromarray(image)
        width = image.width
        height = image.height
        mode = image.mode
        data = base64.b64encode(image.tobytes()).decode('utf-8')
        # width = image.shape[1]
        # height = image.shape[0]
        # data = base64.b64encode(image.tobytes()).decode('utf-8')

    elif component.type == 'pil':
        width = image.width
        height = image.height
        mode = image.mode
        data = base64.b64encode(image.tobytes()).decode('utf-8')

    else:
        width = 0
        height = 0
        data = value

    return {
        'width': width,
        'height': height,
        'type': component.type,
        'mode': mode,
        'data': data,
    }


def to_image(value):
    if not value:
        return value
    type = value.get('type', None)
    image = value.get('data', None)
    width = value.get('width', 0)
    height = value.get('height', 0)
    mode = value.get('mode', 'RGB')
    if type == 'numpy' or type == 'pil':
        image = base64.b64decode(image)
        image = PIL.Image.frombytes(mode, (width, height), image)
        if type == 'numpy':
            image = np.asarray(image)

    return image


def component_to_value(component, value):
    if isinstance(value, DataFrame):
        return value.to_dict()

    if isinstance(component, gr.Image):
        value = parse_image(component, value)

    if isinstance(component, gr.Dropdown):
        if component.type == 'index':
            value = component.choices[value]

    if isinstance(component, gr.Radio):
        if component.type == 'index':
            value = component.choices[value][-1]

    return value


def flow_to_component(flow):
    value = flow['value']
    elem_type = flow['elem_type']

    if elem_type == 'dataframe':
        value = DataFrame.from_dict(value)

    if elem_type == 'image' and value:
        value = to_image(value)

    if elem_type == 'gallery' and value:
        if isinstance(value, list):
            value = [v['name'] for v in value]

    return value

workspace/util.py
Scholar01/sd-webui-workspace
import json
import os.path
from dataclasses import dataclass, field
from modules.ui import paste_symbol, save_style_symbol
from modules.ui_components import ToolButton
from modules.scripts import basedir
from .util import *


@dataclass
class StableDiffusionWorkSpace:
    blocks: gr.Blocks = field(default_factory=gr.Blocks, repr=False)
    tabs: list[gr.Tab] = field(default_factory=list, repr=False)
    components: list[gr.components.IOComponent] = field(default_factory=list, repr=False)
    filename: str = field(default=os.path.join(basedir(), 'workspace.json'))

    def setup_components(self, blocks):
        """
        当ui组件加载完成之后, 建立需要的事件关联
        """
        self.blocks = blocks
        self.tabs = get_tabs(blocks)
        self.components.clear()

        ids = []
        for tab in self.tabs:
            ids.extend(get_all_elem_id(tab))

        components = self.blocks.blocks.values()
        self.components = [c for c in components if c.elem_id in ids]
        # self.components_dict = {c.elem_id: c for c in components if c.elem_id in ids}
        # self.components = list(self.components_dict.values())

    def setup_ui(self):
        """
        初始化ui
        """
        if self.blocks is None or self.blocks.blocks == {}:
            return

        quick_settings: gr.Row = list(self.blocks.blocks.values())[0]
        assert quick_settings.elem_id == 'quicksettings'

        quick_settings.parent.__enter__()

        btn_import = ToolButton(value=paste_symbol, elem_id=self.elem_id('import'), tooltip='import workspace')
        btn_save = ToolButton(value=save_style_symbol, elem_id=self.elem_id('save'), tooltip='save workspace')

        quick_settings.add_child(btn_import)
        quick_settings.add_child(btn_save)
        btn_save.click(fn=self.on_save, inputs=self.components, outputs=None)
        btn_import.click(fn=self.on_load, inputs=None, outputs=self.components)
        quick_settings.parent.__exit__()

    def elem_id(self, elem_id):
        return f'workflow_{elem_id}'

    def on_save(self, *args):
        result = {}
        for i, arg in enumerate(args):
            component = self.components[i]
            value = component_to_value(component, arg)
            result[component.elem_id] = {
                'elem_type': str(component),
                'value': value
            }

        js = json.dumps(result, indent=4)
        with open(self.filename, 'w', encoding='utf8') as f:
            f.write(js)
        print('save workspace success')

    def on_load(self):
        result = []
        if not os.path.exists(self.filename):
            result = [gr.update() for g in range(len(self.components))]
        else:
            workflows = json.load(open(self.filename, encoding='utf8'))
            for component in self.components:
                if component.elem_id in workflows:
                    flow = workflows[component.elem_id]
                    value = flow_to_component(flow)
                    result.append(gr.update(value=value))
                else:
                    result.append(gr.update())
        return result

workspace/workspace.py
Scholar01/sd-webui-workspace