0. 简介

Masked Visual Actions 面向的是机器人操作中的世界模型问题,处理的矛盾很具体:预训练视频模型在互联网视频上积累了大量关于运动、接触、遮挡与持存的先验,但机器人领域惯用的动作信号是关节角、末端位姿或控制向量,这些量与视频模型学习先验时所处的像素空间不对齐,也随机械结构改变而失效。斯坦福、马里兰大学与哈佛的作者没有再给视频模型加一路低维动作编码器,而是把动作直接写成画面里某个实体的时空掩码轨迹:保留机器人所占的像素区域、其余区域涂灰,模型的任务就变成条件视频补全。这样一个接口带来了一个不对称的收益——训练时只喂机器人掩码,推理时换成物体掩码,同一份权重会零样本地反过来补出机械臂行为。全文用 15 小时机器人视频、rank 256 的 LoRA 适配 Wan2.2-Fun-A14B-Control,单一 checkpoint 同时支撑策略评估、基于模型的规划与动作提取三类下游任务。下面结合论文与官方仓库,重点拆解掩码接口为什么能跨具身、前向与逆向为何能共用一份权重、以及这条路线现在的真实边界在哪里。

这项工作的贡献不在视频生成质量本身,而在动作表示的接口设计。底座模型没改、训练器没改,改的是"动作长什么样"这一件事。把它放进机器人世界模型的谱系里看,它站在"低维动作条件"与"稀疏视觉条件"之间,选择了最稠密、最与图像对齐的那一端。

在这里插入图片描述

1. 机器人世界模型面临问题

1.1 动作信号与视频先验的错位

机器人世界模型的通用做法是给生成模型加一路动作条件:末端位姿七维、关节角七维、或者干脆是控制器指令。这条路线的问题不在精度,而在这些量根本不在视频模型学习先验的那个空间里。视频模型在海量真实视频上学到的是"手靠近杯子、杯子会倾倒"这类像素级因果,它从没见过"末端位姿 z 轴下降 3 厘米"这种编码。要让它理解,就得额外学一层从低维向量到画面变化的对应关系,而这层对应关系恰恰是最容易过拟合到训练时那台机器人的部分。

进一步看,把动作可视化到图像上并不自动解决问题。IRASim 把末端位姿画成三维坐标轴叠在画面上,VAP 把机器人骨架画成连杆线段,这两种都是稀疏视觉条件:它们告诉模型"末端在哪",但不告诉模型机械臂的外形有多粗、夹爪张开多少、哪块金属会挡住杯子。模型只能自己脑补几何与遮挡,而脑补的模板来自训练数据里那台机器人。换一台机器人,脑补就会把新机器人"改造"回旧的样子。

在这里插入图片描述

1.2 掩码轨迹给出的第三条路

掩码视觉动作的做法是把实体在整段视频里占据的时空像素区域整块保留下来,其余部分统一置为灰色。它不是单帧轮廓,而是随时间逐帧变化的时空掩码,所以机械臂的外形、位置、尺度、夹爪张合状态、以及它与物体的接触过程全部直接可见。模型要做的事情从"解码一个抽象向量"退化成"把已给出的一部分画面补完整",这正是视频模型在预训练阶段反复做过的动作。

直觉理解(掩码接口像什么):低维动作条件像给画师一串坐标数字让他画手臂,稀疏骨架像给他一副火柴人简笔,而掩码轨迹相当于直接把手臂那块剪影贴在画布上、让画师补背景和被碰到的东西。前两种画师必须"知道你的手臂长什么样",第三种不需要。

1.3 这个名词解决的是接口问题不是能力问题

这里要厘清的是,掩码视觉动作没有给视频模型增加任何新能力。论文在结论里写得很直白:这套方法是 repurpose 模型已有的先验,而不是修改它的能力上限。生成速度、能表达的物理现象范围、对细微接触的刻画精度,全部受限于底座模型。它真正改变的是同一份先验能被多少种机器人、多少种查询方式调用——这是接口层面的收益,不是模型层面的收益,把这两件事分开看,后面的实验数字才不会被误读。

2. 整体框架:一个联合分布的两种切法

2.1 把场景写成实体集合

论文的形式化非常简洁。一段视频 V ∈ R T × H × W × 3 V \in \mathbb{R}^{T \times H \times W \times 3} V∈RT×H×W×3 描绘一个场景 S \mathcal{S} S,而场景被看作实体集合 e 1 , e 2 , … , e n e_1, e_2, \ldots, e_n e1​,e2​,…,en​。每个实体在输出视频里都有一条时空轨迹,论文直接用 e i e_i ei​ 同时指代实体本身和它占据的像素区域。视频模型隐式建模的就是所有实体轨迹的联合分布,包括它们之间的相互作用:

p ( V ) = p ( e 1 , e 2 , … , e n ) p(V) = p(e_1, e_2, \ldots, e_n) p(V)=p(e1​,e2​,…,en​)

这个写法本身没有引入新假设,它只是把"视频模型建模了什么"这件事换了个说法:既然模型能生成整段视频,它就必然隐含了所有实体轨迹的联合分布,包括它们相互碰撞、遮挡、推动时的耦合。有了联合分布,剩下的事情就是从里面取条件。给定一个条件实体子集 S ⊆ { 1 , … , n } \mathcal{S} \subseteq \{1, \ldots, n\} S⊆{1,…,n},得到的条件分布是

p ( { e i } i ∉ S   ∣   { e j } j ∈ S , I 0 ) p\bigl(\{e_i\}_{i \notin \mathcal{S}} \,\big|\, \{e_j\}_{j \in \mathcal{S}}, I_0\bigr) p({ei​}i∈/S​ ​{ej​}j∈S​,I0​)

其中 I 0 I_0 I0​ 是初始场景的参考图,它提供了真实的场景外观,让模型知道桌面、光照、物体材质长什么样。这里的关键是:改变 S \mathcal{S} S 就等于向同一个模型问不同的问题,而不是换一个模型、换一个 head、或者换一套权重。整篇论文的所有应用差异,最后都能归约到"这一次 S \mathcal{S} S 里放了谁"这一个自由度上。

2.2 掩码是条件的实现方式

条件通过掩码落地,这一步是把上面那个抽象的"取条件"变成可计算的操作。令 M ∈ { 0 , 1 } T × H × W M \in \{0,1\}^{T \times H \times W} M∈{0,1}T×H×W 是二值掩码, M t , h , w = 1 M_{t,h,w}=1 Mt,h,w​=1 表示该时空像素被揭示给模型, M t , h , w = 0 M_{t,h,w}=0 Mt,h,w​=0 表示这个位置需要模型自己预测。注意掩码的三个维度覆盖了时间与空间,这是它区别于单帧掩码的地方。对选定的条件集 S \mathcal{S} S,掩码就是被条件化实体所占像素区域的并集:

M ( S ) = ⋃ i ∈ S e i M(\mathcal{S}) = \bigcup_{i \in \mathcal{S}} e_i M(S)=i∈S⋃​ei​

模型接收的输入是掩码后的视频 M ⊙ V M \odot V M⊙V 与参考图 I 0 I_0 I0​,训练时从掩码分布中采样 M M M,学习条件分布 p θ ( V ∣ M ⊙ V , I 0 ) p_\theta(V \mid M \odot V, I_0) pθ​(V∣M⊙V,I0​)。这个设计直接借了 BERT 式掩码建模与 masked-image prompting 的思路:变化掩码输入,同一个模型服务多种任务。论文还引了一批"结构化掩码与条件推断"方向的近期工作作为思想同源,包括反事实世界建模与概率结构整合那条线,说明这不是一次孤立的技巧移植。

难点提示(为什么"涂灰"而不是"留空"):未揭示区域被统一设为均匀灰色背景,而不是随机噪声或纯黑。灰色是 latent 空间里最"无信息"的取值,既不会被模型误读成真实物体(黑色可能被当成阴影或深色物体),也不会像噪声那样干扰 VAE 编码。这是个很小但很实际的工程选择。

2.3 主动实体与被动实体的划分

在机器人语境下,把实体分成两类很方便。主动实体(active)通过自身能动性作用于场景,比如机械臂或人手;被动实体(passive)的运动源自与主动实体的交互,比如被操作的杯子。令 A \mathcal{A} A 索引主动实体, P = { 1 , … , n } ∖ A \mathcal{P} = \{1,\ldots,n\} \setminus \mathcal{A} P={1,…,n}∖A 索引被动实体。这个划分立刻给出两种自然用法。

前向模型取 S = A \mathcal{S} = \mathcal{A} S=A,也就是把主动实体的整条时空轨迹交给模型、让它推断被动实体的响应。这对应标准的动作条件动力学建模:给出机器人怎么动,模型模拟它对场景造成的后果。与以往条件在低维动作指令上的做法不同,这里的主动条件是以掩码视频形式提供的,因此与具身无关。

p ( { e i } i ∈ P   ∣   { e j } j ∈ A , I 0 ) p\bigl(\{e_i\}_{i \in \mathcal{P}} \,\big|\, \{e_j\}_{j \in \mathcal{A}}, I_0\bigr) p({ei​}i∈P​ ​{ej​}j∈A​,I0​)

逆向模型取 S = P \mathcal{S} = \mathcal{P} S=P,条件在被动实体上、预测主动实体。论文强调这个方向在传统动作条件世界模型里没有对应物:用户指定的是世界里想要发生什么,模型负责恢复出与之一致的智能体行为。用低维动作向量作条件的模型根本无法接受这类查询,因为"杯子应该沿这条轨迹移动"没法编码成末端位姿序列。

p ( { e i } i ∈ A   ∣   { e j } j ∈ P , I 0 ) p\bigl(\{e_i\}_{i \in \mathcal{A}} \,\big|\, \{e_j\}_{j \in \mathcal{P}}, I_0\bigr) p({ei​}i∈A​ ​{ej​}j∈P​,I0​)

论文特别强调,主动/被动的区分是描述我们怎么用模型的方便说法,不是模型本身的属性。模型只在掩码视频补全任务上训练,从没有显式的"能动性"概念,推理时任意子集 S \mathcal{S} S 都能选。而事实上作者只用了描绘主动机器人实体的掩码来训练,模型却零样本泛化到了以被动实体为条件的查询。

在这里插入图片描述

3. 第一条核心机制:分割式数据构建

3.1 用 SAM 直接抠机器人

数据集用两条互补的路线构建掩码条件,第一条是分割。给定任意一段视频,用 Segment Anything 配 “A robotic arm” 这个文本提示就能把机械臂从场景里分离出来,不需要相机标定,也不需要知道视频里究竟是哪一款机器人。作者在 DROID 上跑这条路线,好处是极其通用:模型学到的是"把缺失区域补全"这件事本身,等价于建模场景里所有实体的联合分布。

分割路线有两个真实的局限。第一,测试时用户很难手工提供实体的精确分割掩码——你不能指望一个想让机器人"把杯子挪到桌角"的用户先画出机械臂逐帧轮廓。第二,机器人被遮挡的区域会隐式泄漏原视频的场景动力学信息:如果掩码是从真实视频抠出来的,那么"机械臂后面挡着什么"这个信息已经通过掩码边界的形状偷偷传给了模型。这两点合起来就是引入第二条路线的动机。

3.2 数据规模与失败轨迹

训练数据由 DROID 的真实视频与 RoboCasa 的仿真数据混合而成,规模是约 1000 条 DROID 演示(配两路外部相机、每条都过一遍分割与渲染两条流水线)与约 4000 个 RoboCasa 样本(覆盖所有不含导航的任务),总时长约 15 小时。为了保证验证与测试不含训练场景,作者把 CLVR 与 RAD 两个实验室的全部场景从训练集里排除,留作测试。

这里值得注意的一个细节:两个数据集都刻意保留了失败轨迹,RoboCasa 那边甚至专门生成失败案例塞进训练数据。原因是这个模型要被用来做规划——它必须能生成"这条动作没抓住杯子"这种反事实结果,如果只在成功轨迹上训练,模型会学到"机械臂伸过去就一定成功"的偏置,那 Best-of-N 筛选就完全失效了。后面实验里观察到的正向偏置,说明这一步做得还不够彻底。

3.3 提示词被刻意压扁

底座视频模型需要文本提示作为条件,但作者只给了一个高层提示(high-level prompt),目的是确保模型只依赖视觉条件信号来生成视频。这一步很容易被忽略却很关键:如果提示词写得很详细(“机械臂抓起马克杯放到桌上”),模型完全可以忽略掩码、纯靠文本编出一段合理视频,那掩码接口的可控性就无从验证。仓库的 infer.py 里默认的负面提示词是一长串中文画质描述词,正面提示词则由用户传入,示例里也只是 "a robot arm picks up a mug" 这种粗粒度句子。

# inference/infer.py — 默认负面提示词与底座模型 id
DEFAULT_NEGATIVE_PROMPT = (
    "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,"
    "最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,"
    "画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,"
    "杂乱的背景,三条腿,背景人很多,倒着走"
)

BASE_MODEL_ID = "PAI/Wan2.2-Fun-A14B-Control"

这段代码本身没什么算法含量,但它透露了两件事。一是负面提示词整段沿用了 Wan 官方的默认值,作者没有为机器人场景做定制,进一步印证"底座不动、只动接口"的取向。二是 BASE_MODEL_ID 指向 ModelScope 上的 PAI/Wan2.2-Fun-A14B-Control,说明这套东西的复现门槛主要是显存而不是数据——底座权重是公开的。

4. 第二条核心机制:渲染式数据构建

4.1 为什么必须有渲染路线

分割路线解决不了"用户如何在测试时指定任意动作轨迹"这个问题。渲染路线的思路是把机器人的 URDF 网格按记录的状态对齐到输入视频上,用渲染结果当掩码条件。一旦模型学会接受渲染网格作为条件,推理时就可以渲染任意动作轨迹——包括从策略里采样出来、现实中从未执行过的候选动作。这是整套规划应用的前提。

代价是渲染路线需要机器人状态与相机标定。DROID 恰好记录了每一时刻的机器人状态,作者按 PointWorld 的协议细化相机标定,让 URDF 与输入轨迹精确对齐。在 RoboCasa 仿真里更简单:只渲染机器人、排除场景其余部分即可。论文明确说这两条路线是互补关系——分割通用但不可控,渲染可控但需标定且只能渲染机器人本身。

4.2 半透明渲染与红色夹爪指

渲染细节里有两个很具体的设计。第一,机器人只用半透明渲染,目的是让模型看到完整的机器人而不受自遮挡影响——如果按不透明渲染,机械臂前段会挡住后段,模型拿到的信息与分割掩码一样是残缺的。第二,夹爪手指被涂成亮红色,让视频模型能轻易观察到夹爪的开合动作。

工程价值(红色夹爪这一笔):抓取任务的成败几乎全押在夹爪张合的那几帧上,而夹爪在整幅画面里只占几十个像素。给它一个与机械臂主体高对比的颜色,等于在条件信号里给这几十个像素加了权重。这种"用颜色编码语义"的技巧在可控生成里很常见,但放在机器人条件上确实是必要的一步。

在这里插入图片描述

4.3 真机数据里的标定链

真机实验用的是配自制 3D 打印柔性末端执行器的 Franka Panda,通过自制 GELLO 界面遥操作采集演示。每次 rollout 从一个初始场景状态 s 0 s_0 s0​ 出发,同步记录机器人状态与两路 RGB 观测。传感配置是两台 ZED Mini 2,一台腕装提供近末端的自我中心视角,一台外置提供工作空间的第三人称固定视角,都以 720p、30 Hz 记录。这里要留意"自制末端执行器"这个细节——它让这批真机数据天然成为跨具身测试集,因为 DROID 里的 Franka 装的是原厂夹爪。

标定链条围绕一枚桌面固定的 AprilTag 展开。先让 Franka 走一批不同位姿、用腕部相机拍标签:正运动学给出末端坐标系 T E B ( q t ) T^{B}_{E}(q_t) TEB​(qt​),AprilTag 检测给出标签在腕相机坐标系里的位姿,两者联立解出手眼变换 T C w r i s t E T^{E}_{C^{\mathrm{wrist}}} TCwristE​ 与标签在机器人基座系的位姿 T A B T^{B}_{A} TAB​。外置相机再观测同一枚标签,由已知的 T A B T^{B}_{A} TAB​ 直接反推:

T C e x t B = T A B ( T A C e x t ) − 1 T^{B}_{C^{\mathrm{ext}}} = T^{B}_{A}\left(T^{C^{\mathrm{ext}}}_{A}\right)^{-1} TCextB​=TAB​(TACext​)−1

标定质量的验证方式很朴素:把机器人模型投影回相机图像,肉眼检查与观测帧是否对齐。这里的关键是,这条标定链的精度直接决定掩码条件的质量——URDF 渲染偏了几个像素,模型就会学到错位的动作—结果对应关系。论文把标定列为方法的固有限制之一,是诚实的。

4.4 轨迹的存储格式

每次 rollout 存成轨迹 τ = { ( I t e x t , I t w r i s t , q t ) } t = 1 T \tau = \{(I^{\mathrm{ext}}_t, I^{\mathrm{wrist}}_t, q_t)\}_{t=1}^{T} τ={(Itext​,Itwrist​,qt​)}t=1T​,其中两路 RGB 观测加测得的关节状态。所有数据以 HDF5 存放,一个文件对应一次 rollout,内含 RGB 帧、关节状态、时间戳与标定元数据,元数据包括两路相机内参 K e x t , K w r i s t K_{\mathrm{ext}}, K_{\mathrm{wrist}} Kext​,Kwrist​ 以及三个外参变换。需要视觉掩码时,把存下来的关节状态与 Franka 运动学模型、自制末端执行器几何、标定相机参数一起送进渲染器,投影到图像平面。这套组织方式意味着掩码不是预先烘死在数据里的,而是按需重渲染的,换相机视角或换末端几何都不用重采数据。

5. 训练:两个专家、两份 LoRA

5.1 底座是一个双专家 MoE

底座模型 Wan-Fun-Control 2.2 14B 不是单个 DiT,而是一个双专家混合结构:一个 high-noise DiT 负责扩散早期的高噪声区间,一个 low-noise DiT 负责后期的低噪声区间,切分点在时间步边界 0.358。这个结构上的事实直接决定了适配方式——必须训练两份 LoRA,一份灌进 pipe.dit(高噪声专家),一份灌进 pipe.dit2(低噪声专家)。论文正文只说"rank 256 的 LoRA",仓库把这件事说透了。

掩码条件视频用与视频模型相同的 autoencoder 编码,条件注入机制选的是拼接(concatenation)。作者给的理由很直接:条件信号与目标输出在空间上是对齐的,拼接是最自然的选择——不需要 cross-attention 去学"条件的哪一部分对应画面的哪一块",因为它们逐像素一一对应。这个选择顺带省掉了一整套需要调的东西:没有额外的适配器网络、没有注意力层要插、没有条件强度的调度曲线。

5.2 训练脚本读出来的超参

仓库的 training/train_control.sh 是一份对 DiffSynth-Studio 原版 trainer 的薄封装,作者在注释里明确写了"我们没有修改 trainer",这种钉版本的做法对复现很友好,避免了上游 trainer 接口漂移导致的参数名不匹配。脚本的核心逻辑是把同一条 accelerate 命令按专家跑两遍,唯一变化的是时间步边界:

# training/train_control.sh — 按噪声专家分别训练 LoRA
train_expert () {
  local noise=$1 max_b min_b
  if [ "$noise" = "high" ]; then max_b=0.358; min_b=0; else max_b=1; min_b=0.358; fi
  echo "=== training ${noise}-noise expert (timestep boundary ${min_b}-${max_b}) ==="
  accelerate launch --mixed_precision=bf16 --num_processes=${NUM_GPUS} --multi_gpu "${TRAIN_PY}" \
    --dataset_metadata_path "${DATASET_CSV}" \
    --data_file_keys "video,control_video,reference_image" \
    --height 480 --width 832 \
    --learning_rate 1e-4 \
    --num_epochs 5 \
    --lora_base_model "dit" \
    --lora_target_modules "q,k,v,o,ffn.0,ffn.2" \
    --lora_rank 256 \
    --extra_inputs "control_video,reference_image" \
    --max_timestep_boundary ${max_b} \
    --min_timestep_boundary ${min_b}
}

这段脚本把论文里几句话展开成了可执行细节。LoRA 挂在 DiT 的注意力四个投影(q,k,v,o)与 FFN 的首末两层(ffn.0,ffn.2)上,rank 256 在 LoRA 语境里算相当高——通常 8 到 64 就够做风格适配,256 说明作者要学的是一种新的条件语义而不是画风微调。分辨率 480×832、学习率 1e-4、5 个 epoch,配合论文里"8 张 H200、batch size 4、约 10000 步、4 天"的描述,整个训练成本落在学术组能承担的量级。

难点提示(时间步边界为什么要传给训练器):扩散模型的每一步去噪难度不同。高噪声阶段是"从雪花里定大致构图",低噪声阶段是"把边缘和纹理磨清楚"。MoE 把这两件事交给两个网络,训练时就必须告诉每个专家只在自己负责的噪声区间内采样时间步,否则低噪声专家会被高噪声样本带跑。--max_timestep_boundary 与 --min_timestep_boundary 就是这道闸门。

5.3 数据接口只有四列

训练需要的元数据 CSV 只有四列:prompt、reference_image、video、control_video。其中 video 是真实目标片段,control_video 是条件片段(渲染的 URDF 机器人),reference_image 是第一帧真实画面。这个接口的简洁程度值得单独说一句:整套方法对外暴露的契约就是"一段条件视频 + 一张首帧 + 一句话",任何能产出这三样东西的机器人平台都能接进来,不需要暴露动作空间维度、控制频率或运动学结构。

换句话说,掩码视觉动作把跨具身问题从"统一动作空间"降级成了"统一渲染管线"。前者是一个悬置多年的开放问题——不同自由度、不同控制模式、不同底盘构型的机器人怎么共享一个动作空间,至今没有公认答案;后者是一个有成熟工具链的工程任务,URDF、渲染器、手眼标定都是现成的东西。这个降级本身就是这篇论文最实际的贡献。

6. 推理:条件视频进、RGB 视频出

6.1 双 LoRA 装载

推理脚本先构建完整 pipeline(两个专家 DiT + T5 文本编码器 + VAE),然后把两份 LoRA 分别装到两个专家上。这一步是整套推理流程里唯一与常规 Wan 推理不同的地方,其余部分完全走 DiffSynth 的现成管线。把这段代码与论文的形式化并排看,能清楚地看到抽象符号如何落成函数实参:

# inference/infer.py — 双专家 LoRA 装载与生成调用
pipe = build_pipeline(args.base_model_id, low_vram=args.low_vram)
pipe.load_lora(pipe.dit,  args.lora_high, alpha=args.lora_alpha)   # 高噪声专家
pipe.load_lora(pipe.dit2, args.lora_low,  alpha=args.lora_alpha)   # 低噪声专家

control = load_frames(args.control_video, args.height, args.width, args.num_frames)

if args.reference_image and os.path.exists(args.reference_image):
    reference_image = Image.open(args.reference_image).convert("RGB").resize((args.width, args.height))
else:
    reference_image = control[0]      # 缺省退化为条件视频第 0 帧

output_video = pipe(
    prompt=args.prompt,
    negative_prompt=args.negative_prompt,
    control_video=control,            # 掩码视觉动作
    reference_image=reference_image,  # 初始场景 I_0
    height=args.height, width=args.width,
    num_inference_steps=args.num_inference_steps,
    seed=args.seed,
    tiled=True,
)

这段代码把论文的形式化落成了三个实参:control_video 就是 M ⊙ V M \odot V M⊙V,reference_image 就是 I 0 I_0 I0​,prompt 是那个被刻意压扁的高层文本。reference_image 缺省退化为条件视频第 0 帧这个分支很有意思——在纯渲染条件下,第 0 帧只有一个半透明机器人和灰背景,缺少真实场景信息,所以真机使用时必须显式传入真实首帧。

6.2 默认生成配置与成本

脚本默认值给出了这套系统的实际工作点:--num-frames 81、--num-inference-steps 50、--fps 15、分辨率 480×832、--seed 123。81 帧、15 fps 对应约 5.4 秒的片段,这也解释了论文里 VLM 评估协议写的"以 15 fps 采样以便 Gemini 看到全部 81 帧"。50 步去噪 × 14B 双专家,单次生成的开销不小,这正是论文在局限里说"多候选生成成本较高"的来源。

显存这一侧仓库也给了退路。14B 双专家加 T5-XXL 编码器同时驻留显存对消费级卡不现实,--low-vram 分支把闲置权重卸载到磁盘,用速度换显存,代价是每次专家切换都要走一次磁盘读取,在 50 步去噪的过程中这个切换会发生多次,因此实测速度差距不小:

# inference/infer.py — 低显存模式的分级卸载
if low_vram:
    kwargs = dict(
        offload_dtype="disk",   offload_device="disk",     # 闲置权重落盘
        onload_dtype=torch.bfloat16,      onload_device="cpu",
        preparing_dtype=torch.bfloat16,   preparing_device="cuda",
        computation_dtype=torch.bfloat16, computation_device="cuda",
    )

四级流水(disk → cpu → cuda 预备 → cuda 计算)是 DiffSynth 的现成机制,作者只是把开关暴露出来。对想复现的人来说,这行代码的含义是:没有 H200 也能跑推理,只是慢,训练那 8 卡 4 天才是真门槛。脚本里还有一句 vram_limit 按当前卡的总显存减 2 GB 自动设定,属于把 DiffSynth 的调参交给运行环境自己决定。

6.3 权重与复现路径

权重下载脚本明确了 Hugging Face 仓库里只有两个文件——masked_world_lora_high.safetensors 与 masked_world_lora_low.safetensors,底座由 infer.py 经 ModelScope 自动拉取,这个脚本不碰底座。文件名里的 masked_world 也保留了项目早期的内部命名,算是个小小的考古痕迹。脚本还留了 --revision 参数支持按 git tag 取特定版本的权重,方便日后作者更新时锁定复现所用的那一版。

# inference/download_weights.py — 只下载两份 LoRA 适配器
DEFAULT_REPO_ID = "HadiZayer/masked-visual-actions"
FILES = ["masked_world_lora_high.safetensors", "masked_world_lora_low.safetensors"]

paths = [
    hf_hub_download(repo_id=args.repo, filename=f, revision=args.revision, local_dir=args.out)
    for f in FILES
]

…详情请参照古月居

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐