在 Devin 中试用 DeepWiki 处理私有仓库
DeepWiki

AUTOMATIC1111/stable-diffusion-webui

菜单

检查点系统

相关源文件

检查点系统负责管理、加载和切换不同的 Stable Diffusion 模型检查点（也称为权重文件）。该系统负责发现磁盘上的检查点文件、将它们加载到内存中，并进行优化以实现高效的资源利用。有关 VAE 管理的详细信息，请参阅 VAE 管理。

概述

Stable Diffusion Web UI 支持多种模型架构（SD1、SD2、SDXL、SSD、SD3），每种架构都需要不同的权重文件。检查点系统提供了一种统一的方式来处理这些模型文件，无论其内部结构如何。它会自动检测模型类型，应用适当的优化，并以最小的内存开销管理模型切换。

来源： modules/sd_models.py784-875

检查点表示

每个检查点都由 CheckpointInfo 类表示，该类存储有关检查点文件的元数据。

主要属性包括

  • filename：检查点文件的完整路径
  • name：用于显示的相对路径或基本名称
  • hash：用于标识的快速哈希（sha256 的前 8 个字节）
  • sha256：用于精确标识的完整 SHA256 哈希
  • title：UI 中显示的名称，包含哈希
  • metadata：附加信息（用于 safetensors 格式）
  • ids：可用于引用此检查点的标识符列表

来源： modules/sd_models.py56-128

检查点发现

系统在几个位置搜索检查点文件

  1. 主模型目录 (models/Stable-diffusion/)
  2. 由 --ckpt_dir 命令行选项指定的自定义目录
  3. 由 --ckpt 指定的特定检查点路径

检查点发现是在 list_models() 函数中实现的，该函数填充了

  • checkpoints_list：一个将标题映射到 CheckpointInfo 对象的字典
  • checkpoint_aliases：一个将各种 ID 映射到 CheckpointInfo 对象的字典

来源： modules/sd_models.py153-179

检查点加载过程

加载检查点涉及几个步骤

  1. 选择：找到要加载的适当检查点
  2. 状态字典加载：从磁盘读取权重
  3. 模型创建：根据配置实例化模型类
  4. 权重应用：将权重应用于模型
  5. 优化：应用内存优化和精度设置
  6. 设备传输：将模型移动到适当的设备

此顺序由 load_model() 和 reload_model_weights() 函数处理。

检查点选择

select_checkpoint() 函数确定要加载哪个检查点

  1. 尝试在设置中查找与名称匹配的检查点
  2. 如果未指定，则回退到第一个检查点
  3. 如果找不到检查点，则引发错误

来源： modules/sd_models.py218-240

状态字典加载

get_checkpoint_state_dict() 函数加载权重值

  1. 检查检查点是否已在内存缓存中
  2. 如果不在，则使用适当的格式（safetensors 或 PyTorch）从磁盘读取
  3. 如果启用了缓存，则存储在缓存中

来源： modules/sd_models.py332-347 modules/sd_models.py312-329

模型创建和权重应用

load_model_weights() 函数将权重应用于模型

  1. 确定模型类型（SD1、SD2、SDXL、SSD、SD3）
  2. 应用具有适当转换的权重
  3. 应用内存优化（半精度、channels_last 格式）
  4. 加载适当的 VAE

来源： modules/sd_models.py410-539

内存优化技术

检查点系统实现了多种内存优化策略

检查点缓存

系统可以将多个模型检查点保留在内存中

  • 由 sd_checkpoint_cache 设置控制
  • 使用 OrderedDict (checkpoints_loaded) 来跟踪已加载的检查点
  • 实现了 LRU（最近最少使用）的移除策略

精度选项

控制模型精度以平衡质量和内存使用

  • 全精度模式 (--no-half)
  • 默认半精度（FP16）
  • FP8 存储选项，用于最大限度地节省内存

VRAM 管理模式

适应各种 GPU VRAM 容量的不同模式

  • 标准模式：将模型完全保留在 GPU 上
  • 低 VRAM 模式：在不使用时卸载部分模型
  • Medvram 模式：性能和内存使用之间的折衷

VAE 集成

检查点系统与 VAE 系统集成

  1. 每个模型检查点都可以关联一个 VAE
  2. VAE 可以通过以下方式指定
    • 用户元数据
    • 与检查点附近具有匹配名称的文件
    • 全局设置
    • 自动检测

load_model_weights() 函数调用 sd_vae.load_vae() 来在加载模型后应用适当的 VAE。

来源： modules/sd_models.py535-539 modules/sd_vae.py132-234

用户元数据

检查点可以包含用户定义的元数据，包括

  • 描述
  • 备注
  • 首选 VAE

此元数据与检查点文件分开存储，并可以通过 UI 进行编辑。

来源： modules/ui_extra_networks_checkpoints_user_metadata.py1-66

模型类型检测

系统通过检查关键模式自动检测状态字典中的模型类型

  1. SD3：检查 "model.diffusion_model.x_embedder.proj.weight"
  2. SDXL：检查 "conditioner" 属性
  3. SSD：对于 SDXL 模型，检查是否缺少 transformer 块
  4. SD2：检查 cond_stage_model 是否具有 "model" 属性
  5. SD1：如果其他类型都不匹配，则为默认值

来源： modules/sd_models.py379-402

API 和集成

检查点系统为代码库的其他部分提供了几个关键功能

  • get_closet_checkpoint_match()：按名称查找检查点
  • model_hash()：为检查点标识生成快速哈希
  • reload_model_weights()：高效地切换到新的检查点
  • apply_token_merging()：应用 token 合并优化

该系统被以下部分使用

  • 用于模型选择的 Web UI
  • 用于远程模型操作的 API
  • 用于高级模型操作的扩展

来源： modules/sd_models.py183-200 modules/sd_models.py940-1008

检查点格式支持

系统支持两种主要的检查点格式

  1. SafeTensors (.safetensors)：一种安全格式，可防止代码执行
    • 提供更好的安全性
    • 可以包含元数据
    • 通常加载速度更快
  2. PyTorch (.ckpt, .pt)：传统的 PyTorch 检查点格式
    • 广泛用于旧模型
    • 与所有模型类型兼容

格式根据文件扩展名自动检测。

来源： modules/sd_models.py312-329