kang
6ed16be0a6
- LucidFlux (ICLR'26) 深度解析 - FLUX-IR (TPAMI'25) 深度解析 - 带文件+行号证据 - 两个深挖:Modulation 实现 + Reinforce ODE 采样 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
660 lines
32 KiB
HTML
660 lines
32 KiB
HTML
<!DOCTYPE html>
|
||
<html lang="zh-CN">
|
||
<head>
|
||
<meta charset="UTF-8">
|
||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||
<title>FLUX 图像修复双雄源码解析 · LucidFlux & FLUX-IR</title>
|
||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/atom-one-dark.min.css">
|
||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"></script>
|
||
<script>document.addEventListener('DOMContentLoaded',()=>{hljs.highlightAll();});</script>
|
||
<style>
|
||
:root{
|
||
--bg:#0b0d12;
|
||
--bg-2:#11141b;
|
||
--bg-3:#161a23;
|
||
--border:#242938;
|
||
--fg:#d8dee9;
|
||
--fg-dim:#8b93a7;
|
||
--fg-muted:#5c6478;
|
||
--accent:#7aa2f7;
|
||
--accent-2:#bb9af7;
|
||
--accent-3:#9ece6a;
|
||
--warn:#e0af68;
|
||
--danger:#f7768e;
|
||
--code-bg:#0f1218;
|
||
}
|
||
*{box-sizing:border-box;margin:0;padding:0}
|
||
html{scroll-behavior:smooth}
|
||
body{
|
||
font-family:-apple-system,BlinkMacSystemFont,"Segoe UI","PingFang SC","Hiragino Sans GB",sans-serif;
|
||
background:var(--bg);color:var(--fg);
|
||
line-height:1.75;font-size:15px;
|
||
-webkit-font-smoothing:antialiased;
|
||
}
|
||
.layout{display:grid;grid-template-columns:260px 1fr;min-height:100vh}
|
||
/* Sidebar */
|
||
aside{
|
||
position:sticky;top:0;height:100vh;overflow-y:auto;
|
||
background:var(--bg-2);border-right:1px solid var(--border);
|
||
padding:28px 20px;font-size:13px;
|
||
}
|
||
aside .brand{
|
||
font-weight:700;font-size:15px;color:var(--accent);
|
||
padding-bottom:14px;margin-bottom:14px;border-bottom:1px solid var(--border);
|
||
letter-spacing:.3px;
|
||
}
|
||
aside .brand small{display:block;color:var(--fg-muted);font-weight:400;margin-top:4px;font-size:11px;letter-spacing:0}
|
||
aside nav ul{list-style:none}
|
||
aside nav>ul>li{margin-bottom:2px}
|
||
aside nav a{
|
||
display:block;color:var(--fg-dim);text-decoration:none;
|
||
padding:6px 10px;border-radius:6px;transition:all .15s;
|
||
border-left:2px solid transparent;
|
||
}
|
||
aside nav a:hover{background:var(--bg-3);color:var(--fg)}
|
||
aside nav .sub{padding-left:14px;margin:2px 0 6px}
|
||
aside nav .sub a{font-size:12px;padding:4px 10px;color:var(--fg-muted)}
|
||
aside nav .sub a:hover{color:var(--fg-dim)}
|
||
aside .meta{margin-top:24px;padding-top:16px;border-top:1px solid var(--border);color:var(--fg-muted);font-size:11px;line-height:1.8}
|
||
aside .meta code{background:var(--bg-3);padding:1px 5px;border-radius:3px;font-size:10px}
|
||
|
||
main{padding:56px 64px 120px;max-width:960px}
|
||
h1{font-size:32px;font-weight:700;margin-bottom:8px;letter-spacing:-.5px;color:#fff}
|
||
h1 .emoji{margin-right:8px}
|
||
.subtitle{color:var(--fg-dim);font-size:15px;margin-bottom:8px}
|
||
.tags{display:flex;gap:8px;flex-wrap:wrap;margin-bottom:32px;margin-top:16px}
|
||
.tag{
|
||
font-size:11px;padding:3px 10px;border-radius:999px;
|
||
background:var(--bg-3);border:1px solid var(--border);color:var(--fg-dim);
|
||
font-family:"SF Mono",Monaco,Menlo,Consolas,monospace;
|
||
}
|
||
.tag.accent{border-color:var(--accent);color:var(--accent)}
|
||
.tag.purple{border-color:var(--accent-2);color:var(--accent-2)}
|
||
.tag.green{border-color:var(--accent-3);color:var(--accent-3)}
|
||
|
||
h2{
|
||
font-size:24px;font-weight:700;margin:56px 0 16px;
|
||
padding-bottom:10px;border-bottom:1px solid var(--border);
|
||
color:#fff;scroll-margin-top:20px;
|
||
}
|
||
h2 .num{color:var(--accent);margin-right:12px;font-family:"SF Mono",monospace;font-size:20px}
|
||
h3{
|
||
font-size:18px;font-weight:600;margin:32px 0 12px;
|
||
color:var(--accent-2);scroll-margin-top:20px;
|
||
}
|
||
h4{font-size:15px;font-weight:600;margin:20px 0 8px;color:var(--fg)}
|
||
p{margin:10px 0;color:var(--fg)}
|
||
strong{color:#fff;font-weight:600}
|
||
em{color:var(--warn);font-style:normal}
|
||
a{color:var(--accent);text-decoration:none;border-bottom:1px dotted var(--accent)}
|
||
a:hover{border-bottom-style:solid}
|
||
|
||
/* Inline code & code blocks */
|
||
code{
|
||
font-family:"JetBrains Mono","SF Mono",Monaco,Menlo,Consolas,monospace;
|
||
font-size:13px;
|
||
}
|
||
p code,li code,td code{
|
||
background:var(--bg-3);padding:2px 6px;border-radius:4px;
|
||
color:var(--accent-3);font-size:12.5px;
|
||
border:1px solid var(--border);
|
||
}
|
||
pre{
|
||
background:var(--code-bg);border:1px solid var(--border);
|
||
border-radius:8px;padding:16px 20px;margin:16px 0;
|
||
overflow-x:auto;font-size:13px;line-height:1.65;
|
||
position:relative;
|
||
}
|
||
pre code{background:transparent !important;padding:0;border:0;color:inherit}
|
||
|
||
/* File ref badges */
|
||
.fileref{
|
||
display:inline-block;background:var(--bg-3);color:var(--accent);
|
||
padding:1px 8px;border-radius:4px;font-size:11.5px;
|
||
font-family:"JetBrains Mono",monospace;
|
||
border:1px solid var(--border);
|
||
border-bottom:1px solid var(--border) !important;
|
||
}
|
||
|
||
/* Lists */
|
||
ul,ol{margin:12px 0;padding-left:26px}
|
||
li{margin:6px 0;color:var(--fg)}
|
||
li::marker{color:var(--fg-muted)}
|
||
|
||
/* Tables */
|
||
table{
|
||
width:100%;border-collapse:collapse;margin:20px 0;
|
||
background:var(--bg-2);border-radius:8px;overflow:hidden;
|
||
font-size:13.5px;
|
||
}
|
||
th,td{
|
||
padding:10px 14px;text-align:left;
|
||
border-bottom:1px solid var(--border);
|
||
}
|
||
th{background:var(--bg-3);font-weight:600;color:var(--fg);font-size:12.5px;text-transform:uppercase;letter-spacing:.5px}
|
||
tr:last-child td{border-bottom:0}
|
||
td code{font-size:11.5px}
|
||
|
||
/* Callouts */
|
||
.callout{
|
||
padding:14px 20px;border-radius:8px;margin:20px 0;
|
||
border-left:3px solid var(--accent);
|
||
background:color-mix(in srgb, var(--accent) 8%, var(--bg-2));
|
||
}
|
||
.callout.warn{border-left-color:var(--warn);background:color-mix(in srgb, var(--warn) 8%, var(--bg-2))}
|
||
.callout.danger{border-left-color:var(--danger);background:color-mix(in srgb, var(--danger) 8%, var(--bg-2))}
|
||
.callout.success{border-left-color:var(--accent-3);background:color-mix(in srgb, var(--accent-3) 8%, var(--bg-2))}
|
||
.callout strong:first-child{color:var(--accent)}
|
||
.callout.warn strong:first-child{color:var(--warn)}
|
||
.callout.danger strong:first-child{color:var(--danger)}
|
||
.callout.success strong:first-child{color:var(--accent-3)}
|
||
|
||
/* Hero card */
|
||
.hero{
|
||
background:linear-gradient(135deg, color-mix(in srgb, var(--accent) 12%, var(--bg-2)), var(--bg-2));
|
||
border:1px solid var(--border);
|
||
border-radius:12px;padding:24px 28px;margin:24px 0 40px;
|
||
}
|
||
.hero p{margin:6px 0;color:var(--fg-dim);font-size:14px}
|
||
|
||
/* Two-column compare */
|
||
.compare-grid{display:grid;grid-template-columns:1fr 1fr;gap:20px;margin:20px 0}
|
||
.compare-grid .card{
|
||
background:var(--bg-2);border:1px solid var(--border);
|
||
border-radius:10px;padding:20px;
|
||
}
|
||
.compare-grid .card h4{margin-top:0;color:var(--accent)}
|
||
.compare-grid .card.right h4{color:var(--accent-2)}
|
||
.compare-grid ul{padding-left:18px;font-size:13px}
|
||
.compare-grid li{color:var(--fg-dim);margin:4px 0}
|
||
|
||
/* Mobile */
|
||
@media (max-width: 900px){
|
||
.layout{grid-template-columns:1fr}
|
||
aside{position:static;height:auto;border-right:0;border-bottom:1px solid var(--border)}
|
||
main{padding:32px 20px 80px}
|
||
.compare-grid{grid-template-columns:1fr}
|
||
h1{font-size:26px}
|
||
h2{font-size:20px}
|
||
}
|
||
|
||
/* Scrollbar */
|
||
::-webkit-scrollbar{width:10px;height:10px}
|
||
::-webkit-scrollbar-track{background:var(--bg)}
|
||
::-webkit-scrollbar-thumb{background:var(--border);border-radius:5px}
|
||
::-webkit-scrollbar-thumb:hover{background:var(--fg-muted)}
|
||
</style>
|
||
</head>
|
||
<body>
|
||
<div class="layout">
|
||
|
||
<aside>
|
||
<div class="brand">
|
||
FLUX 图像修复双雄
|
||
<small>LucidFlux × FLUX-IR · 源码深度解析</small>
|
||
</div>
|
||
<nav>
|
||
<ul>
|
||
<li><a href="#overview">概述 & 标题党验证</a></li>
|
||
<li><a href="#lucidflux">一、LucidFlux 解析</a>
|
||
<ul class="sub">
|
||
<li><a href="#lf-arch">1. 整体架构</a></li>
|
||
<li><a href="#lf-train">2. 训练入口</a></li>
|
||
<li><a href="#lf-infer">3. 推理入口</a></li>
|
||
<li><a href="#lf-diff">4. 与原 FLUX 差异</a></li>
|
||
<li><a href="#lf-repro">5. 可复现性</a></li>
|
||
</ul>
|
||
</li>
|
||
<li><a href="#fluxir">二、FLUX-IR 解析</a>
|
||
<ul class="sub">
|
||
<li><a href="#fi-method">1. 核心方法论</a></li>
|
||
<li><a href="#fi-arch">2. Unified 模型结构</a></li>
|
||
<li><a href="#fi-task">3. 任务条件化</a></li>
|
||
<li><a href="#fi-train">4. 训练与推理</a></li>
|
||
<li><a href="#fi-vs">5. Task-specific vs Unified</a></li>
|
||
<li><a href="#fi-repro">6. 可复现性</a></li>
|
||
</ul>
|
||
</li>
|
||
<li><a href="#compare">三、两篇对比</a></li>
|
||
<li><a href="#deep1">四、深挖 1 · Modulation</a></li>
|
||
<li><a href="#deep2">五、深挖 2 · Reinforce ODE</a></li>
|
||
<li><a href="#essence">六、创新本质</a></li>
|
||
</ul>
|
||
</nav>
|
||
<div class="meta">
|
||
<div>日期 · 2026-04-13</div>
|
||
<div>本地路径 · <code>~/Projects/research/<br>20260413-flux-image-restoration/</code></div>
|
||
<div style="margin-top:10px">证据原则 · 所有论断带<br>文件路径 + 行号,不糊弄</div>
|
||
</div>
|
||
</aside>
|
||
|
||
<main>
|
||
<h1><span class="emoji">🎨</span>FLUX 图像修复双雄源码解析</h1>
|
||
<p class="subtitle">把 12B 参数的预训练扩散模型,改造成通用图像修复器的两种正交路线</p>
|
||
<div class="tags">
|
||
<span class="tag accent">LucidFlux · ICLR'26</span>
|
||
<span class="tag purple">FLUX-IR · TPAMI'25</span>
|
||
<span class="tag green">HKUST</span>
|
||
<span class="tag">Black Forest FLUX.1-dev</span>
|
||
</div>
|
||
|
||
<div class="hero">
|
||
<p><strong>起因</strong>:刷短视频刷到"FLUX 和 Wan 预训练扩散模型秒变图像修复神器",想验真伪。</p>
|
||
<p><strong>结论</strong>:FLUX 这条路线实打实,Wan 是视频生成模型,跟图像修复无关——标题党。</p>
|
||
<p><strong>本文</strong>:深啃两个 FLUX 修复代表作的源码,带文件路径 + 行号证据。</p>
|
||
</div>
|
||
|
||
<h2 id="overview"><span class="num">§0</span>概述 & 标题党验证</h2>
|
||
<p>两个仓库都 <code>git clone</code> 下来全量扫过,<code>grep -r "[Ww]an2\.?[12]"</code> 零命中。Wan 跟这两个工作毫无关系。真正能把 FLUX 拽成修复器的,是下面这两条完全不同的技术路线:</p>
|
||
<ul>
|
||
<li><strong>LucidFlux</strong>(HKUST-GZ,W2GenAI-Lab)——架构创新路线:冻结 FLUX,加双分支 conditioner + SigLIP caption-free 语义对齐</li>
|
||
<li><strong>FLUX-IR</strong>(HKUST,Zhu Zhiyu)——训练算法创新路线:ControlNet 改造 + 强化 ODE 轨迹对齐 + 轨迹蒸馏</li>
|
||
</ul>
|
||
|
||
<h2 id="lucidflux"><span class="num">§1</span>LucidFlux 深度解析</h2>
|
||
|
||
<h3 id="lf-arch">1. 整体架构</h3>
|
||
<p><strong>核心策略</strong>:<em>冻结 FLUX.1 主干</em>,只训练双分支 conditioner + SigLIP 语义对齐模块。这是最"温和"的改造——对原生 FLUX 能力破坏极小。</p>
|
||
|
||
<ul>
|
||
<li><span class="fileref">src/flux/lucidflux.py:312-343</span> 加载时对 FLUX 主干 <code>requires_grad_(False)</code>,只有 condition_branch 进入训练</li>
|
||
<li><span class="fileref">src/flux/model.py:138-215</span> FLUX 前向接收 <code>block_controlnet_hidden_states</code> 参数,166-194 行把 conditioner 输出直接加到每个 DoubleStreamBlock</li>
|
||
<li>原 FLUX 的 double/single stream block 保持原样(<span class="fileref">src/flux/modules/layers.py</span>),未魔改。hidden_size=3072, num_heads=24, depth=19 都保留</li>
|
||
</ul>
|
||
|
||
<h4>双分支 conditioner</h4>
|
||
<p><span class="fileref">src/flux/lucidflux.py:189-244</span></p>
|
||
<ul>
|
||
<li><code>condition_lq</code>: 处理低质量(LQ)原图</li>
|
||
<li><code>condition_pre</code>: 处理 SwinIR 预超分结果</li>
|
||
<li>两支各含 2 个 DoubleStreamBlock + <code>zero_module</code> 初始化的 ControlNet block(78-80 行)</li>
|
||
<li>通过 Modulation 模块自适应融合(232-243 行)</li>
|
||
</ul>
|
||
|
||
<h4>Modulation(本文深挖 1)</h4>
|
||
<p><span class="fileref">src/flux/lucidflux.py:161-186</span></p>
|
||
<ul>
|
||
<li>timestep + control_index 联合编码</li>
|
||
<li>LayerNorm 后输出 shift/scale 系数做特征调制</li>
|
||
<li>control_index 在 0-19 范围遍历,支持<strong>分层</strong>自适应</li>
|
||
</ul>
|
||
|
||
<h4>SigLIP 语义对齐</h4>
|
||
<p><span class="fileref">src/flux/flux_prior_redux_ir.py:44-104</span></p>
|
||
<ul>
|
||
<li>SigLIP 视觉编码器吃 SwinIR 前置输出 → Redux 编码器转 1024 token FLUX 兼容嵌入</li>
|
||
<li>训练时(<span class="fileref">train.py:476-487</span>)和推理时(<span class="fileref">inference.py:160-195</span>)与文本 embedding 拼接</li>
|
||
</ul>
|
||
<div class="callout success">
|
||
<strong>这就是 "caption-free" 的关键</strong> —— 不需要用户写 prompt,SigLIP 从图像里自动提取语义,再通过 Redux 编码器适配到 FLUX 的文本 token 空间。
|
||
</div>
|
||
|
||
<h3 id="lf-train">2. 训练入口</h3>
|
||
<ul>
|
||
<li><span class="fileref">train.py:300-548</span> 主流程,<code>train_configs/train_LucidFlux.yaml</code> 配置</li>
|
||
<li>超参:批大小 1, 最大 100k 步, lr=2e-5, <strong>Adafactor</strong> 优化器</li>
|
||
<li><span class="fileref">train.py:324-351</span> 模型初始化 — FLUX 主干 + VAE 冻结,加载预训 conditioner 权重,<em>只 conditioner 进 <code>.train()</code></em></li>
|
||
<li><strong>损失</strong>:L2 MSE 预测噪声残差,target = <code>noise - packed_latents</code> (<span class="fileref">train.py:95-100, 502</span>)</li>
|
||
<li><strong>数据</strong>:<span class="fileref">image_datasets/lq_gt_dataset.py:34-68</span> LQ/GT 配对,中心裁剪 512×512</li>
|
||
</ul>
|
||
|
||
<p><strong>训练 pipeline 步骤</strong>:</p>
|
||
<ol>
|
||
<li>VAE 编码 LQ/GT 到 latent</li>
|
||
<li>SwinIR 8× 前置超分</li>
|
||
<li>SigLIP + Redux 提取语义特征</li>
|
||
<li>DualConditionBranch 输出残差</li>
|
||
<li>FLUX 接收残差 → 预测噪声</li>
|
||
<li>MSE loss 回传,只更新 conditioner</li>
|
||
</ol>
|
||
|
||
<h3 id="lf-infer">3. 推理入口</h3>
|
||
|
||
<h4>标准推理 · <code>inference.py</code></h4>
|
||
<ul>
|
||
<li><span class="fileref">inference.py:93</span> <code>load_flow_model</code> 加载 FLUX</li>
|
||
<li><span class="fileref">inference.py:94</span> VAE 来自 FLUX 官方</li>
|
||
<li><span class="fileref">inference.py:112</span> 输入分辨率 16 倍对齐</li>
|
||
</ul>
|
||
|
||
<h4>2K 推理 · <code>inference-2k.py</code></h4>
|
||
<div class="callout">
|
||
<strong>关键差异</strong> · <span class="fileref">inference-2k.py:130</span> 改用 UltraFlux VAE,不再用 FLUX 官方 VAE:
|
||
</div>
|
||
<pre><code class="language-python">from src.ultraflux.autoencoder_kl import AutoencoderKL
|
||
ae = AutoencoderKL.from_pretrained("./weights/ultraflux", subfolder="vae")</code></pre>
|
||
<ul>
|
||
<li>分辨率对齐 32 倍(<span class="fileref">inference-2k.py:148</span>)</li>
|
||
<li>输入预处理 <code>width//2, height//2</code>(<span class="fileref">inference-2k.py:180</span>)</li>
|
||
</ul>
|
||
|
||
<h4>推理 Pipeline</h4>
|
||
<p><span class="fileref">inference.py:139-224</span></p>
|
||
<ol>
|
||
<li>LQ → 标准化到 [-1, 1]</li>
|
||
<li>SwinIR 超分(151-158 行)</li>
|
||
<li>SigLIP + Redux 特征(173-190 行)</li>
|
||
<li>生成噪声 + embedding(162-198 行)</li>
|
||
<li><code>denoise_lucidflux()</code> 逐步去噪(<span class="fileref">sampling.py:96-151</span>):每步 DualConditionBranch 输出残差加到 FLUX</li>
|
||
<li>VAE 解码 + 小波重构对齐</li>
|
||
</ol>
|
||
|
||
<h3 id="lf-diff">4. 与原 FLUX 的差异</h3>
|
||
<h4>新增文件</h4>
|
||
<ul>
|
||
<li><span class="fileref">src/flux/condition.py:33-223</span> SingleConditionBranch</li>
|
||
<li><span class="fileref">src/flux/lucidflux.py</span> DualConditionBranch + Modulation + ConditionBranchWithRedux</li>
|
||
<li><span class="fileref">src/flux/flux_prior_redux_ir.py</span> SigLIP + Redux 集成</li>
|
||
<li><span class="fileref">src/flux/swinir.py</span> SwinIR 前置</li>
|
||
</ul>
|
||
<h4>原 FLUX 魔改范围</h4>
|
||
<div class="callout success">
|
||
<strong>几乎零魔改</strong> —— 只在 <span class="fileref">model.py:146, 166-194</span> 加了 <code>block_controlnet_hidden_states</code> 参数支持。DoubleStreamBlock、SingleStreamBlock 内部结构完全不动。这是 LucidFlux 最大的优势:未来 FLUX 官方升级,迁移成本极低。
|
||
</div>
|
||
|
||
<h3 id="lf-repro">5. 可复现性</h3>
|
||
<table>
|
||
<tr><th>项目</th><th>值</th></tr>
|
||
<tr><td>主模型</td><td><code>black-forest-labs/FLUX.1-dev</code> (HF)</td></tr>
|
||
<tr><td>SigLIP</td><td><code>siglip2-so400m-patch16-512</code></td></tr>
|
||
<tr><td>Conditioner</td><td><code>weights/lucidflux/lucidflux.pth</code></td></tr>
|
||
<tr><td>2K VAE</td><td><code>weights/ultraflux/vae/</code>(额外)</td></tr>
|
||
<tr><td>训练数据</td><td>HF <code>W2GenAI/LucidFlux</code> → <code>LucidFlux-Training-Data.tar.gz</code></td></tr>
|
||
<tr><td>数据过滤</td><td><code>tools/filtering_pipeline.py</code></td></tr>
|
||
<tr><td>依赖</td><td><code>diffusers==0.32.2</code>, <code>transformers==4.43.3</code>, <code>liger_kernel==0.6.1</code>, <code>deepspeed==0.18.8</code></td></tr>
|
||
<tr><td>显存 · 标准</td><td>~28GB(offload)</td></tr>
|
||
<tr><td>显存 · 2K</td><td>~38GB</td></tr>
|
||
<tr><td>显存 · ComfyUI</td><td>8-12GB(量化)</td></tr>
|
||
</table>
|
||
|
||
<h2 id="fluxir"><span class="num">§2</span>FLUX-IR 深度解析</h2>
|
||
|
||
<h3 id="fi-method">1. 核心方法论 · ODE 轨迹学习</h3>
|
||
<p>论文宣称的两个核心创新:<strong>Reinforced ODE Alignment</strong> 和 <strong>Distillation Cost-Aware ODE Acceleration</strong>。这才是 FLUX-IR 的灵魂——不是架构创新,是<em>训练目标</em>创新。</p>
|
||
|
||
<h4>Reinforced ODE Alignment</h4>
|
||
<p><span class="fileref">Unified_restoration/src/flux/sampling.py:241-356</span> 的 <code>denoise_controlnet_rein()</code>:</p>
|
||
<ul>
|
||
<li><strong>双轨迹 reward</strong>:
|
||
<ul>
|
||
<li><code>X_ode_t</code> 纯 ODE 求解轨迹</li>
|
||
<li><code>X_sde_t</code> 在 <code>sde_step</code> 控制的步骤注入随机的 SDE 轨迹</li>
|
||
</ul>
|
||
</li>
|
||
<li>用 SDE 作为 reward 信号对齐 ODE 轨迹</li>
|
||
</ul>
|
||
|
||
<h4>Distillation Cost-Aware ODE Acceleration</h4>
|
||
<p><span class="fileref">sampling.py:690-789</span> 的 <code>denoise_controlnet_distill()</code>:</p>
|
||
<ul>
|
||
<li>把 10 步蒸馏到 5-8 步</li>
|
||
<li>损失设计:<code>loss_t2O</code>(蒸馏 vs 原轨迹)+ <code>loss_t2G</code>(vs GT)+ 基础像素损失</li>
|
||
<li>成本感知:<code>t_stride ∈ {10, 8, 6}</code> 加权不同阶段(<span class="fileref">model.py:99-193</span>)</li>
|
||
</ul>
|
||
|
||
<h4>两阶段训练</h4>
|
||
<p><span class="fileref">Task_specific_restoration/model/model.py:203-227</span></p>
|
||
<ul>
|
||
<li><code>optimize_parameters_reinforce()</code> → 第一阶段:训练 ODE/SDE 对齐</li>
|
||
<li><code>optimize_parameters_distill()</code> → 第二阶段:蒸馏加速</li>
|
||
</ul>
|
||
|
||
<h3 id="fi-arch">2. Unified 模型结构 · ControlNet 改造 FLUX-dev</h3>
|
||
|
||
<h4>ControlNet 集成</h4>
|
||
<p><span class="fileref">src/flux/controlnet.py:33-223</span></p>
|
||
<ul>
|
||
<li><code>ControlNetFlux</code> 继承 FLUX 架构,加 2 层 DoubleStreamBlock</li>
|
||
<li><span class="fileref">controlnet.py:27-30</span> <code>zero_module()</code> 零初始化保证初期无扰(ControlNet 标准做法)</li>
|
||
<li><span class="fileref">controlnet.py:174-179</span> 条件图像压缩到 latent + 位置编码融合</li>
|
||
<li><span class="fileref">model.py:199-200</span> 残差加法: <code>img = img + block_controlnet_hidden_states[index_block % 2]</code></li>
|
||
</ul>
|
||
|
||
<h4>双 VAE encoder(关键创新)</h4>
|
||
<p><span class="fileref">xflux_pipeline.py:43-44</span></p>
|
||
<pre><code class="language-python">self.ae2 = load_ae(model_type, device=...)
|
||
self.ae2.encoder.load_state_dict(torch.load('checkpoints/encoder_lq.bin'))</code></pre>
|
||
<ul>
|
||
<li><code>ae</code> 吃 GT</li>
|
||
<li><code>ae2</code> 吃 LQ,<strong>独立训练</strong>的 encoder</li>
|
||
<li>目的:学习品质差异。原生 FLUX VAE 对 LQ 图像编码效果差,专门训一个 LQ encoder 能把退化信息更好地映射到 latent 空间</li>
|
||
</ul>
|
||
|
||
<h4>RAM 标签(可选增强)</h4>
|
||
<p><span class="fileref">main.py:30-41</span> / <span class="fileref">main_.py:92-104</span></p>
|
||
<ul>
|
||
<li>Recognize Anything Model 自动打图像标签</li>
|
||
<li><code>prompts = f"{ram_prompt}, {args.prompt}"</code>(<code>main_.py</code> 版本)</li>
|
||
<li><code>main.py</code> 里被注释,作为可选增强</li>
|
||
</ul>
|
||
|
||
<h3 id="fi-task">3. 任务条件化</h3>
|
||
<p><strong>无显式任务 embedding</strong> —— 全靠 prompt 文本 + <code>control_weight</code> 隐式区分:</p>
|
||
<table>
|
||
<tr><th>任务</th><th>Prompt</th><th>control_weight</th></tr>
|
||
<tr><td>超分</td><td><code>high-resolution, ultra-sharp, detailed</code></td><td>0.8</td></tr>
|
||
<tr><td>去噪</td><td><code>noise-free, clean, smooth</code></td><td>0.8</td></tr>
|
||
<tr><td>低光</td><td><code>bright, clear, vivid</code></td><td>0.9</td></tr>
|
||
<tr><td>去雨</td><td><code>remove raindrops, clean</code></td><td>0.9</td></tr>
|
||
</table>
|
||
|
||
<h4>控制权重缩放</h4>
|
||
<p><span class="fileref">sampling.py:312</span></p>
|
||
<pre><code class="language-python">block_controlnet_hidden_states = [i * controlnet_gs for i in ...]</code></pre>
|
||
<p><code>controlnet_gs</code> 直接映射命令行 <code>--control_weight</code>。任务越"重"(去雨、低光),weight 越大。</p>
|
||
|
||
<h3 id="fi-train">4. 训练与推理架构</h3>
|
||
|
||
<h4>两个 main 脚本</h4>
|
||
<ul>
|
||
<li><code>main.py</code> 中心裁剪 1024×1024 单图推理(<span class="fileref">main.py:220-221</span>)</li>
|
||
<li><code>main_.py</code> 分块推理 + 加权融合,支持任意尺寸(<span class="fileref">main_.py:23-77</span>)</li>
|
||
</ul>
|
||
|
||
<h4>推理流程</h4>
|
||
<p><span class="fileref">xflux_pipeline.py:239-500</span></p>
|
||
<ol>
|
||
<li>图像 → VAE 编码 → 4×64×64 latent</li>
|
||
<li><code>denoise_controlnet()</code> ODE 求解(21 步)</li>
|
||
<li>latent → VAE 解码 → 输出图像</li>
|
||
<li>支持 <code>offload=True</code> 降显存</li>
|
||
</ol>
|
||
|
||
<h4>损失函数</h4>
|
||
<p><span class="fileref">Task_specific_restoration/model/model.py:99-193</span></p>
|
||
<ul>
|
||
<li><strong>蒸馏阶段</strong>:<code>MSE(distilled_traj, GT_traj) + MSE(output, target)</code></li>
|
||
<li><strong>强化阶段</strong>:<code>MSE(X_ode, X_sde) + MSE(final_output, GT)</code></li>
|
||
<li>像素损失占主导,轨迹损失辅助收敛</li>
|
||
</ul>
|
||
|
||
<h3 id="fi-vs">5. Task-specific vs Unified 对比</h3>
|
||
<div class="compare-grid">
|
||
<div class="card">
|
||
<h4>Task-specific</h4>
|
||
<ul>
|
||
<li>每任务独立模型(LOLv1/Raindrop/Underwater)</li>
|
||
<li>纯 DDPM U-Net,<strong>未改 FLUX</strong></li>
|
||
<li><code>core/</code> 只有指标计算代码</li>
|
||
<li>轻量但要训 N 个模型</li>
|
||
</ul>
|
||
</div>
|
||
<div class="card right">
|
||
<h4>Unified</h4>
|
||
<ul>
|
||
<li>单 FLUX-dev + ControlNet 适配</li>
|
||
<li>基于 X-FLUX 魔改(<span class="fileref">README.md:102</span> 致谢)</li>
|
||
<li>参数共享,一模型搞 10+ 任务</li>
|
||
<li>新增:<code>xflux_pipeline.py</code> 管道 · <code>controlnet.py</code> 新类 · <code>model.py:165-200</code> ControlNet 残差注入</li>
|
||
</ul>
|
||
</div>
|
||
</div>
|
||
|
||
<h3 id="fi-repro">6. 可复现性</h3>
|
||
<table>
|
||
<tr><th>项目</th><th>值</th></tr>
|
||
<tr><td>权重</td><td><code>checkpoints/FluxIR.bin</code> + <code>checkpoints/encoder_lq.bin</code></td></tr>
|
||
<tr><td>下载</td><td><a href="https://drive.google.com/drive/folders/1CFWxmxOwcp6ARRX-y9yYsXSwpIRAgK37">Google Drive</a></td></tr>
|
||
<tr><td>显存 @1024² (bf16)</td><td>FLUX-dev ~21GB + VAE 解码 ~8GB</td></tr>
|
||
<tr><td>建议 GPU</td><td>H100 / A100 / L40S(40GB+)</td></tr>
|
||
<tr><td>依赖</td><td>PyTorch 2.4+, CUDA 12+, diffusers / transformers / accelerate / deepspeed</td></tr>
|
||
<tr><td>推理步数</td><td>21(蒸馏后)</td></tr>
|
||
</table>
|
||
|
||
<h2 id="compare"><span class="num">§3</span>两篇正面对比</h2>
|
||
<table>
|
||
<tr><th>维度</th><th>LucidFlux</th><th>FLUX-IR</th></tr>
|
||
<tr><td>核心创新</td><td>双分支 conditioner + SigLIP caption-free</td><td>RL-ODE 对齐 + 轨迹蒸馏</td></tr>
|
||
<tr><td>训练目标</td><td>只训 conditioner</td><td>ControlNet + LQ VAE encoder</td></tr>
|
||
<tr><td>任务区分</td><td>不区分(通用)</td><td>Prompt + control_weight</td></tr>
|
||
<tr><td>预处理</td><td>SwinIR 前置超分</td><td>RAM 标签(可选)</td></tr>
|
||
<tr><td>基础代码</td><td>黑森林 FLUX.1 官方</td><td>X-FLUX 魔改</td></tr>
|
||
<tr><td>推理步数</td><td>标准</td><td>21 步(蒸馏后)</td></tr>
|
||
<tr><td>场次</td><td>ICLR 2026</td><td>TPAMI 2025</td></tr>
|
||
<tr><td>对原 FLUX 侵入</td><td>几乎零</td><td>中等(双 VAE + ControlNet)</td></tr>
|
||
</table>
|
||
<div class="callout">
|
||
<strong>本质差异</strong> · LucidFlux 侧重"<em>条件注入怎么设计最不侵入 FLUX</em>"——架构创新。FLUX-IR 侧重"<em>采样轨迹怎么优化最快最准</em>"——训练算法创新。两条路线基本正交,理论可组合:LucidFlux 的 conditioner + FLUX-IR 的蒸馏轨迹。
|
||
</div>
|
||
|
||
<h2 id="deep1"><span class="num">§4</span>深挖 1 · LucidFlux Modulation 怎么实现分层自适应</h2>
|
||
<p>这段代码是 LucidFlux 能用"只有 2 个 DoubleStreamBlock 输出的 conditioner"喂"FLUX 19 层主干"的魔法所在。</p>
|
||
|
||
<h4>完整实现</h4>
|
||
<p><span class="fileref">src/flux/lucidflux.py:161-186</span></p>
|
||
<pre><code class="language-python">class Modulation(nn.Module):
|
||
def __init__(self, dim, bias=True):
|
||
self.silu = nn.SiLU()
|
||
self.linear = nn.Linear(dim, 2 * dim, bias=bias)
|
||
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=dim)
|
||
self.control_index_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=dim)
|
||
# ↑ 关键:多加了一个 control_index 嵌入器
|
||
|
||
def forward(self, x, timestep, control_index):
|
||
timesteps_proj = self.time_proj(timestep * 1000)
|
||
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=x.dtype))
|
||
|
||
# control_index 也走正弦位置编码 + MLP
|
||
control_index_proj = self.time_proj(control_index)
|
||
control_index_emb = self.control_index_embedder(control_index_proj.to(dtype=x.dtype))
|
||
|
||
# 两个 embedding 相加,一起产生调制参数
|
||
timesteps_emb = timesteps_emb + control_index_emb
|
||
emb = self.linear(self.silu(timesteps_emb))
|
||
shift_msa, scale_msa = emb.chunk(2, dim=1)
|
||
return self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]</code></pre>
|
||
|
||
<h4>关键洞察</h4>
|
||
<div class="callout success">
|
||
和 DiT 里经典的 adaLN-zero <em>一模一样</em>,但<strong>多加了一个 <code>control_index</code> 嵌入</strong>——这个 index 就是要注入的目标 block 编号(0-18)。两个 embedding 相加后过 linear 产生 shift/scale,实现:<br>
|
||
「在第 <code>t</code> 步、针对第 <code>k</code> 个 block,调 conditioner 信号的尺度和偏移」
|
||
</div>
|
||
|
||
<h4>调用处:拉伸 2 → 19</h4>
|
||
<p><span class="fileref">src/flux/lucidflux.py:232-243</span></p>
|
||
<pre><code class="language-python">out = []
|
||
num_blocks = 19
|
||
for i in range(num_blocks // 2 + 1): # i=0..9
|
||
for control_index, (lq, pre) in enumerate(zip(out_lq, out_pre)):
|
||
control_index = torch.tensor(control_index, device=timesteps.device, dtype=timesteps.dtype)
|
||
lq = self.modulation_lq(lq, timesteps, i * 2 + control_index)
|
||
if len(out) == num_blocks:
|
||
break
|
||
pre = self.modulation_pre(pre, timesteps, i * 2 + control_index)
|
||
out.append(lq + pre) # 两支相加,按层位置注入 FLUX
|
||
return out</code></pre>
|
||
<div class="callout">
|
||
<strong>核心技巧</strong> · conditioner 只输出 2 个残差(因为只有 2 个 DoubleStreamBlock),但 FLUX 主干有 19 个 DoubleStreamBlock 要注入。<br><br>
|
||
<strong>解决方案</strong>:同一个 conditioner 输出 + 19 个不同的 <code>(timestep, control_index)</code> 组合 → <strong>19 份不同调制结果</strong>,相当于用 <code>control_index</code> 把 2 个原始信号"拉伸"成 19 个位置感知的信号。<br><br>
|
||
<strong>收益</strong>:参数量极省——如果直接让 conditioner 输出 19 个残差,训练参数至少翻 9 倍。
|
||
</div>
|
||
|
||
<h2 id="deep2"><span class="num">§5</span>深挖 2 · FLUX-IR Reinforce ODE 用 SDE 当 reward</h2>
|
||
|
||
<h4>数据流概览</h4>
|
||
<p><span class="fileref">Unified_restoration/src/flux/sampling.py:241-355</span></p>
|
||
<ul>
|
||
<li><code>x_0 = img</code> 纯噪声</li>
|
||
<li><code>x_1 = controlnet_cond[1:2]</code> 训练时是 <strong>GT 图像 latent</strong>,推理用 LQ</li>
|
||
<li><code>img = (1-t)*x_1 + t*x_0</code> <em>线性插值初始化</em> —— rectified flow 的标准套路</li>
|
||
<li><code>X_ode_t</code> / <code>X_sde_t</code> 两个列表记录两条轨迹</li>
|
||
</ul>
|
||
|
||
<h4>采样循环</h4>
|
||
<p><span class="fileref">sampling.py:291-353</span></p>
|
||
<pre><code class="language-python">for t_curr, t_prev in zip(timesteps[t_start:], timesteps[t_start+1:]):
|
||
block_res_samples = controlnet(...) # ControlNet 残差
|
||
pred = model(...,
|
||
block_controlnet_hidden_states=[i * controlnet_gs for i in block_res_samples])
|
||
|
||
# CFG 分支被硬关(if ... and False),节省显存
|
||
if i == sde_step and sample_sde:
|
||
img = sde_sampling(t_curr, t_curr - t_prev, img, pred, seed) # 注入 SDE 噪声
|
||
else:
|
||
img = img + (t_prev - t_curr) * pred # 普通 ODE 步
|
||
|
||
X_sde_t.append(img)
|
||
i += 1</code></pre>
|
||
|
||
<h4>SDE 注入细节</h4>
|
||
<p><span class="fileref">sampling.py:471-496</span> 的 <code>sde_sampling()</code>:</p>
|
||
<pre><code class="language-python">def sde_sampling(t_curr, deltaT, x_curr, pred_noise_residual, seed):
|
||
eplson = get_noise(1, 1024, 1024, device=x_curr.device, dtype=torch.float32, seed=seed)
|
||
eplson = rearrange(eplson, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||
|
||
alpha = torch.tensor([1 + np.random.random() * 1]) # α ∈ [1, 2]
|
||
|
||
# 方差补偿公式(保证 SDE 边际分布与 ODE 一致)
|
||
beta_ = ((t_curr - deltaT)**2 * (1 - (t_curr - alpha * deltaT))**2 /
|
||
(1 - (t_curr - deltaT))**2) - (t_curr - alpha * deltaT)**2
|
||
beta = torch.sqrt(beta_)
|
||
|
||
while beta_ < 0: # 数值稳定:重采样 alpha
|
||
alpha = torch.tensor([np.random.randint(2, 10)])
|
||
beta_ = ((t_curr - deltaT)**2 * (1 - (t_curr - alpha * deltaT))**2 /
|
||
(1 - (t_curr - deltaT))**2) - (t_curr - alpha * deltaT)**2
|
||
beta = torch.sqrt(beta_)
|
||
if beta_ > 0:
|
||
break
|
||
# ... 后续:alpha 控制漂移项,beta 控制噪声项</code></pre>
|
||
|
||
<h4>方法论精髓</h4>
|
||
<div class="callout success">
|
||
<strong>"Reinforce" 的真实含义不是 PPO 那种 RL</strong>,而是:
|
||
<ul>
|
||
<li>ODE 确定性采样:精度高但可能卡在局部最优</li>
|
||
<li>SDE 随机注入:能跳出局部最优探索更好轨迹</li>
|
||
<li>Loss = <code>MSE(X_sde_t, X_ode_t)</code> + 其他项</li>
|
||
<li><em>用 SDE 的随机探索当自监督 teacher,蒸馏回确定性 ODE</em></li>
|
||
</ul>
|
||
<strong>收益</strong>:推理时纯 ODE 就能跑出 SDE 级质量,21 步搞定。
|
||
</div>
|
||
|
||
<div class="callout warn">
|
||
<strong>冷知识</strong> · <span class="fileref">sampling.py:316</span> 那句 <code>if i >= timestep_to_start_cfg and False:</code> —— CFG 分支被<strong>硬关</strong>了,只剩条件路径。这是为了减显存(去雨/去噪不太需要负 prompt)。代码里留着原始 CFG 实现以防需要重启。
|
||
</div>
|
||
|
||
<h2 id="essence"><span class="num">§6</span>两个创新点的本质</h2>
|
||
<table>
|
||
<tr><th></th><th>LucidFlux Modulation</th><th>FLUX-IR Reinforce</th></tr>
|
||
<tr><td>解决的问题</td><td>2 个 conditioner 输出怎么喂 19 个 FLUX block</td><td>ODE 确定性采样怎么学到 SDE 的鲁棒性</td></tr>
|
||
<tr><td>代价</td><td>增加 2 个 Modulation 模块(几 MB 参数)</td><td>训练时 2 份前向,推理不变</td></tr>
|
||
<tr><td>可迁移性</td><td>可插入任何 DiT-style 模型</td><td>可用于任何 rectified flow 模型</td></tr>
|
||
<tr><td>能否组合</td><td colspan="2" style="text-align:center;color:var(--accent-3);font-weight:600">✓ 理论正交,可拼</td></tr>
|
||
</table>
|
||
|
||
<div class="callout">
|
||
<strong>最后一句话</strong> · 两篇都是"把 FLUX 当通用先验"的实践,但思路完全不同。LucidFlux 告诉你<em>最不动 FLUX 怎么做修复</em>,FLUX-IR 告诉你<em>把 FLUX 改狠一点、但采样更快怎么做修复</em>。下次看到类似标题党,记得先 <code>grep</code> 一下源码。
|
||
</div>
|
||
|
||
<div style="margin-top:80px;padding-top:24px;border-top:1px solid var(--border);color:var(--fg-muted);font-size:12px;text-align:center">
|
||
源码解析 · 2026-04-13 · 所有论断带文件:行号,欢迎自行 <code>grep</code> 验证
|
||
</div>
|
||
|
||
</main>
|
||
</div>
|
||
</body>
|
||
</html>
|