krita_ai_diffusion插件分析

109 阅读19分钟

ai_diffusion.desktop文件

ai_diffusion.desktop 这个文件名暗示它可能是一个桌面快捷方式(通常用于Linux系统)。我会打开文件,查看其内容,然后提供分析。让我们开始吧。

文件 ai_diffusion.desktop 是一个桌面快捷方式,通常用于Linux系统。它指定了以下属性:

  • Type: Service (服务类型)
  • ServiceTypes: Krita/PythonPlugin (服务类型是 Krita 的 Python 插件)
  • X-KDE-Library: ai_diffusion (用于 KDE 的库名称是 ai_diffusion)
  • X-Python-2-Compatible: false (不兼容 Python 2)
  • X-Krita-Manual: manual.html (Krita 手册文件名为 manual.html)
  • Name: AI Image Diffusion (名称是 AI 图像扩散)
  • Comment: Workflows supported by generative AI (Stable Diffusion) (备注说明该服务支持由生成式 AI(稳定扩散)提供的工作流)

这个快捷方式是用于启动与 Krita 图形编辑器相关的一个名为 AI Image Diffusion 的服务或插件。

__init__.py文件

__init__.py这个文件通常是一个 Python 包的一部分,用于将目录标记为包含 Python 模块。文件内容可能会包含包的初始化代码或仅仅是空的,用于标记。我将打开并读取这个文件的内容,以便进一步分析。让我们开始吧。

文件 __init__.py 是 Krita 的生成式 AI 插件(使用稳定扩散 Stable Diffusion)的一部分。这里是主要内容和其含义:

  1. 文档字符串:

    • 文件顶部有一个文档字符串 "Generative AI plugin for Krita using Stable Diffusion",描述了该文件的用途。
  2. 版本号:

    • __version__ = "1.12.0" 表示该插件的版本号是 1.12.0。
  3. 导入和条件检查:

    • 导入了 importlib.util 用于后续的模块检查。
    • 通过 importlib.util.find_spec 检查是否安装了名为 websockets.src 的模块。如果没有找到,会抛出一个 ImportError,说明 websockets 模块未与插件一起安装,并提供了一个链接到 GitHub 以下载正确的插件发布包。
  4. 依赖于 Krita 环境的导入:

    • 文件下半部分包含了一个条件导入语句,只有当代码在 Krita 内部运行时才会导入。这部分代码导入了 AIToolsExtension 类,但仅当找到 krita 模块时(即代码运行在 Krita 环境中)。

这个 __init__.py 文件负责设置插件,并确保必要的依赖项存在。

ai_diffusion.action 文件

ai_diffusion.action 根据文件扩展名,这个文件可能与自动化任务、脚本或特定应用程序的动作相关。我将打开这个文件,读取并分析其内容。让我们开始吧。

文件 ai_diffusion.action 是一个 XML 格式的文件,定义了一系列与 AI 图像生成相关的动作。这些动作似乎是为某个软件或应用程序(可能是图形编辑器如 Krita)设计的。下面是一些主要的动作和它们的描述:

  1. Generate image (生成图像) :

    • 动作名:ai_diffusion_generate
    • 描述:生成图像。
  2. Cancel current job (取消当前任务) :

    • 动作名:ai_diffusion_cancel
    • 描述:取消当前任务。
  3. Cancel queued jobs (取消排队的任务) :

    • 动作名:ai_diffusion_cancel_queued
    • 描述:取消排队的任务。
  4. Cancel all jobs (取消所有任务) :

    • 动作名:ai_diffusion_cancel_all
    • 描述:取消所有任务。
  5. Apply result (应用结果) :

    • 动作名:ai_diffusion_apply
    • 描述:应用结果。
  6. Switch workspace: Generate (切换工作区:生成) :

    • 动作名:ai_diffusion_switch_workspace_generation
    • 描述:切换到生成图像的工作区。
  7. Switch workspace: Upscale (切换工作区:放大) :

    • 动作名:ai_diffusion_switch_workspace_upscaling
    • 描述:切换到放大图像的工作区。
  8. Toggle workspace (切换工作区) :

    • 动作名:ai_diffusion_toggle_workspace
    • 描述:切换工作区。
  9. Show Settings (显示设置) :

    • 动作名:ai_diffusion_settings
    • 描述:显示设置。

这个文件定义了一系列的用户界面动作,每个动作都包含了标签、名称和其他可能的属性。这些动作似乎用于控制 AI 图像处理的不同阶段,包括生成、取消和应用结果等功能。

attention_edit.py 文件

attention_edit.py 包含了用于选择和编辑文本中特定部分的函数。这些函数可能用于一个更大的应用中,比如用于处理自然语言或者代码编辑器中的文本。这里是主要的功能和类的简要描述:

函数:

  1. select_current_parenthesis_block:

    • 描述:选择光标所指向的当前括号块。
    • 参数:text (str), cursor_pos (int), open_brackets (list[str]), close_brackets (list[str])
    • 返回:括号块的开始和结束索引的元组,或者在找不到括号块时返回 None。

函数 select_current_parenthesis_block,用于在文本中找到光标当前指向的括号块。这个函数特别有用于文本编辑器或编程IDE中,用于选择或高亮配对的括号。下面是逐行的详细分析:

pythonCopy code
def select_current_parenthesis_block(
    text: str, cursor_pos: int, open_brackets: list[str], close_brackets: list[str]
) -> Tuple[int, int] | None:
  • 这行定义了函数 select_current_parenthesis_block
  • 函数接收四个参数:text (文本), cursor_pos (光标位置), open_brackets (开括号列表), close_brackets (闭括号列表)。
  • 返回类型是 Tuple[int, int] | None,意味着函数返回一个元组(包含两个整数,表示括号块的开始和结束位置)或者 None(如果没有找到有效的括号块)。
pythonCopy code
    """Select the current parenthesis block that the cursor points to."""
  • 这是函数的文档字符串,说明了函数的目的和作用。
pythonCopy code
    # Ensure cursor position is within valid range
    cursor_pos = max(0, min(cursor_pos, len(text)))
  • 这行代码确保传入的 cursor_pos 在有效范围内(即在文本长度的范围内)。如果光标位置小于 0,它会被设置为 0。如果光标位置大于文本长度,它会被设置为文本的最大索引。
pythonCopy code
    # Find the nearest '(' before the cursor
    start = -1
    for open_bracket in open_brackets:
        start = max(start, text.rfind(open_bracket, 0, cursor_pos))
  • 初始化 start 变量为 -1,用于存储找到的最近的开括号位置。
  • 遍历 open_brackets 列表,使用 text.rfind 方法找到每种开括号在光标位置之前最近出现的位置。
  • 更新 start 为最近的开括号位置。
pythonCopy code
    # If '(' is found, find the corresponding ')' after the cursor
    end = -1
    if start != -1:
        open_parens = 1
        for i in range(start + 1, len(text)):
            if text[i] in open_brackets:
                open_parens += 1
            elif text[i] in close_brackets:
                open_parens -= 1
                if open_parens == 0:
                    end = i
                    break
  • 初始化 end 变量为 -1,用于存储找到的配对的闭括号位置。
  • 如果找到了开括号(即 start 不为 -1),则初始化 open_parens 计数为 1。
  • 遍历文本,从开括号之后的位置开始。
  • 如果遇到开括号,则增加 open_parens 计数;如果遇到闭括号,则减少 open_parens 计数。
  • 如果 open_parens 计数降到 0,说明找到了配对的闭括号,记录这个位置到 end,并跳出循环。
pythonCopy code
    # Return the indices only if both '(' and ')' are found
    if start != -1 and end >= cursor_pos:
        return start, end + 1
    else:
        return None
  • 检查是否找到了有效的括号块(即找到了开括号和对应的闭括号,并且闭括号的位置在光标位置之后)。
  • 如果找到了有效的括号块,则返回括号块的开始和结束位置(注意 end + 1 是为了包含闭括号)。
  • 如果没有找到有效的括号块,返回 None

整体而言,这个函数的逻辑是通过计数开闭括号的方法,找到包围光标位置的最近的括号块。这种方法在处理嵌套括号时非常有效。

  1. select_current_word:

    • 描述:选择光标所指向的当前单词。
    • 参数:text (str), cursor_pos (int)
    • 返回:单词的开始和结束索引的元组。

函数 select_current_word,用于在文本中找到光标当前指向的单词。这个函数在文本编辑器或编程IDE中非常有用,用于选择或高亮光标所在的单词。以下是逐行的详细分析:

pythonCopy code
def select_current_word(text: str, cursor_pos: int) -> Tuple[int, int]:
  • 这行定义了函数 select_current_word
  • 函数接收两个参数:text (文本) 和 cursor_pos (光标位置)。
  • 返回类型是 Tuple[int, int],意味着函数返回一个元组,包含两个整数,表示单词的开始和结束位置。
pythonCopy code
    """Select the word the cursor points to."""
  • 这是函数的文档字符串,说明了函数的目的和作用。
pythonCopy code
    delimiters = r".,/!?%^*;:{}=`~()<> " + "\t\r\n"
  • 这行代码定义了一个字符串 delimiters,包含了被视为单词分隔符的字符。这些分隔符包括常见的标点符号、空格、制表符 (\t)、回车 (\r) 和换行符 (\n)。
pythonCopy code
    start = end = cursor_pos
  • 初始化 startend 变量为 cursor_pos 的值,这两个变量将用来存储找到的单词的开始和结束位置。
pythonCopy code
    # seek backward to find beginning
    while start > 0 and text[start - 1] not in delimiters:
        start -= 1
  • 从光标位置开始,向文本的开头方向遍历,寻找单词的开始位置。
  • 如果当前检查的字符不是分隔符,start 减一,向前移动。
  • 当遇到分隔符或到达文本开头时停止。
pythonCopy code
    # seek forward to find end
    while end < len(text) and text[end] not in delimiters:
        end += 1
  • 从光标位置开始,向文本的末尾方向遍历,寻找单词的结束位置。
  • 如果当前检查的字符不是分隔符,end 加一,向后移动。
  • 当遇到分隔符或到达文本末尾时停止。
pythonCopy code
    return start, end
  • 返回找到的单词的开始和结束位置。这两个值定义了文本中的一个范围,包含了光标所在的单词。

整体而言,这个函数的逻辑是通过检查分隔符来找到光标所在位置的单词的边界。

  1. select_on_cursor_pos:

    • 描述:基于光标位置返回文本中的范围。它首先尝试选择括号块,如果没有找到括号块,则选择单词。
    • 参数:text (str), cursor_pos (int)
    • 返回:基于光标位置的文本范围的开始和结束索引的元组。

函数 select_on_cursor_pos,用于根据光标的位置在文本中选择一个范围。该函数首先尝试选择由括号包围的块,如果没有找到,它将选择光标所在的单词。以下是逐行的详细分析:

pythonCopy code
def select_on_cursor_pos(text: str, cursor_pos: int) -> Tuple[int, int]:
  • 这行定义了函数 select_on_cursor_pos
  • 函数接收两个参数:text (文本) 和 cursor_pos (光标位置)。
  • 返回类型是 Tuple[int, int],意味着函数返回一个元组,包含两个整数,表示选择的范围的开始和结束位置。
pythonCopy code
    """Return a range in the text based on the cursor_position."""
  • 这是函数的文档字符串,说明了函数的目的和作用。
pythonCopy code
    return select_current_parenthesis_block(
        text, cursor_pos, ["(", "<"], [")", ">"]
    ) or select_current_word(text, cursor_pos)
  • 这行代码首先调用 select_current_parenthesis_block 函数,尝试在文本中选择光标当前指向的括号块。括号类型包括圆括号 () 和尖括号 <>
  • 如果 select_current_parenthesis_block 返回 None(即没有找到括号块),则 or 表达式的右边部分会被求值。这时,函数调用 select_current_word,尝试选择光标所在的单词。
  • 函数返回的是两个函数中返回的第一个非 None 的结果。如果光标位于一个括号块内,它将返回该括号块的范围;如果光标不在括号块内,它将返回光标所在单词的范围。

整体而言,select_on_cursor_pos 函数提供了一种灵活的方式来根据光标位置选择文本,优先选择括号内的文本块,如果没有找到,则选择单词。这对于文本编辑和代码分析特别有用,因为它允许用户快速选择和操作代码结构的特定部分。

类:

  1. ExprNode:

    • 描述:表示一个表达式节点,可以是文本或表达式(带有权重)。
    • 属性:type (str), value (str or None), weight (float), children (list[ExprNode])

ExprNode 这个类用于表示一个表达式节点,可以是纯文本或者是一个带有权重和子节点的复杂表达式。以下是逐行的详细分析:

pythonCopy code
class ExprNode:
  • 定义了一个名为 ExprNode 的类。这个类用于表示一个表达式节点,这个节点可以是纯文本或者是一个复杂的表达式。
pythonCopy code
    def __init__(self, type, value, weight=1.0, children=None):
  • 这是类的初始化方法 __init__
  • 接收四个参数:type (节点类型), value (节点值), weight (节点权重,默认为 1.0), children (子节点列表,默认为 None)。
pythonCopy code
        self.type = type  # 'text' or 'expr'
  • 将传入的 type 参数赋值给实例变量 self.typetype 用于区分节点是纯文本 ("text") 还是表达式 ("expr")。
pythonCopy code
        self.value = value  # text or sub-expression
  • 将传入的 value 参数赋值给实例变量 self.value。对于文本节点 ("text"),value 是实际的文本内容;对于表达式节点 ("expr"),value 可以是子表达式。
pythonCopy code
        self.weight = weight  # weight for 'expr' nodes
  • 将传入的 weight 参数赋值给实例变量 self.weightweight 通常用于表达式节点,表示该节点在整个表达式中的相对重要性或权重。
pythonCopy code
        self.children = children if children is not None else []  # child nodes
  • 将传入的 children 参数赋值给实例变量 self.children。如果 children 不是 None,则直接使用传入的值;如果是 None,则将 self.children 初始化为空列表。children 存储的是当前节点的子节点列表。
pythonCopy code
    def __repr__(self):
  • 定义了 __repr__ 方法。这个方法用于提供对象的字符串表示,通常用于调试和日志。
pythonCopy code
        if self.type == "text":
            return f"Text('{self.value}')"
  • 如果节点类型是 "text",则返回格式化的字符串表示,其中包含节点的文本内容。
pythonCopy code
        else:
            assert self.type == "expr"
            return f"Expr({self.children}, weight={self.weight})"
  • 如果节点类型不是 "text",则断言节点类型必须是 "expr"
  • 返回格式化的字符串表示,其中包含节点的子节点和权重信息。

整体而言,ExprNode 类用于创建和表示一个表达式树的节点。这些节点可以是纯文本或者是复杂的表达式,带有权重和子节点。__repr__ 方法提供了一种方便的方式来查看节点的内容,特别是在调试时。

主要逻辑:

  1. parse_expr:

    • 描述:解析表达式,构建一个由 ExprNode 组成的列表。它处理带有权重的括号块,例如 (text:0.9)
    • 参数:expression (str)
    • 返回:ExprNode 列表。

函数 parse_expr,用于解析一个特定的表达式语言,生成一个由 ExprNode 组成的列表。这种表达式语言允许在括号内带有权重的子表达式,以及纯文本。以下是逐行的详细分析:

pythonCopy code
def parse_expr(expression: str) -> List[ExprNode]:
  • 定义了函数 parse_expr
  • 函数接收一个参数:expression (表达式文本)。
  • 返回类型是 List[ExprNode],意味着函数返回一个包含 ExprNode 对象的列表。
pythonCopy code
    """
    Parses following attention syntax language.
    expr = text | (expr:number)
    expr = text + expr | expr + text
    """
  • 这是函数的文档字符串,说明了函数的目的和解析的表达式语法。表达式可以是纯文本或者是带有权重的子表达式。子表达式由括号包围,并带有一个数字表示权重。
pythonCopy code
    def parse_segment(segment):
  • 定义了一个内部函数 parse_segment,用于解析单个段落。
pythonCopy code
        match = re.match(r"^[([{<](.*?):(-?[\d.]+)[]})>]$", segment)
  • 使用正则表达式匹配段落格式。这个正则表达式寻找一个开括号,后跟任意内容,然后是一个冒号和一个数字(可能是负数或小数),最后是一个闭括号。
pythonCopy code
        if match:
  • 如果匹配成功,说明段落是一个表达式。
pythonCopy code
            inner_expr = match.group(1)
            number = float(match.group(2))
  • 提取括号内的表达式 (inner_expr) 和权重 (number)。
pythonCopy code
            return ExprNode("expr", None, weight=number, children=parse_expr(inner_expr))
  • 递归调用 parse_expr 来解析括号内的表达式,然后创建一个类型为 "expr"ExprNode 对象。
pythonCopy code
        else:
  • 如果匹配失败,说明段落是纯文本。
pythonCopy code
            return ExprNode("text", segment)
  • 创建一个类型为 "text"ExprNode 对象。
pythonCopy code
    segments = []
    stack = []
    start = 0
    bracket_pairs = {"(": ")", "<": ">"}
  • 初始化 segments 列表用于存储解析的段落。
  • 初始化 stack 用于跟踪括号匹配。
  • 初始化 start 变量用于记录当前段落的开始位置。
  • 定义 bracket_pairs 字典用于存储配对的括号。
pythonCopy code
    for i, char in enumerate(expression):
  • 遍历表达式的每个字符。
pythonCopy code
        if char in bracket_pairs:
  • 如果当前字符是开括号。
pythonCopy code
            if not stack:
                if start != i:
                    segments.append(ExprNode("text", expression[start:i]))
                start = i
  • 如果 stack 为空,说明是一个新段落的开始。将之前的文本作为纯文本段落添加到 segments 列表。
pythonCopy code
            stack.append(bracket_pairs[char])
  • 将闭括号添加到 stack 以跟踪括号匹配。
pythonCopy code
        elif stack and char == stack[-1]:
  • 如果当前字符是期望的闭括号。
pythonCopy code
            stack.pop()
            if not stack:
                node = parse_segment(expression[start : i + 1])
                if node.type == "expr":
                    segments.append(node)
                    start = i + 1
                else:
                    stack.append(char)
  • stack 中移除闭括号。如果 stack 为空,说明找到了一个完整的段落。解析这个段落并将结果添加到 segments 列表。
pythonCopy code
    if start < len(expression):
        remaining_text = expression[start:].strip()
        if remaining_text:
            segments.append(ExprNode("text", remaining_text))
  • 检查并处理表达式末尾的剩余文本。
pythonCopy code
    return segments
  • 返回解析得到的段落列表。

整体而言,parse_expr 函数解析一个特定的表达式语言,生成一个由 ExprNode 组成的列表,这些节点可以是纯文本或带有权重和子节点的复杂表达式。这种结构非常适合处理嵌套的数据结构,如数学表达式、编程语言的语法树等。

  1. edit_attention:

    • 描述:编辑文本中的注意力。它调整文本或括号块内的权重,并返回编辑后的文本。
    • 参数:text (str), positive (bool) - 表示是否增加权重。
    • 返回:编辑后的文本。

函数 edit_attention,用于编辑文本提示内的注意力权重。它解析文本以识别权重,并根据 positive 参数增加或减少权重。以下是逐行的详细分析:

pythonCopy code
def edit_attention(text: str, positive: bool) -> str:
  • 定义了函数 edit_attention
  • 函数接收两个参数:text (要编辑的文本) 和 positive (布尔值,指示是增加 (True) 还是减少 (False) 注意力权重)。
  • 返回类型是 str,意味着函数返回修改后的文本。
pythonCopy code
    """Edit the attention of text within the prompt."""
  • 这是函数的文档字符串,说明了函数的目的和作用。
pythonCopy code
    if text == "":
        return text
  • 如果传入的文本为空,则直接返回文本。这是一个快速退出条件,避免对空字符串进行不必要的处理。
pythonCopy code
    segments = parse_expr(text)
  • 调用之前定义的 parse_expr 函数解析文本,并将结果存储在 segments 变量中。
pythonCopy code
    if len(segments) == 1 and segments[0].type == "expr":
        attention_string = text[1 : text.rfind(":")]
        weight = segments[0].weight
        open_bracket = text[0]
        close_bracket = text[-1]
  • 如果解析后只有一个段落,并且这个段落的类型是 "expr",则从文本中提取注意力字符串、权重和括号。
  • attention_string 是括号内的文本。
  • weight 是段落的权重。
  • open_bracketclose_bracket 分别是文本的开括号和闭括号。
pythonCopy code
    elif text[0] == "<":
        attention_string = text[1:-1]
        weight = 1.0
        open_bracket = "<"
        close_bracket = ">"
  • 如果文本以 < 开头,假设文本是一个特殊的表达式,没有指定权重。
  • attention_string 是括号内的文本。
  • 将权重设置为 1.0。
  • open_bracketclose_bracket 分别设置为 <>
pythonCopy code
    else:
        attention_string = text
        weight = 1.0
        open_bracket = "("
        close_bracket = ")"
  • 如果文本不符合上述两种情况,将整个文本视为注意力字符串,权重默认为 1.0,括号默认为 ()
pythonCopy code
    weight = weight + 0.1 * (1 if positive else -1)
  • 根据 positive 参数增加或减少权重。如果 positiveTrue,增加 0.1;如果为 False,减少 0.1。
pythonCopy code
    weight = max(weight, -2.0)
    weight = min(weight, 2.0)
  • 将权重限制在 -2.0 到 2.0 的范围内。
pythonCopy code
    return (
        attention_string
        if weight == 1.0 and open_bracket == "("
        else f"{open_bracket}{attention_string}:{weight:.1f}{close_bracket}"
    )
  • 如果权重为 1.0 且开括号为 (,则返回原始的注意力字符串,不包括括号和权重。
  • 否则,返回格式化的字符串,其中包含开括号、注意力字符串、权重和闭括号。

整体而言,edit_attention 函数提供了一种修改文本中注意力权重的方法。它解析文本以识别权重,并根据 positive 参数增加或减少权重,最后返回修改后的文本。这对于处理需要根据用户输入动态调整的文本提示特别有用。

client.py 文件

client.py 用于实现客户端逻辑,可能与网络通信或API请求有关。我将打开这个文件,读取并分析其内容。让我们开始吧。

文件 client.py 包含一个名为 Client 的类,这个类似乎用于与一个服务端进行 HTTP 或 WebSocket 通信。它涉及到多种资源的管理和与远端服务器的交互。以下是文件中定义的主要组件和它们的功能:

类和数据结构:

  1. DeviceInfo:

    • 用于存储和处理设备信息,例如类型、名称和 VRAM。
  2. CheckpointInfo:

    • 用于管理和表示模型检查点的信息,包括文件名、版本、是否用于填充(inpainting)和是否是细化器(refiner)。
  3. Client:

    • 主要的客户端类,用于与 ComfyUI 服务器进行通信。
    • 包含各种资源的管理,例如 checkpoints(检查点)、vae_models(VAE模型)、upscalers(放大器)、control_model(控制模型)等。
    • 提供与服务器交互的方法,例如连接到服务器、获取系统信息、检查缺失的包等。

主要功能和逻辑:

  1. async connect:

    • 静态异步方法用于连接到指定的 URL,并获取系统信息和节点信息。
  2. _ensure_supported_style:

    • 私有方法用于确保支持的样式存在,如果没有,则创建默认样式。
  3. _extract_message_png_image:

    • 私有方法用于从消息中提取 PNG 图像。
  4. _extract_pose_json:

    • 私有方法用于从消息中提取姿势信息(可能是 JSON 格式)。
  5. _validate_executed_node:

    • 私有方法用于验证执行的节点,检查接收的图像数量和类型。

整体上看,client.py 脚本是一个复杂的客户端实现,用于处理与远程服务器的交互,并管理与生成式模型相关的各种资源。