ZImageControlNetPipeline 工作流详解

总体流程

数据流向简图

点击左侧流程图节点即可跳转到对应函数详解。

flowchart TD
    A["prompt<br/><i>[str]</i>"] -->|"encode_prompt()"| B["prompt_embeds<br/><i>[list[Tensor(seq_i,D)]</i>"]

    C["control_image<br/><i>[PIL/Tensor(3,H,W)]</i>"] -->|"prepare_image()"| D["resized_image<br/><i>[Tensor(B,3,H,W)]</i>"]
    D -->|"VAE encode"| E["control_image_latent<br/><i>[Tensor(B,C,1,Hl,Wl)]</i>"]

    F["scheduler"] -->|"calculate_shift()"| G["timesteps<br/><i>[Tensor(N)]</i>"]

    H["random noise"] -->|"prepare_latents()"| I["latents<br/><i>[Tensor(B,C,Hl,Wl)] fp32</i>"]

    I -->|"unsqueeze(2)"| J["latent_input<br/><i>[Tensor(B,C,1,Hl,Wl)]</i>"]
    G -->|"expand(B)"| K["timestep<br/><i>[Tensor(B)]</i>"]

    B --> L["controlnet()"]
    E --> L
    J --> L
    K --> L

    L -->|"block_samples<br/><i>[list[Tensor]]</i>"| M["transformer()"]
    B --> M
    J --> M
    K --> M

    M -->|"model_out_list<br/><i>[list[Tensor(B,C,1,Hl,Wl)]]</i>"| N["CFG postprocess"]
    N -->|"noise_pred<br/><i>[Tensor(B,C,Hl,Wl)]</i>"| O["scheduler.step()"]
    O -->|"× N 循环"| I

    O -->|"final latents<br/><i>[Tensor(B,C,Hl,Wl)]</i>"| P["VAE decode"]
    P -->|"pixel<br/><i>[Tensor(B,3,H,W)]</i>"| Q["postprocess()"]
    Q --> R["最终图像<br/><i>[PIL.Image]</i>"]

ControlNet 做了什么

ControlNet 在这个 Pipeline 中的核心角色是将外部条件图像（如姿态、深度图、Canny 边缘、线稿等）注入到去噪过程中，从而精确控制生成图像的结构和布局。其工作可分为以下几个关键步骤：

1. 条件图像编码（预处理阶段）

在去噪循环开始之前，ControlNet 先将条件图像转为 latent 空间：

1. Resize：将输入 control_image（PIL 图像或 Tensor）缩放到与目标图像相同的尺寸 (B, 3, H, W)；
2. VAE Encode：通过 VAE 编码器将其映射到 latent 空间；
3. Argmax 采样：使用 argmax 模式取 latent 分布的众数，得到确定性压缩表示 (B, C, H_l, W_l)；
4. Shift + Scale：应用 VAE 的 shift_factor 和 scaling_factor 进行归一化；
5. Unsqueeze：在第 2 维插入一个维度，使形状变为 (B, C, 1, H_l, W_l)，与 latent 输入的 3D 格式对齐。

2. 去噪循环中的控制信号生成（第 6c 步）

在每一轮去噪步骤中，ControlNet 接收四个输入：

ControlNet 内部结构复制了 Transformer 的若干编码层（通常为前半部分 blocks），在这些层的输出上额外引入条件图像信息。具体来说：

• ControlNet 的每个 block 接受对应 Transformer block 的 latent 特征，同时也通过**零卷积（zero-convolution）**将条件图像特征融合进来；
• 零卷积初始化为零，确保训练初期 ControlNet 不会干扰原始模型的生成能力，随后逐步学会注入有用控制信号；
• 每个 block 输出的 controlnet_block_samples[i] 是一个残差信号（residual），代表"为了让生成结果符合条件图像，该 block 的特征应该往哪个方向偏移"。

3. 与 Transformer 的融合（第 6d 步）

ControlNet 输出的 controlnet_block_samples（一组残差 list）被传入 Transformer：

1
2
3
4
5
6

model_out_list = self.transformer(
    latent_model_input_list,
    timestep_model_input,
    prompt_embeds_model_input,
    controlnet_block_samples=controlnet_block_samples,  # ← ControlNet 的控制信号
)[0]

在 Transformer 内部，每个对应 block 的特征会和 controlnet_block_samples[i] 相加：

其中 $h_i$ 是原本的隐层特征，$c_i$ 是 ControlNet block 输出的残差，$\alpha$ 即 controlnet_conditioning_scale（默认 0.75）。

4. conditioning_scale 的作用

controlnet_conditioning_scale 控制条件图像的施加力度：

• 接近 0：几乎忽略条件图像，等价于纯文生图；
• 0.5~0.75（推荐）：在文字语义和条件结构之间取得平衡；
• 接近 1.0 或更高：强约束生成结果的结构，但可能削弱文字语义的响应。

总结

条件图像 → VAE Encode → ControlNet → 各层残差 → Transformer 融合 → 结构受控的生成结果

函数详解

encode_prompt() — 文本编码

对应流程图节点： prompt → prompt_embeds

位于 ZImageControlNetPipeline 中的入口方法，负责将自然语言 prompt 转为模型可用的 embedding 列表。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def encode_prompt(
    self,
    prompt: str | list[str],
    device: torch.device | None = None,
    do_classifier_free_guidance: bool = True,
    negative_prompt: str | list[str] | None = None,
    prompt_embeds: list[torch.FloatTensor] | None = None,
    negative_prompt_embeds: torch.FloatTensor | None = None,
    max_sequence_length: int = 512,
):

工作流程：

1. 对每个 prompt 调用 _encode_prompt() 获取正样本 embedding；
2. 若启用 CFG (do_classifier_free_guidance=True)：
   • 若未提供 negative_prompt，默认使用空字符串 "" 作为负样本；
   • 对负样本同样调用 _encode_prompt()；
3. 返回 (prompt_embeds, negative_prompt_embeds) 两个列表，每个元素为 [seq_i, D] 的 Tensor。

_encode_prompt() — 底层文本编码

1
2
3
4
5
6
7

def _encode_prompt(
    self,
    prompt: str | list[str],
    device: torch.device | None = None,
    prompt_embeds: list[torch.FloatTensor] | None = None,
    max_sequence_length: int = 512,
) -> list[torch.FloatTensor]:

工作流程：

1. Chat Template：将 prompt 包装为 [{"role": "user", "content": prompt}]，调用 tokenizer.apply_chat_template() 生成带特殊 token 的完整文本；
2. Tokenize：tokenizer(padding="max_length", max_length=512, truncation=True) → input_ids 和 attention_mask；
3. Text Encoder 前向：self.text_encoder(input_ids, attention_mask, output_hidden_states=True)；
4. 取倒数第二层：.hidden_states[-2]，这是 Z-Image 使用的关键设计——不使用最后一层，而是取倒数第二层 hidden states；
5. Mask 裁剪：只保留 attention_mask=True 位置的 embedding，去掉 padding token，得到变长序列 [seq_i, D]。

为什么取倒数第二层？ 倒数第二层 hidden states 相比最后一层保留了更丰富的中间语义信息，实验表明这对生成质量更有利。

prepare_image() — 条件图像预处理

对应流程图节点： control_image → resized_image → control_image_latent

将用户输入的条件图像（PIL 或 Tensor）统一预处理为 pipeline 可用的格式。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def prepare_image(
    self,
    image,
    width, height,
    batch_size,
    num_images_per_prompt,
    device, dtype,
    do_classifier_free_guidance=False,
    guess_mode=False,
):

工作流程：

1. 类型判断：若为 Tensor 则直接使用，否则调用 self.image_processor.preprocess() 进行 resize + normalize；
2. Batch 扩展：根据 batch_size 和 num_images_per_prompt 对图像进行 repeat_interleave；
3. CFG 复制：若启用 CFG，将图像复制一份（正负样本各一份），形状变为 [2B, 3, H, W]；
4. 返回预处理后的 Tensor。

后续在 __call__ 中继续处理：

1
2
3

control_image = retrieve_latents(self.vae.encode(control_image), sample_mode="argmax")
control_image = (control_image - shift_factor) * scaling_factor
control_image = control_image.unsqueeze(2)  # → [B, C, 1, Hl, Wl]

• retrieve_latents(..., sample_mode="argmax")：取 VAE 潜变量分布的众数（而非随机采样），确保条件信号确定性；
• shift + scale：应用 VAE 的归一化参数，使 latent 分布与 Transformer 训练时的输入分布一致；
• unsqueeze(2)：在第 2 维插入一个维度，从 [B, C, Hl, Wl] 变为 [B, C, 1, Hl, Wl]，与 3D 序列格式对齐。

prepare_latents() — 随机潜变量初始化

对应流程图节点： random noise → latents

生成去噪过程的初始噪声 latent。

1
2
3
4
5
6
7
8
9

def prepare_latents(
    self,
    batch_size,
    num_channels_latents,
    height, width,
    dtype, device,
    generator,
    latents=None,
):

工作流程：

1. 计算 latent 尺寸：height // vae_scale, width // vae_scale（vae_scale = vae_scale_factor × 2，通常为 16）；
2. 若 latents 为 None，调用 randn_tensor() 生成标准高斯噪声 [B, C, Hl, Wl]，dtype 固定为 fp32 以保证数值精度；
3. 若提供了已有 latent，则进行 shape 校验后直接使用（用于 img2img 或继续去噪）。

时间步调度

对应流程图节点： scheduler → timesteps

涉及三个辅助函数，在去噪循环前一次性计算好所有时间步。

calculate_shift() — 时间偏移

1
2
3
4
5
6
7

def calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
):

根据图像分辨率动态调整时间步调度。公式为：

• 高分辨率图像 → 更大的 image_seq_len → 更大的 mu → 更多去噪步骤集中在高噪声阶段，有利于生成全局结构；
• 低分辨率图像 → 更小的 mu → 去噪更均匀分布。

get_default_z_image_sigmas() — 默认 sigma 序列

1
2

def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
    return torch.linspace(1.0, 1 / num_inference_steps, num_inference_steps).tolist()

生成从 1.0 到 1/N 的线性递减 sigma 序列，用于 Flow Matching 调度器。

retrieve_timesteps() — 统一调度器接口

1
2
3
4
5
6
7
8

def retrieve_timesteps(
    scheduler,
    num_inference_steps=None,
    device=None,
    timesteps=None,
    sigmas=None,
    **kwargs,
):

支持三种调用方式：

• 传入 num_inference_steps → 自动生成默认 timesteps；
• 传入 sigmas → 使用自定义 sigma 序列；
• 传入 timesteps → 直接指定时间步。

最终调用 scheduler.set_timesteps() 并将 mu 作为额外参数传入。

controlnet() — ControlNet 前向

对应流程图节点： controlnet()

在每轮去噪中调用，生成控制信号残差。

1
2
3
4
5
6
7

controlnet_block_samples = self.controlnet(
    latent_model_input_list,   # list[Tensor(B, C, 1, Hl, Wl)] — 当前加噪 latent
    timestep_model_input,      # Tensor(B) — 当前时间步
    prompt_embeds_model_input, # list[Tensor(seq_i, D)] — 文本嵌入
    control_image_input,       # Tensor(B, C, 1, Hl, Wl) — 条件图像 latent
    conditioning_scale=0.75,   # float — 控制信号强度
)

内部结构：

• ControlNet 是 Transformer 的"副本"，只复制了前半部分（通常 19 个 block 中的前几个）；
• 每个 block 在原 Transformer 特征基础上，通过零卷积层将 control_image 特征融合进来；
• 零卷积（zero-convolution）初始化为 0，训练后学会输出非零残差；

输出： controlnet_block_samples — 一个 list，每个元素对应一个 Transformer block 的残差信号，形状与对应 block 的 hidden states 一致。

详细原理见上方 ControlNet 做了什么。

transformer() — Transformer 去噪

对应流程图节点： transformer()

融合文本语义、ControlNet 控制信号，预测当前时间步的噪声。

1
2
3
4
5
6

model_out_list = self.transformer(
    latent_model_input_list,                          # 当前加噪 latent
    timestep_model_input,                             # 时间步
    prompt_embeds_model_input,                        # 文本嵌入
    controlnet_block_samples=controlnet_block_samples, # ControlNet 残差
)[0]

内部流程：

1. 将 latent 和 prompt_embeds 分别转为 token 序列，拼接后送入 DiT（Diffusion Transformer）blocks；
2. 在每个 block 中，若传入了 controlnet_block_samples，则将其加到对应 block 的 hidden states 上： $$h_i' = h_i + \alpha \cdot c_i$$
3. 经过所有 blocks 后，输出预测的噪声/速度场 model_out_list。

输出： model_out_list — 一个 list，每个元素为 [B, C, 1, Hl, Wl]，代表预测的噪声。

CFG 后处理

对应流程图节点： CFG postprocess

Classifier-Free Guidance 的后处理，将正负样本的预测结果合并为最终噪声预测。

1
2

# 核心公式
pred = pos + current_guidance_scale * (pos - neg)

工作流程：

1. 拆分正负样本：model_out_list 的前半部分为 pos_out（正 prompt 预测），后半部分为 neg_out（负 prompt 预测）；
2. CFG 外推：对每个 batch 元素计算 pos + scale × (pos - neg)，guidance_scale 越大，越远离负样本方向；
3. 可选 Renormalization（cfg_normalization=True）：
   • 计算 pos 的 L2 范数 ori_pos_norm；
   • 计算 pred 的 L2 范数 new_pos_norm；
   • 若 new_pos_norm > ori_pos_norm × cfg_normalization，则按比例缩放 pred，防止 CFG 导致数值爆炸；
4. Squeeze：noise_pred.squeeze(2) 去除第 2 维，从 [B, C, 1, Hl, Wl] 变回 [B, C, Hl, Wl]；
5. 取反：noise_pred = -noise_pred，因为 Z-Image 使用 Flow Matching 范式，预测的是速度场，需要取反得到噪声方向。

CFG Truncation（可选）：

1
2

if t_norm > self._cfg_truncation:
    current_guidance_scale = 0.0  # 关闭 CFG

• 在去噪后期（t_norm > cfg_truncation），关闭 CFG 直接使用正样本预测，避免过度引导破坏细节；
• cfg_truncation=1.0 表示不截断，始终开启 CFG。

scheduler.step() — Scheduler 步进

对应流程图节点： scheduler.step()

使用 Flow Match Euler Discrete Scheduler 从当前 latent 减去预测噪声，得到下一步 latent。

1
2
3
4
5
6

latents = self.scheduler.step(
    noise_pred.to(torch.float32),  # 噪声预测（fp32）
    t,                              # 当前时间步
    latents,                        # 当前 latent（fp32）
    return_dict=False,
)[0]

数学原理（Flow Matching）：

在 Flow Matching 框架下，模型预测的是速度场 $v_\theta(x_t, t)$，Scheduler 通过 Euler 步进更新：

其中 $\Delta t$ 由 scheduler 根据当前 sigma 和下一步 sigma 自动计算。

注意： 所有数值均在 fp32 下计算，以保证去噪过程的数值稳定性。

VAE 解码

对应流程图节点： VAE decode

将去噪完成后的 latent 解码为像素空间图像。

1
2
3

latents = latents.to(self.vae.dtype)
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
image = self.vae.decode(latents, return_dict=False)[0]

工作流程：

1. 类型转换：将 fp32 latent 转为 VAE 的 dtype（通常 bf16/fp16）；
2. Reverse Scaling/Shifting：先除以 scaling_factor 再加回 shift_factor，这是预处理阶段 shift+scale 的逆操作；
3. VAE Decode：将 latent [B, C, Hl, Wl] 解码为像素 [B, 3, H, W]（其中 H, W 是 latent 尺寸 × vae_scale_factor）。

后处理

对应流程图节点： postprocess() → 最终图像

1

image = self.image_processor.postprocess(image, output_type=output_type)

工作流程：

1. 将 Tensor 像素值从 [-1, 1] 归一化到 [0, 255]；
2. 根据 output_type 转换格式：
   • "pil" → PIL.Image.Image 列表；
   • "pt" / "np" → 保持 Tensor / numpy 数组；
3. 调用 self.maybe_free_model_hooks() 释放模型内存（offload）。

辅助函数：retrieve_latents()

1
2
3
4
5

def retrieve_latents(
    encoder_output: torch.Tensor,
    generator: torch.Generator | None = None,
    sample_mode: str = "sample",
):

统一从 VAE 编码器输出中提取 latent：

• sample_mode="sample"：从 latent 分布中随机采样（用于初始噪声生成）；
• sample_mode="argmax"：取分布的众数（用于条件图像编码，保证确定性）；
• 也兼容直接返回 .latents 的编码器。