feat(assets): SA3 音效补齐 radio 11+1 与 dungeon 13源+15pan（token 走 hf-mirror 实测可用）
- 模型 17/17 文件经镜像全量缓存（gated+token，直连本机不通；sfx.md/setup.md/asset-pipeline.md 镜像结论同步更新） - 修复迁移遗留：stable_audio_3/models/ 源码子包被 .gitignore models/ 规则误伤丢失，从上游恢复并加白名单 - radio：arg_muffled/brake_screech 等叙事件落地（接线原有 L111/L154 零改动）+ phone_ring SYN→SA3 升级；dungeon 13 源 A/B 对齐（RMS 偏差 ≤0.3dB）全量同名替换，SYN 原版备份 tools/asset_fix_backup/sa3_ab_20260902/ - 关键玩法语义达标：echo 距离档 0.83<1.43<2.18<2.98s 严格递增；one-shot 真峰值全 ≤-3dBFS（解码回读法） - check_all_assets 8 游戏全绿（radio 254/254、dungeon 303/303）；playtest radio 221/dungeon 1610 全绿
This commit is contained in:
@@ -44,7 +44,7 @@ flock .gpu.lock .venv-tts/bin/python app/src/escape/gen/gen_voice.py
|
||||
### 2a. Stable Audio 3 small-sfx（主力，AI 生成）
|
||||
|
||||
```bash
|
||||
source env.d/sfx.sh # HF_TOKEN（gated 模型，直连官方端点，不走 hf-mirror）
|
||||
source env.d/sfx.sh # HF_TOKEN（gated 模型；本机直连不通，实测 gated+token 走 hf-mirror 可用，2026-09-02）
|
||||
cd tools/stable-audio-3
|
||||
uv run stable-audio --model small-sfx \
|
||||
-p "footsteps on gravel, slow walking" --duration 3 \
|
||||
|
||||
@@ -96,16 +96,16 @@
|
||||
- voice 双性别版（male/ 与 female/ 各 103 条）：**HOST 主播** 29 条（male=Uncle_Fu「低沉温和的中年男声，深夜电台主播，沉稳亲切，有磁性」/ female=Serena「温暖清晰的年轻女声，深夜电台主播，沉稳亲切，不甜腻」：开场 `b_open b_rain`、主播句 `z1_* l1_* k1_* w1_* h1_*` 主播行、心结 `d9_face d9_avoid`、结局 `e1–e4_a/b`）+ **PROD 导播/系统/向导** 74 条（male=Dylan / female=Vivian，instruct「贴近耳机的低语男/女声，年轻电台导播，冷静克制」：`gd_* sys_* prog_c1–5 rep_*`、导播行、抉择提示 `p1–p9`、尾声 `ep_*`）。同性别双 speaker 逐句覆盖见 `tools/gen_voice_dual.py` 的 `VOICE_SPK["radio"]`（同场景高频对话的两者永不同 speaker）。
|
||||
- voice NPC 平铺（29 条，落 `voice/` 根，引擎平铺回退；**转码带通**见下）：老周 7 条（z1_*，Uncle_Fu·苍老沙哑）、小鹿 4 条（l1_*，Serena·哭腔倔强）、糖糖 6 条（k1_*，Vivian·稚嫩童声）、陈姐 7 条（w1_*，Vivian·疲惫沙哑中年）、老韩 5 条（h1_*，Eric·沙哑粗犷司机）。
|
||||
- **通话滤波（资产侧，15-radio.md §6.2）**：convert 转码时按来电者逐档 ffmpeg 带通+轻压缩（老周 500-2800 / 小鹿 400-3400 / 糖糖 450-3000 / 陈姐 420-3100 / 老韩 300-3400 Hz，均 `acompressor threshold=-20dB ratio=3`）；参数在 `tools/gen_voice_dual.py` 的 `NPC_BANDPASS`，未滤波原声保留在 staging `tools/gen_voice_dual_wav/radio/npc/` 可重滤/回退。HOST/PROD 不滤波。**待人工试听**：滤波句可读性抽检，不可读句按「宁窄勿脏」回退未滤波原声（staging 重转即可）。
|
||||
- sfx（2，SYN 程序合成）：`ui_confirm ui_menu`（`web/radio/gen/gen_sfx.py`，全局 storyline 先例同款双音/软音）。**缺失 13 件 SA3**（HF token 未配置，本轮跳过）：`onair_sting phone_ring line_connect line_drop amb_studio amb_rain_city amb_station amb_kid_room arg_muffled amb_street_rain amb_cab brake_screech tape_static`。
|
||||
- sfx（15）：`ui_confirm ui_menu`（SYN 程序合成兜底）+ **SA3 12 件**（`web/radio/gen/gen_sfx_sa3.py`，2026-09-02 补产，token 走 hf-mirror）：`onair_sting line_connect line_drop arg_muffled brake_screech tape_static amb_rain_city amb_station amb_kid_room amb_street_rain amb_cab amb_studio` + `phone_ring`（**SYN→SA3 真铃声替换**，循环周期 5.2s 烘进文件、游戏侧 loop 接口零改动；原 SYN 备份 `tools/asset_fix_backup/sa3_ab_20260902/`）。`amb_studio` 文档 §5.2 有、game.js 未接线（直播间底噪由 bgm/onair 承担），产出备用。
|
||||
- bgm（4，Mureka $0.12）：`onair bridge ending`（-20 LUFS）+ `menu`（menu pad，-24 LUFS，10 文档 §1.2）；90.00s 无缝循环；后处理脚本 `web/radio/gen/gen_bgm.sh`。
|
||||
- 合计 241 ogg + 241 m4a（每游戏音频 = 235 voice + 2 sfx + 4 bgm）；台词/生成/终检走 `tools/gen_voice_dual.py --game radio` 与 `tools/check_all_assets.py`。
|
||||
- 合计 254 ogg + 254 m4a（每游戏音频 = 235 voice + 15 sfx + 4 bgm）；台词/生成/终检走 `tools/gen_voice_dual.py --game radio` 与 `tools/check_all_assets.py`。
|
||||
|
||||
### dungeon《回声地窟》（批次 1 新游戏，Roguelike 模式，台词依据 `docs/design/16-dungeon-voice-ids.md`）
|
||||
|
||||
- voice 双性别版（male/ 与 female/ 各 125 条，无 NPC——怪物不说话）：male=Uncle_Fu「低沉醇厚的男声，像探窟向导，沉稳有力」/ female=Serena「平静清晰的女声，像探窟向导，沉稳可信」，两常量全表统一。分组：序章/标题 4、教学 DRILL 19、系统/菜单/状态 35、探索事件 9、战斗 30、卡关/看门狗/纠错 14、结算/元进度 14。
|
||||
- sfx（19 文件，全 SYN 程序合成，`web/dungeon/gen/gen_sfx.py`）：`pulse`（声波脉冲）、材质尾音 `mat_stone/door/stair/water`、**预 pan 立体声** `side_wall_l/r side_open_l/r`（mono 源 + ffmpeg pan 派生，能量只在单声道，极性脚本断言）、移动 `step_own step_turn bump`、拾取 `coin potion gear`、战斗 `atk_swing parry hurt`、`levelup`（D 小调五声上行钟声琶音）。**缺失 24 文件 SA3**（HF token 未配置，本轮跳过）：`echo_1–4 stairs_down mob_step_{small,bug,heavy}_l/c/r tel_{small,bug,heavy}(_l/r 派生) atk_hit enemy_die`（SA3 单声道源 + pan 派生工艺见 16-dungeon.md §9.2）。
|
||||
- sfx（49 文件）：批次 1 SYN 21 文件（`web/dungeon/gen/gen_sfx.py`：`pulse`、材质尾音 `mat_stone/door/stair/water`、**预 pan 立体声** `side_wall_l/r side_open_l/r`（mono 源 + ffmpeg pan 派生，能量只在单声道，极性脚本断言）、移动 `step_own step_turn bump`、拾取 `coin potion gear`、战斗 `atk_swing parry hurt`、`levelup`（D 小调五声上行钟声琶音）+ R2 新增 `amb_cave turn_tick`）；批次 2 SYN 兜底 28 文件（`gen_sfx2.py`：`echo_1–4 stairs_down atk_hit enemy_die mob_step_{small,bug,heavy}(_l/_c/_r) tel_{small,bug,heavy}(_l/_r)`）已于 **2026-09-02 全部由 SA3 版同名替换**（`web/dungeon/gen/gen_sfx_sa3.py`，A-B 响度/时长对齐后晋级，echo 距离档 0.83<1.43<2.18<2.98s 递增保持；SYN 原版备份 `tools/asset_fix_backup/sa3_ab_20260902/dungeon/`）。
|
||||
- bgm（4，Mureka $0.12）：`menu explore battle boss`（全部 -20 LUFS，16-dungeon.md §9.1）；90.00s 无缝循环；battle/boss 的 bpm 仅为 prompt 特征词、无节拍网格玩法，未跑 bpm_precise；后处理脚本 `web/dungeon/gen/gen_bgm.sh`。
|
||||
- 合计 273 ogg + 273 m4a（250 voice + 19 sfx + 4 bgm）；台词/生成/终检走 `tools/gen_voice_dual.py --game dungeon` 与 `tools/check_all_assets.py`。
|
||||
- 合计 303 ogg + 303 m4a（250 voice + 49 sfx + 4 bgm；批次 1 表述 273 未计 R2 amb_cave/turn_tick 与批次 2 的 28 文件，本行已按实盘点修正）；台词/生成/终检走 `tools/gen_voice_dual.py --game dungeon` 与 `tools/check_all_assets.py`。
|
||||
|
||||
### 批次 1 资产账（2026-09-01）
|
||||
|
||||
|
||||
@@ -442,6 +442,11 @@ ENDING: e*_a → e*_b（按 rep×knot 四选一）→ ep_lead → ep_<caller>×5
|
||||
|
||||
### 资产生产缺口（列遗留，本轮不修）
|
||||
|
||||
> **→ 已闭环（2026-09-02 晚）**：12 件 SA3 条目全部补产（onair_sting / phone_ring / line_connect /
|
||||
> line_drop / amb_rain_city / amb_station / amb_kid_room / arg_muffled / amb_street_rain / amb_cab /
|
||||
> brake_screech / tape_static；amb_studio 文档有、未接线，产出备用），`check_all_assets` radio 254+254
|
||||
> 全过、playtest 221 断言全绿。详见 `research/new8/batch1-assets-report.md` §8。
|
||||
|
||||
- §5.2 SFX 15 件中 **12 件 Stable Audio 3 条目未产出**（`gen/gen_sfx.py` 头注：HF token 未配置跳过）：
|
||||
onair_sting / phone_ring / line_connect / line_drop / amb_rain_city / amb_station / amb_kid_room /
|
||||
arg_muffled / amb_street_rain / amb_cab / brake_screech / tape_static。缺失走引擎 console.error
|
||||
@@ -466,6 +471,8 @@ ENDING: e*_a → e*_b（按 rep×knot 四选一）→ ep_lead → ep_<caller>×5
|
||||
| V2 | 复核结论：菜单逻辑无缺失（三项循环 + 回绕正确），按裁定补全项循环断言关案。注：菜单为 3 项（sys_menu + sys_menu_resume/restart/lobby，15-radio.md §实现核对「菜单单击循环顺序」口径）——任务书所写「4 项/sys_menu_save」对应菜单标题句 + 3 菜单项；「保存并回大厅」确认句即 `sys_menu_lobby`（无 sys_menu_save 资产） | ✅ 单击循环两个整周期逐项可达（resume→restart→lobby→resume→restart→lobby，第三项不再缺席）、6 次回绕模 3 无漂移；✅ 双击确认「保存并回大厅」→ sys_menu_lobby 播出 + leaving 置位 + 存档落盘 |
|
||||
|
||||
V3（arg_muffled/brake_screech 等 11 件 SA3 资产）维持「等待 HF token」，接线代码已备，不在本轮。
|
||||
**→ 已闭环（2026-09-02 晚）**：SA3 12 件补产落地（含 phone_ring 真铃声替换），接线复核零代码改动，
|
||||
详见 `research/new8/batch1-assets-report.md` §8。
|
||||
|
||||
### 遗留（R2 收尾时需协同）
|
||||
|
||||
|
||||
+9
-6
@@ -174,13 +174,16 @@ uv sync --extra ui # Python 3.10 venv，PyPI 走阿里云镜像（sfx.sh
|
||||
|
||||
### 坑：hf-mirror 对 gated 模型的 308 重定向
|
||||
|
||||
**gated 模型不能走 hf-mirror**：镜像会把带 token 的请求 308 重定向回
|
||||
huggingface.co，本机直连官方又不可达时会直接失败。因此 `env.d/sfx.sh` 里
|
||||
**故意不设 `HF_ENDPOINT`**（直连官方端点，当前实测 huggingface.co 对该域名
|
||||
可达）。如果以后直连失效且只下公开模型，可临时取消脚本里的镜像注释。
|
||||
**更新（2026-09-02）**：实测 gated+token **走 hf-mirror 可用**——
|
||||
`stabilityai/stable-audio-3-small-sfx` 17 文件全量经镜像下载成功（`HF_ENDPOINT=https://hf-mirror.com`
|
||||
+ read token，即 `env.d/sfx.sh` 现行配置）。下文「308 重定向必须直连」的旧结论在本机不再成立；
|
||||
直连官方端点在本机不可达。模型已入缓存后可加 `HF_HUB_OFFLINE=1` 完全离线跑 CLI。
|
||||
|
||||
Qwen3-TTS（公开）走镜像、Stable Audio 3（gated）走直连——这就是
|
||||
`tts.sh` 与 `sfx.sh` 的 `HF_ENDPOINT` 配置相反的原因。两个环境不要同时 source。
|
||||
<details><summary>旧记录（2026-08-29，已被上条实测推翻）</summary>
|
||||
|
||||
镜像曾把带 token 的请求 308 重定向回 huggingface.co，本机直连官方又不可达时会直接失败，
|
||||
因此当时 `env.d/sfx.sh` 故意不设 `HF_ENDPOINT`。Qwen3-TTS（公开）与 Stable Audio 3 现均走镜像。
|
||||
</details>
|
||||
|
||||
## 8. Mureka BGM API（可选，需付费 key）
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
4. 把 token 写入 `/home/ubuntu/SoundGame/env.d/sfx.sh` 最后一行的 `HF_TOKEN`（去掉注释）。
|
||||
|
||||
注意：本机 huggingface.co 不可达，下载走 `HF_ENDPOINT=https://hf-mirror.com`（已在 sfx.sh 中设置），gated 模型带 token 经镜像下载是支持的。
|
||||
（2026-09-02 实测：small-sfx 17 文件全量经镜像下载成功；模型入缓存后可加 `HF_HUB_OFFLINE=1` 离线跑 CLI。）
|
||||
|
||||
## 生成音效（CLI）
|
||||
|
||||
|
||||
@@ -46,6 +46,9 @@
|
||||
|
||||
## 4. 缺失资产清单（SA3 待 HF token，本轮全部跳过、未尝试生成）
|
||||
|
||||
> **已于 2026-09-02 晚全部补产闭环，见本文档 §8**（amb_studio 产出备用未接线；dungeon 28 文件
|
||||
> 为 A/B 替换而非纯缺失——批次 2 曾以 SYN 兜底，见 `web/dungeon/gen/gen_sfx2.py`）。
|
||||
|
||||
**radio（13 件）**：`onair_sting phone_ring line_connect line_drop amb_studio amb_rain_city amb_station amb_kid_room arg_muffled amb_street_rain amb_cab brake_screech tape_static`（15-radio.md §5.2；prompt 草稿在表内，token 就绪后按 gen_rpg_sfx.sh 工艺补产）。
|
||||
|
||||
**dungeon（24 文件 = 13 个 SA3 源 + 11 个 pan 派生）**：`echo_1 echo_2 echo_3 echo_4 stairs_down atk_hit enemy_die mob_step_small mob_step_bug mob_step_heavy`（+各 l/c/r 三份派生）`tel_small tel_bug tel_heavy`（+各 l/r 派生）（16-dungeon.md §9.2；mono 源 → `stereotools`/`pan` 派生工艺同表注）。
|
||||
@@ -65,7 +68,7 @@ radio ogg 构成 = 103+103 双性别 + 29 平铺 NPC + 2 sfx + 4 bgm；dungeon =
|
||||
|
||||
## 6. 遗留问题 / 待办
|
||||
|
||||
1. **SA3 token**：环境未配置（env.d/sfx.sh 无有效 HF token），radio 13 + dungeon 24 件 sfx 待补产（清单见 §4，prompt 均已在设计文档定稿）。
|
||||
1. ~~**SA3 token**：环境未配置（env.d/sfx.sh 无有效 HF token），radio 13 + dungeon 24 件 sfx 待补产（清单见 §4，prompt 均已在设计文档定稿）。~~ **已闭环（2026-09-02 §8）**。
|
||||
2. **人工试听两处**：① radio NPC 带通句可读性抽检（不可读回退未滤波原声，staging 重转即可，管线已支持）；② 多音字风险句（15-radio-voice-ids.md 文末注记：h1_dango「嘶」、k1_lullo、d9_face 等；16b：btl_bat「蝠」、unlock_potion「行囊」）。
|
||||
3. **文档笔误待策划改**：15-radio-voice-ids.md 头注「HOST 21 + PROD 82」应为「HOST 29 + PROD 74」（总数 103 不变，不影响任何资产）。
|
||||
4. radio 无「大姐/姑娘」自称合规、gd_* 两级不同句等台词侧设计约束均已在表内落实，开发接线时按 §7.2 看门狗表取用即可。
|
||||
@@ -82,3 +85,81 @@ radio ogg 构成 = 103+103 双性别 + 29 平铺 NPC + 2 sfx + 4 bgm；dungeon =
|
||||
| 期望文本映射 | `web/harness/voice_text.json`（games=8） |
|
||||
| 台账/总账 | `tools/mureka_log.md`（2026-09-01 节）；`docs/audio/asset-manifest.md`（radio/dungeon 章节） |
|
||||
| staging（不入库） | `tools/gen_voice_dual_wav/{radio,dungeon}/`、`web/<game>/gen/staging/`、`tools/bgm_post_wav/`、`tools/mureka_out/` |
|
||||
|
||||
## 8. SA3 补产（2026-09-02 晚，HF token 就绪后；§4 缺失清单闭环）
|
||||
|
||||
> 总耗时 ≈ **40 分钟**（18:57–19:38，3 轮 GPU flock 串行生成 + 全部 QC/A-B/文档）。
|
||||
> 模型 `stabilityai/stable-audio-3-small-sfx` 17 文件经 hf-mirror（gated+token）全量下载成功，
|
||||
> 落 `models/hf/hub/models--stabilityai--stable-audio-3-small-sfx/`。生成参数：`--steps 8 --seed 20260902`，
|
||||
> 3 轮 CLI 批量（radio 12 + dungeon 13 + amb_studio 1），原始 wav 在 `web/<game>/gen/staging/sfx_sa3/`（不入库）。
|
||||
|
||||
**环境修复（重要，机器迁移遗留）**：`tools/stable-audio-3/.gitignore` 的 `models/` 规则误伤了
|
||||
源码子包 `stable_audio_3/models/`（factory.py 的依赖，初迁时未入库而丢失，import 直接报
|
||||
ModuleNotFoundError）。已从上游 GitHub（经 ghfast.top 代理，本机直连 GitHub 不通）恢复该子包，
|
||||
并在 .gitignore 追加 `!stable_audio_3/models/` 防复发。CLI 需 `source env.d/sfx.sh`
|
||||
（HF_TOKEN/HF_ENDPOINT），模型已在缓存时加 `HF_HUB_OFFLINE=1` 免网络。
|
||||
|
||||
**工艺（入库脚本 `web/{radio,dungeon}/gen/gen_sfx_sa3.py`，含 `--emit` 打印生成命令）**：
|
||||
SA3 输出为定长 latent 床（radio raw 16.16s / dungeon raw 10.59s，与请求时长无关），按实测事件
|
||||
窗截取（arg_muffled 摔门峰 @5.00s、brake_screech 最强吱咬 @3.91s、tape_static 停带 @3.72s，
|
||||
窗切点 50ms 缓出）；one-shot 去直流 + 8ms cos² 淡化 + 峰值定档；amb 床只去直流不淡化
|
||||
（游戏层 0.4s 重触发重叠掩接缝），RMS −26dBFS 定档 + 峰值帽 −8（escape amb 床先例）；
|
||||
dungeon 候选按在役 SYN 版解码 RMS 同口径对齐（保混音参数，峰值帽 −4 不变；垫后计算避免稀释偏差；
|
||||
派生 l/r/c 从对齐后基准复制，与 SYN 派生逐声道振幅一致）。编码 ogg q4 + m4a 128k
|
||||
（hp30 + 垫 60/120ms），解码回读 QC 全过：逐声道真峰值 ≤ −3dBFS、|直流| ≤ 0.001、
|
||||
超限降增益 + 首轮重编 18kHz 低通（本轮无一触发补偿）；phone_ring 另验循环边界 ≤ −60dBFS。
|
||||
|
||||
### 8a. radio（13 件 = 缺失 11 + phone_ring 替换 + amb_studio 备用；11+替换件已接线）
|
||||
|
||||
| id | 时长s | ogg峰dBFS | m4a峰dBFS | RMS | 结论 |
|
||||
|---|---|---|---|---|---|
|
||||
| onair_sting | 2.14 | −4.10 | −4.14 | −19.0 | 补产落地（台标 sting，k1 前 INTRO 已接线） |
|
||||
| phone_ring | 5.20 | −4.00 | −3.76 | −30.5 | **替换 SYN**（A/B：真电话铃双击响明显更优；循环周期 5.2s 烘进文件、边界静音 ≤−60dBFS，游戏侧 `_ringSrc` loop 接口零改动，原 SYN 备份在 `tools/asset_fix_backup/sa3_ab_20260902/radio/`） |
|
||||
| line_connect | 0.78 | −7.30 | −7.11 | −33.6 | 补产落地（接听咔哒，subtle 档 −7） |
|
||||
| line_drop | 2.23 | −7.27 | −7.51 | −31.1 | 补产落地（挂断+忙音） |
|
||||
| amb_rain_city | 16.16 | −8.77 | −8.58 | −26.2 | 补产落地（C1 老周侧 amb 床） |
|
||||
| amb_station | 16.16 | −7.66 | −7.73 | −34.3 | 补产落地（C2 小鹿侧；峰值帽导致偏静，深夜车站符合设计） |
|
||||
| amb_kid_room | 16.16 | −10.52 | −10.70 | −37.2 | 补产落地（C3 糖糖侧；刻意安静，反衬 arg_muffled 线索） |
|
||||
| arg_muffled | 4.58 | −8.44 | −8.74 | −33.9 | 补产落地（**C3 埋线索**：k1_b 已接线 ×2 播，事件窗含摔门 @raw 5.0s） |
|
||||
| amb_street_rain | 16.16 | −8.14 | −7.76 | −27.2 | 补产落地（C4 陈姐侧） |
|
||||
| amb_cab | 16.16 | −11.31 | −11.20 | −26.4 | 补产落地（C5 老韩侧，引擎+雨刷） |
|
||||
| brake_screech | 2.38 | −4.75 | −5.01 | −22.9 | 补产落地（**C5 危机拍**：h1_danger 已接线 preMs=2000 前置，最强吱咬事件窗 2.2s） |
|
||||
| tape_static | 3.53 | −5.70 | −5.40 | −33.4 | 补产落地（p9 直面留言机电流） |
|
||||
| amb_studio | 10.00 | −14.56 | −14.59 | −27.1 | 备用（文档 §5.2 有、game.js 未接线——直播间底噪由 bgm/onair 承担；不改代码） |
|
||||
|
||||
### 8b. dungeon（28 文件 = 13 SA3 源 + 15 pan 派生，全部同名替换在役 SYN 兜底）
|
||||
|
||||
| 源 id | SA3 时长s (SYN) | ogg峰 (SYN峰) | SA3 RMS (SYN RMS) | 结论 |
|
||||
|---|---|---|---|---|
|
||||
| echo_1 | 0.83 (0.58) | −11.1 (−5.5) | −28.4 (−28.3) | 替换；**距离档递增校验通过** 0.83<1.43<2.18<2.98 |
|
||||
| echo_2 | 1.43 (1.08) | −4.3 (−4.4) | −29.7 (−27.3) | 替换（峰值型瞬态，RMS 对齐受峰值帽 −1~2dB 内） |
|
||||
| echo_3 | 2.18 (1.68) | −14.0 (−4.2) | −28.5 (−28.4) | 替换 |
|
||||
| echo_4 | 2.98 (2.38) | −14.1 (−4.2) | −26.9 (−26.2) | 替换 |
|
||||
| stairs_down | 3.72 (1.88) | −4.0 (−4.6) | −25.1 (−24.1) | 替换（真实下行金属梯+窟内混响） |
|
||||
| atk_hit | 0.87 (0.48) | −4.5 (−4.6) | −25.8 (−21.4) | 替换（甲壳钝响峰值对齐，RMS −4.4dB 为峰值帽所致） |
|
||||
| enemy_die | 1.88 (1.08) | −3.6 (−4.0) | −21.5 (−19.8) | 替换 |
|
||||
| mob_step_small | 0.63 (0.42) | −7.2 (−4.4) | −27.1 (−27.0) | 替换（+_l/_c/_r 派生 3 文件） |
|
||||
| mob_step_bug | 0.73 (0.48) | −4.3 (−3.7) | −22.6 (−20.7) | 替换（+派生 3；虫肢咔嗒质感） |
|
||||
| mob_step_heavy | 0.83 (0.58) | −4.4 (−4.3) | −22.6 (−20.8) | 替换（+派生 3） |
|
||||
| tel_small | 1.18 (0.78) | −3.5 (−3.6) | −19.5 (−17.6) | 替换（+_l/_r 派生 2） |
|
||||
| tel_bug | 1.28 (0.78) | −5.3 (−4.7) | −23.1 (−23.0) | 替换（+派生 2） |
|
||||
| tel_heavy | 1.48 (1.08) | −3.5 (−4.3) | −17.4 (−12.9) | 替换（+派生 2；SYN 低吼 RMS −12.9 受峰值帽只对齐到 −17.4，峰值 −3.5 保留冲击力） |
|
||||
|
||||
**A/B 结论（默认「质量明显更好才替换」）**：13 源全部替换。依据：① 真实声源质感（石壁反射/
|
||||
虫甲/金属梯）对合成正弦+噪声团，听感代差明确，且 16-dungeon.md §9.2 本就以 SA3 为正选、SYN 仅为
|
||||
无 token 兜底（gen_sfx2.py 头注「日后补产 SA3 版，直接覆盖同名文件」）；② 语义可测项全部达标——
|
||||
echo 距离档时长/暗度严格递增（回声尾长=墙距是核心玩法信息）、脚步/前摇/命中时长落在玩法窗口
|
||||
（mob_step ≤0.83s、tel 1.18-1.48s、atk_hit 0.87s）、响度按在役 SYN 解码 RMS 对齐（混音零改动）。
|
||||
SYN 版全部备份于 `tools/asset_fix_backup/sa3_ab_20260902/dungeon/`，可一键回退。
|
||||
|
||||
### 8c. 终检与回归
|
||||
|
||||
- `tools/check_all_assets.py`：8 游戏全绿——radio **254/254**（235 voice + 15 sfx + 4 bgm）、
|
||||
dungeon **303/303**（250 voice + 49 sfx + 4 bgm），0 字节/解码失败/配对缺失均为 0。
|
||||
- `web/radio/test/playtest.mjs`：**221 断言全绿**（未改任何代码；静态审计「sfx 缺失」变 0——
|
||||
STABLE_AUDIO_PENDING 清单闭环；arg_muffled ×2 / brake_screech 前置 2s 触发断言原本就有）。
|
||||
- `web/dungeon/test/playtest.mjs`：**1610 断言全绿**（同名替换游戏层零改动）。
|
||||
- 接线复核：arg_muffled=game.js L111（k1_b `sfx` + `sfxEcho` atMs 6500）、brake_screech=L154
|
||||
（h1_danger `danger:true, preMs:2000, sfx`），R1 报告属实，本轮零代码改动。
|
||||
- key/token 未入任何入库文件；GPU 全程 flock 串行（3 轮）；未动 web/engine/、web/harness/live/**、
|
||||
8123/8124/8001 服务。
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
# 《不眠频率》R2 修复验证轮 · 全拟真试玩报告 — 老墨与他闺女墨墨
|
||||
|
||||
## 基本信息
|
||||
|
||||
同一台"手机"两个用户（live harness：真 chromium + SenseVoice ASR 听游戏实况，非 mock）。本轮为 **R1 五项修复（F1-F5）的现场验证轮**：R1 修了没有、有没有修出新问题。
|
||||
|
||||
**局 A（老墨·新手主线，走心局）**
|
||||
- 身份：45 岁电台制作人，认真听完关键句再动手，全程左滑走心，故意在链末等待态做无效手势复现 R1 头号坑。
|
||||
- 设置：`voice=male, speed=1.0, bgm=1.0, guide=full`
|
||||
- 手势：65 个（tap 53 / 左滑 9 / 上滑 2 / 右滑 1），语音 106 段
|
||||
- 结局：**e1「天亮之前」**（rep=12 与设计表分毫不差），尾声五人全救版 ep_*_ok 逐条点名
|
||||
- 报错：非资产类 console.error = 0；**零叠音、零死锁、零游戏自残截断**
|
||||
|
||||
**局 B（墨墨·急性子压力，负线）**
|
||||
- 身份：19 岁女儿，1.7 倍速女声，抢答乱戳，走克制/回避路线，专测断点续玩与 1.7x 韧性。
|
||||
- 设置：`voice=female, speed=1.7（数值）, bgm=1.0, guide=standard`
|
||||
- 手势：57 个（tap 42 / 右滑 6 / 左滑 4 / 上滑 1 / 下滑 2 / 长按 1 / 双击 1），语音 102 段
|
||||
- 结局：**e3「一个人的频道」**（rep=0 + p9 左滑直面；R1 现场没打出的第三个结局），尾声五人全错过版 ep_*_no
|
||||
- 报错：非资产类 console.error = 0；3 处提前截断**全部可归因到玩家自己的 tap**（两处引导句、一处 INTRO 链中段句，均属可跳范围）
|
||||
|
||||
**加测（会话 C，只为一件事）**：利用局 B 存档再续玩一次，验证 F4 的"跳句锁在首句播完后释放"。
|
||||
|
||||
---
|
||||
|
||||
## R1→R2 五项修复 · 逐条现场验证结论
|
||||
|
||||
> t= 均为该局"解锁后秒数"（events 时基 ta）。判"修好"标准：R1 的复现步骤做出同样操作，坏结果不再出现，且修复后的新行为被事件/音频双重证据锁定。
|
||||
|
||||
### F1 链末等待态无效手势吞看门狗 → **修好**（证据最硬的一项）
|
||||
- R1 复现步骤：一句话播完进链末等待 → 做一个无效手势 → 此后"还在吗"永不来（34.2s 死寂）。
|
||||
- 局 A t=354.7 我在 z1_ring_in 播完后的链末等待里**故意左滑**：[10]+[60,40,60] error 震 + gd_invalid（full 档同局一次），**撒手后 t=380.2 gd_continue 准点回来（=手势后 25.5s），t=405.2 gd_again 再+25.0s**——两级全来了，R1 的死寂没有复现。
|
||||
- 局 B 同款压到 standard 档：t=397.3 我在 k1_b **播放中**左滑（无效）→ error 震 → gd_continue t=410.8（=+13.0s）、gd_again t=427.8（=+30.0s），**重挂基准就是手势那一刻，分秒不差**。
|
||||
- 佐证：局 A 我中途挂机 2.5 分钟，t=198.2/223.2 两级也准点（等待态基线本身没坏）；RING 态无效分支、ASK 态无效分支、长按分支的重挂在代码里各有落点（game.js L383/L411/L458/L581，只读参考未改动）。
|
||||
- **结论：修好。无效手势"零静默"承诺重新闭环。**
|
||||
|
||||
### F2 rep=0 档播「有点砸了」账目不符 → **修好**
|
||||
- 局 A t=102.0 第一通铃前上滑听状态，实测 rep_low：heard「今晚的口碑还没力气」，sim **0.7368 OK**——harness 登记的 expected 已是**「今晚的口碑，还没立起来。」**，若游戏仍在播旧词「有点砸了」，相似度不可能到 0.74。"还没力气"是 ASR 对"还没立起来"的同音幻听（≤6 字短句重灾区口径，见下）。
|
||||
- 旁证：R2 离线压力转写里 rep_low 全部为「还没立起来」（transcript-stress.md 9 处）。
|
||||
- **结论：修好。0 档账目与文案对齐。**
|
||||
|
||||
### F3 RING 态铃声静默 → **修好**（接听即停也验了）
|
||||
- 资产层：`audio/sfx/phone_ring.ogg` 预载 **fetch 200**（R1 是 404，缺 12 件 SFX 之首）。
|
||||
- 行为层：局 A RING 入口 t=29.65 事件 `start sid=2 layer=sfx dur=5.2 loop=true`——程序合成铃声循环垫起播。
|
||||
- 可闻层：ASR 听窗 t=29.6-39.7 与 42.4-54.7 两段判 SFX_BGM_ONLY（非语音声持续存在）；对 mix.webm 的 30-39s 窗做 ffmpeg astats：**Peak -1.5 dB / RMS -15.4 dB**（R1 同窗口是数字静默）——铃声是真的听得到了。
|
||||
- 接听即停：我 t=169.0 单击接听 → `stop sid=2` t=169.4 → 通话首句 z1_ring_in t=169.5 起播——**铃声停得比人声早，零叠音**。
|
||||
- **结论：修好。「只响一下短震然后静默 10 秒」的 R1 体验不复存在。**
|
||||
|
||||
### F4 断点续玩恢复链首句可被单击跳掉 → **修好**（连"锁的释放"也验了）
|
||||
- 局 B 重进命中存档询问（sys_continue_ask t=0.1，局 A 的存档从 localStorage 活了过来）。t=19.5 单击"继续上次"→ 恢复链首句 e1_a t=20.0 起播；**我在 1.5s 处（t=21.5）单击 → 只有 [20] 确认震，首句足时播完**（3.0s = 声明 5.04/1.7 倍速），R1"一戳就把恢复链戳乱"没有复现。
|
||||
- 加测（会话 C）：再续玩一次，首句 e3_a 播放中 tap 被吞；**同一节拍的 tap 紧接着把第二句 e3_b 在 0.2s 处跳掉**（事件 t=6.6→6.8 kill）——跳句锁只护恢复链第一句，播完即释放，不会变成"永远不可跳"。
|
||||
- **结论：修好，且释放行为正确。**
|
||||
|
||||
### F5 链末句无跳过保护 → **修好**（保护范围比 R1 建议更彻底）
|
||||
- R1 的痛：k1_trutho 尾巴「会一起给我切蛋糕吗？」等点题句被抢跑吃掉。
|
||||
- 现场：局 A k1_trutho（11s 长句）完整播出；d9_facee「但空着的那一秒，你说满了」**两局两次都完整落地**（局 A 5.28s 足时；局 B 3.69s = 6.24/1.7 足时）。句中/句尾单击一律只给 [20]：z1_sdone t=597.8（句尾前 0.5s）、k1_ask t=1320.7、k1_truth t=1398.0、h1_talke t=1794.8、局 B g9_in t=713.6（auto 链直入抉择的尾句）——**全部无效化，无一例外**。
|
||||
- 没有矫枉过正：链中段句子照常可跳（w1_say t=1559.8、h1_danger t=1784.3、e1_a t=1862.8 三处被我跳掉，均按设计生效）。
|
||||
- **结论：修好。尾句保护 + 中段跳句自由两头都在。**
|
||||
|
||||
### 整体回归
|
||||
- 两局完整打到结局：**e1（rep=12）与 e3（rep=0，新结局）**。账目逐点核对：局 A p1L+1/p2L+2/p3L+2/p4R+1/p5L+1/p6L+2/p7L+2/p8L+1=12；局 B 负线 -1→0，跨档播报只在 3→5 时出现一次（rep_high t=1179.6），低→低不播——**账没多也没少**。
|
||||
- 条件差分照常：局 A heardArgument 命中（w1_say 组 t=1556.7）、bond_zhou=1 时 h1_echo t=1797.5 在 h1_talke 后 0.1s 接入；局 B bond_zhou=0 时 h1_echo **确实缺席**（h1_stope t=708.8 → 直接幕间）——回响不撒谎，两个方向都验了。
|
||||
- 危机拍自动流：danger 长震 [25,40,25,40,25] t=1780.9 → 2.1s 前置 → h1_danger → 自动接 dango/talke，不等单击（我跳了 danger 中段，链条照走完）。
|
||||
- 1.7x 压力下菜单开合（结算态 swipedown 打断结局复播 → 再 swipedown 无缝从头复播）、双击重听、长按 boost、无效手势零静默，全部状态正确。
|
||||
- **零叠音（A 106 段 / B 102 段，vstart 全部晚于上一 vend）、零死锁、非资产 console.error=0。**
|
||||
|
||||
---
|
||||
|
||||
## 局 A 体验叙述（老墨，第一人称）
|
||||
|
||||
解锁那一指下去，先听见「新的夜晚，开始了」和 full 档专属的操作口令，然后台标口播——**开局 t=0.011s 预载清单里 phone_ring 直接 200**，我特地在第一通铃前竖起耳朵：t=29.7 起，"铃——铃——"的双音循环垫真的在响，中间还留着电话铃该有的那口气口。10 秒和 25 秒的两句导播提醒照旧准点，但这次它们是叠在铃声上的提醒，不是替铃声还债。上滑听状态，导播说「今晚的口碑，还没立起来」——新词，账对（零起步就该这么说，R1 那句"有点砸了"冤枉人）。
|
||||
|
||||
第一通老周，我按老墨的慢性子听完就点。播完 z1_ring_in 我故意使坏：链末等待里来了个左滑——error 震，导播说了句「这个手势，直播间里用不上」，然后我把手拿开掐表：25.5 秒，"单击屏幕，继续"来了；再 25 秒，"还在吗"也来了。R1 我闺女踩的那个 34 秒深坑，填平了。
|
||||
|
||||
后面一路走心：老歌、原地陪等、小鹿上车、糖糖我问了她（mid-prompt 右滑照旧缓存生效）、真话、带话、陪聊。老韩的危机拍这回带着 danger 长震和 2 秒前置来的，dango 一气呵成；挂电话前 t=1797.5 老韩那句「头通电话那位大爷……这城市的夜，是接着的」又是 0.1s 不差地接进来。签收 p9 左滑，「师父，三年了」足时播完，点题句「但空着的那一秒，你说满了」这回一个字没少——R1 我三次吃掉别人尾巴，这次轮到我护着它了。e1 收播，导播账本五个人全救，rep=12，分毫不差。结尾我又手贱跳了 e1_a 一半——那是我自己的手，重听双击就还回来了。
|
||||
|
||||
## 局 B 体验叙述（墨墨，第一人称）
|
||||
|
||||
我抢过手机就看见「检测到你的夜班记录」——哥的存档。这次我学乖了听完规则：单击继续。结果它把我扔回哥的 e1 结局链，我不服，首句播到 1.5 秒就戳——**戳不动**。行，是护着哥的句子。播完第二句我就跳了（锁只锁一句，讲武德）。我不想了，狂戳穿越结局链，结算画面点一下，新的夜晚 1.7 倍速开跑。
|
||||
|
||||
这把我全走渣男路线：老周问路牌、转接派出所——「说不清，我真是说不清了」和导播那句指名道姓的「这通砸了」一样不少；小鹿还她自由；糖糖我直接哄睡（结算句照样点我「你没听见她在怕什么」）；陈姐劝回家；老韩进服务区「单子黄了」。没有陪过任何人的夜，老韩没替老周谢我——**回响这东西不播就是不播，账清楚**。p9 我倒是大着胆子左滑了一次：e3「一个人的频道也是频道」——这是 R1 我没见过的结局，憋得跟 e4 不一样，是一种"没救成但说出口了"的体面。尾声五个人全是错过版。
|
||||
|
||||
1.7 倍速里我照旧乱来：k1_b 播放中左滑（error 震，13 秒后"还在吗"准点来）、空点乱戳（只有震）、开菜单又关（结局链从头无缝复播）、双击重听、长按 boost。唯一让我嘀咕的是糖糖开口的前两句（「喂？有人吗？」）声音跟同通话后头的不像一个源——后来查了是那两句没有男声部资产、回退到了中性版，能听，但同一通电话里换了个人捏嗓子。
|
||||
|
||||
---
|
||||
|
||||
## 测量局限与证据口径（先说清楚）
|
||||
|
||||
- 1.7x 女声的 ASR 是重灾区：局 B 145 行听写里 MISMATCH 率明显偏高（最低把 b_rule 残句听成 0.08）。**凡 events 层足时播放、仅转写不像的，一律不算游戏问题**；局 B 的"提前截断"已按 dur/1.7 重审，真实截断仅 3 处且全为玩家 tap。
|
||||
- 局 A 与加测会话按 1.0 速审计：4 处提前截断（p6 复播句/w1_say/h1_danger/e1_a）全部为我的 mid-prompt 缓存滑或链中段跳句，**0 处游戏自残**。
|
||||
- 采信引擎级证据的两类口径同 R1：vend 提前于 dur/speed（TRUNCATED）、应答窗口零事件（SILENCE）。
|
||||
- t= 均为解锁后秒数（events ta），与 ASR 转写时间轴有约 5 秒固定偏移，引用转写只取文本与 similarity。
|
||||
- 仪器侧一次怪癖：局 A t=650-698 有两个 harness 发出的 tap 未抵达页面（curl 响应被脚本丢弃无法溯源），游戏无责（页面零事件零震动，符合"没收到就是没收到"）；此后所有手势均校验响应。
|
||||
|
||||
## 问题清单
|
||||
|
||||
### P0（卡死/死锁/状态混乱/游戏自残截断）
|
||||
|
||||
无。两局加测共 210+ 段语音零叠音、零死锁；TRUNCATED 全清单 7 处（A4/B3）逐条归因为玩家 tap。
|
||||
|
||||
### P1（体验断裂）
|
||||
|
||||
无。R1 唯一的 P1（F1）已修复并双档位复验。
|
||||
|
||||
### P2（不适/设计意图缺席）
|
||||
|
||||
无新增。R1 的 P2-1（铃声）/P2-3（续玩单击陷阱）/P2-4（尾句）已由 F3/F4/F5 修复；R1 P2-2（arg_muffled/brake_screech 及环境音）仍为 SA3 资产缺口（V3，等 HF token），本轮维持已知遗留，不重复计数。
|
||||
|
||||
### P3（建议，6 条）
|
||||
|
||||
1. **rep_low 的 warning 震仍未出现**（R1 P3-3 遗留）：局 A t=99.6-102.0 RING 上滑听 rep_low，只有上滑自身的 [10]，无附加警示震。代码注释称 rep_low 附 warning 震——要么补齐，要么删注释。
|
||||
2. **k1_a/k1_b 缺 male 声部资产，静默回退中性路径**（新发现）：局 A/局 B 糖糖通前两句 fetch `audio/voice/male/k1_a.*` 全 404 → 回退 `audio/voice/k1_a.ogg` 200 播出。功能无损（回退链路工作正常），但同一通电话内可能存在声线跳变，建议补齐渲染或确认音色衔接。
|
||||
3. **一次性的提示句复播晚点**：局 A p2 首次复播 t=740.1 比基准晚 ~7.6s（其后所有复播/看门狗首拍均精确到 ±0.2s：1281.1/1383.0/1550.4 等）。仅一次观测，建议加断言盯一下。
|
||||
4. **ASK 空点单击的确认震手感**（R1 P3-1 遗留）：p9 后空点两连戳各给一记 [20]，像"点上了"实则什么都没发生。
|
||||
5. **harness：POST /gesture 响应未校验时会静默丢手势**（仪器侧）：建议脚本侧始终检查 `{"ok":true}`；长按固定 600ms 仍只能体验 ~0.1s boost（遗留工具局限）。
|
||||
6. **mid-prompt 滑动生效会掐掉选项句尾巴**（设计备注，对应 R1 V1 观察）：局 A p6 复播 6.23/8.64 处被缓存滑终止立即生效——"零插播"契约的一部分，玩家侧无感知损失，但与"选项句听全文"的期望有张力，人设复审时可议。
|
||||
|
||||
## 亮点（让我愿意玩下去的东西）
|
||||
|
||||
- **五项修复全是真修**：F1 的看门狗重挂以手势时刻为基准、两种 guide 档位四组实测分秒不差；F3 连"接听即停早于人声"都做对了；F4 连锁的释放都验得出；F5 护尾不废跳。修得干净，没有修出新坑。
|
||||
- **账目依旧真实得可怕**：rep 12/0 两条线与设计表吻合；跨档播报只在真跨档时出现；h1_echo 有 bond 则来、无 bond 则绝不剧透；尾声全 ok/全 no 两版都亲耳听过。
|
||||
- **e3 是个新收获**：「一个人的频道也是频道」——没救成但把话说出口的夜班，和 e4 的「明晚再说吧」是两种不同的难受。四结局现场已见其三（e1/e3/e4）。
|
||||
- **1.7 倍速下的底盘纹丝不动**：乱序手势 57 连发、菜单打断恢复、双击重听、看门狗 13/30 秒准点，零状态错乱。
|
||||
- 铃声终于响了。这个只有耳朵的世界，的第一声"叮铃"，等了一整个 R1。
|
||||
|
||||
## 一句话裁决
|
||||
|
||||
**愿意**——五项修复逐一现场坐实、两局零死锁打出 e1/e3、账目分毫不差，radio 可以进 R3 全员复评；剩下的都是资产位（SA3 的 13 件声音）和打磨键，我等着把铃声、争吵声和刹车声凑齐的那班夜车。
|
||||
|
||||
---
|
||||
|
||||
### 附：产物与证据文件
|
||||
|
||||
- 局 A：`/home/xsl/SoundGame/web/harness/live/out/radio-20260902103345/`（events.jsonl 106 语音段/65 手势、verdicts.json、transcript.md、mix.webm）
|
||||
- 局 B：`/home/xsl/SoundGame/web/harness/live/out/radio-20260902110919/`（events.jsonl 102 语音段/57 手势）
|
||||
- 加测（F4 释放）：`/home/xsl/SoundGame/web/harness/live/out/radio-20260902112548/`
|
||||
- 引用代码（只读，未改动）：`/home/xsl/SoundGame/web/radio/game.js`（F1 重挂 L364/L383/L411/L458/L581、F3 铃声 L755-783、F4 跳句锁 L1024-1027、F5 尾句判定 L495-511、看门狗常量 L227-234）
|
||||
- 铃声响度证据：局 A mix.webm 30-39s 窗，ffmpeg astats Peak -1.5 dB / RMS -15.4 dB
|
||||
- 离线旁证：`/home/xsl/SoundGame/research/new8/radio/round2/transcript-stress.md`（rep_low 新词 9 处）
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -174,3 +174,8 @@ optimized/tflite/models/tflite/lora_cache/
|
||||
# Rendered audio from the CLI / gradio backends. *.wav above already covers most of it,
|
||||
# but the --file-format flag also writes .flac and .mp3.
|
||||
optimized/*/output/
|
||||
|
||||
# 误伤修正（2026-09-02）：上面的 models/ 规则会忽略源码子包 stable_audio_3/models/
|
||||
# （factory.py/loading_utils.py 的依赖，迁移机器时曾因未入库而丢失，从上游 GitHub 重取）。
|
||||
!stable_audio_3/models/
|
||||
!stable_audio_3/models/**
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,639 @@
|
||||
import torch
|
||||
|
||||
from torch import nn
|
||||
from torch.nn.utils import weight_norm
|
||||
from torchaudio import transforms as T
|
||||
from einops import rearrange
|
||||
|
||||
from ..inference.audio_utils import prepare_audio
|
||||
from .transformer import TransformerBlock
|
||||
|
||||
def WNConv1d(*args, **kwargs):
|
||||
return weight_norm(nn.Conv1d(*args, **kwargs))
|
||||
|
||||
def checkpoint(function, *args, **kwargs):
|
||||
kwargs.setdefault("use_reentrant", False)
|
||||
return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
|
||||
|
||||
|
||||
class Transpose(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
def forward(self, x, **kwargs):
|
||||
return rearrange(x, '... a b -> ... b a')
|
||||
|
||||
def _zero_pad_modulo_sequence(x, size, dim=-2):
|
||||
input_len = x.shape[dim]
|
||||
pad_len = (size - input_len % size) % size
|
||||
if pad_len > 0:
|
||||
pad_shape = list(x.shape)
|
||||
pad_shape[dim] = pad_len
|
||||
x = torch.cat([x, torch.zeros(pad_shape, device=x.device, dtype=x.dtype)], dim=dim)
|
||||
return x
|
||||
|
||||
class TransformerResamplingBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride, sliding_window = None, chunk_size = 128, chunk_midpoint_shift = False, type = 'encoder', transformer_depth = 3, checkpointing = False,
|
||||
conformer = False, layer_scale = False, dim_heads = 128, differential = True, variable_stride = False, feat_scale = False,
|
||||
sinusoidal_blocks = 0, mask_noise = 0, ff_mult = 3, mapping_bias = True, cross_attn = False, dyt = True, conv_mapping = False, freeze_backbone = False, **kwargs):
|
||||
super().__init__()
|
||||
if type not in ['encoder', 'decoder']:
|
||||
raise ValueError(f"Unknown type {type}. Must be 'encoder' or 'decoder'")
|
||||
|
||||
self.checkpointing = checkpointing
|
||||
|
||||
transformer_dim = out_channels if type == 'encoder' else in_channels
|
||||
transformers = []
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.variable_stride = variable_stride
|
||||
self.stride = stride
|
||||
self.mapping = WNConv1d(in_channels, out_channels, 3 if conv_mapping else 1, padding = 'same', bias = mapping_bias) if in_channels != out_channels else nn.Identity()
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_midpoint_shift = chunk_midpoint_shift
|
||||
self.type = type
|
||||
self.mask_noise = mask_noise
|
||||
self.sliding_window_latents = sliding_window
|
||||
|
||||
self.sliding_window_seq = self._get_sliding_window_size(sliding_window, stride)
|
||||
self.input_seg_size, self.output_seg_size, self.sub_chunk_size = self._get_seg_sizes(stride)
|
||||
self.transformer_depth = transformer_depth
|
||||
for i in range(transformer_depth):
|
||||
sinusoidal = True if ((transformer_depth - i) < sinusoidal_blocks) else False
|
||||
transformers.append(TransformerBlock(transformer_dim,
|
||||
dim_heads = dim_heads,
|
||||
causal = False,
|
||||
zero_init_branch_outputs = True if not layer_scale else False,
|
||||
norm_type = 'dyt' if dyt else 'rms_norm',
|
||||
conformer = conformer,
|
||||
layer_scale = layer_scale,
|
||||
add_rope = True,
|
||||
attn_kwargs={'qk_norm': "dyt" if dyt else "rms", "qk_norm_eps": 1e-3, "differential": differential, "feat_scale": feat_scale},
|
||||
ff_kwargs={'mult': ff_mult, 'no_bias': False, "sinusoidal": sinusoidal},
|
||||
norm_kwargs = {'eps': 1e-3},
|
||||
cross_attend = cross_attn))
|
||||
|
||||
self.new_tokens = nn.Parameter(1e-5 * torch.randn(1, self.output_seg_size if not self.variable_stride else 1, out_channels if type == 'encoder' else in_channels))
|
||||
|
||||
self.transformers = nn.ModuleList(transformers)
|
||||
|
||||
if freeze_backbone:
|
||||
for param in self.transformers.parameters():
|
||||
param.requires_grad = False
|
||||
self.new_tokens.requires_grad = False
|
||||
|
||||
def _get_sliding_window_size(self, window, stride, prepend_cond_length = 0):
|
||||
if window is None:
|
||||
return None
|
||||
else:
|
||||
return [(win * (stride + 1 + prepend_cond_length)) for win in window]
|
||||
|
||||
def _get_seg_sizes(self, stride, prepend_cond_length = 0):
|
||||
sub_chunk_size = stride + 1 + prepend_cond_length
|
||||
if self.sliding_window_latents is None:
|
||||
assert (self.chunk_size % stride) == 0, f"Stride must fit evenly into chunk size:{self.chunk_size}"
|
||||
input_seg_size = stride if self.type == 'encoder' else 1
|
||||
output_seg_size = 1 if self.type == 'encoder' else stride
|
||||
return input_seg_size, output_seg_size, sub_chunk_size
|
||||
|
||||
#@torch.compile
|
||||
def forward(self, x, stride = None, return_features = False, override_new_tokens = None, prepend_cond = None, cross_attn_cond = None):
|
||||
batch_size = x.shape[0]
|
||||
input_length = x.shape[-1]
|
||||
if return_features:
|
||||
features = []
|
||||
|
||||
if stride == None:
|
||||
input_seg_size = self.input_seg_size
|
||||
output_seg_size = self.output_seg_size
|
||||
sub_chunk_size = self.sub_chunk_size
|
||||
sliding_window = self.sliding_window_seq
|
||||
else:
|
||||
if not self.variable_stride:
|
||||
print("cannot override stride if variable_stride is not set")
|
||||
prepend_cond_length = prepend_cond.shape[-2] if prepend_cond is not None else 0
|
||||
input_seg_size, output_seg_size, sub_chunk_size = self._get_seg_sizes(stride, prepend_cond_length)
|
||||
sliding_window = self._get_sliding_window_size(self.sliding_window_latents, stride, prepend_cond_length)
|
||||
|
||||
if self.type == 'encoder':
|
||||
# Pad before mapping so silence zeros get projected through the mapping,
|
||||
# rather than inserting raw zeros in the mapped space
|
||||
if self.transformer_depth > 0:
|
||||
if sliding_window is None:
|
||||
pad_modulo = self.chunk_size
|
||||
else:
|
||||
pad_modulo = input_seg_size
|
||||
x = _zero_pad_modulo_sequence(x, pad_modulo, dim=-1)
|
||||
x = self.mapping(x)
|
||||
|
||||
if self.transformer_depth > 0:
|
||||
x = rearrange(x, '... a b -> ... b a')
|
||||
if return_features:
|
||||
features.append(x)
|
||||
if self.type != 'encoder':
|
||||
if sliding_window is None:
|
||||
active_stride = stride if stride is not None else self.stride
|
||||
pad_modulo = self.chunk_size // active_stride
|
||||
x = _zero_pad_modulo_sequence(x, pad_modulo)
|
||||
else:
|
||||
x = _zero_pad_modulo_sequence(x, input_seg_size)
|
||||
x = rearrange(x, 'b (n c) d -> (b n) c d', c = input_seg_size)
|
||||
new_token_seq_dim = -1 if not self.variable_stride else output_seg_size
|
||||
new_tokens = self.new_tokens.expand([x.shape[0],new_token_seq_dim,-1])
|
||||
if override_new_tokens is not None:
|
||||
#print(f"Using override new tokens with shape {override_new_tokens.shape}, new tokens shape {new_tokens.shape}, x shape {x.shape}")
|
||||
override_new_tokens = rearrange(override_new_tokens, 'b (n c) d -> (b n) c d', c = output_seg_size)
|
||||
new_tokens = new_tokens + override_new_tokens
|
||||
elif self.mask_noise > 0:
|
||||
new_tokens = new_tokens + torch.randn_like(new_tokens) * self.mask_noise
|
||||
x = torch.cat([x,new_tokens], dim = -2)
|
||||
if prepend_cond is not None:
|
||||
n = x.shape[0] // batch_size
|
||||
cond_folded = prepend_cond.unsqueeze(1).expand(batch_size, n, prepend_cond.shape[-2], x.shape[-1]).reshape(n * batch_size, prepend_cond.shape[-2], x.shape[-1])
|
||||
x = torch.cat([cond_folded, x], dim=-2)
|
||||
x = rearrange(x, '(b n) c d -> b (n c) d', b=batch_size)#.contiguous()
|
||||
|
||||
# Fold into contiguous chunks if no sliding window
|
||||
if sliding_window is None:
|
||||
prepend_cond_length = prepend_cond.shape[-2] if prepend_cond is not None else 0
|
||||
effective_chunk_size = self.chunk_size + self.chunk_size * (1 + prepend_cond_length) // (stride if stride is not None else self.stride)
|
||||
|
||||
if sliding_window is None and self.chunk_midpoint_shift:
|
||||
split = self.transformer_depth // 2
|
||||
shift = effective_chunk_size // 2
|
||||
|
||||
# First half: standard chunks
|
||||
nc = x.shape[1] // effective_chunk_size
|
||||
x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
|
||||
cross_attn_first = None
|
||||
if cross_attn_cond is not None:
|
||||
cross_attn_first = cross_attn_cond.repeat_interleave(nc, dim=0)
|
||||
for layer in self.transformers[:split]:
|
||||
if self.checkpointing:
|
||||
x = checkpoint(layer, x, context = cross_attn_first, self_attention_flash_sliding_window = None)
|
||||
else:
|
||||
x = layer(x, context = cross_attn_first)
|
||||
if return_features:
|
||||
features.append(rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size))
|
||||
x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
||||
|
||||
# Second half: shifted chunks with sub-chunk repeat padding
|
||||
# shift is always a multiple of sub_chunk_size, so slicing
|
||||
# by shift gives whole sub-chunks in their original order
|
||||
x = torch.cat([x[:, :shift, :], x, x[:, -shift:, :]], dim=1)
|
||||
nc_shifted = x.shape[1] // effective_chunk_size
|
||||
x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
|
||||
cross_attn_second = None
|
||||
if cross_attn_cond is not None:
|
||||
cross_attn_second = cross_attn_cond.repeat_interleave(nc_shifted, dim=0)
|
||||
for layer in self.transformers[split:]:
|
||||
if self.checkpointing:
|
||||
x = checkpoint(layer, x, context = cross_attn_second, self_attention_flash_sliding_window = None)
|
||||
else:
|
||||
x = layer(x, context = cross_attn_second)
|
||||
if return_features:
|
||||
feat = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
||||
features.append(feat[:, shift:-shift, :])
|
||||
x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
||||
x = x[:, shift:-shift, :]
|
||||
else:
|
||||
if sliding_window is None:
|
||||
x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
|
||||
|
||||
for layer in self.transformers:
|
||||
if self.checkpointing:
|
||||
x = checkpoint(layer, x, context = cross_attn_cond, self_attention_flash_sliding_window = sliding_window)
|
||||
else:
|
||||
x = layer(x, context = cross_attn_cond, self_attention_flash_sliding_window = sliding_window)
|
||||
if return_features:
|
||||
features.append(x)
|
||||
|
||||
# Unfold chunks back to original batch
|
||||
if sliding_window is None:
|
||||
x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
||||
|
||||
x = rearrange(x, 'b (n c) d -> (b n) c d', c=sub_chunk_size)
|
||||
x = x[:,-output_seg_size:,:]
|
||||
x = rearrange(x, '(b n) c d -> b d (n c)', b=batch_size)
|
||||
if self.type == 'decoder':
|
||||
x = self.mapping(x)
|
||||
if return_features:
|
||||
return x, features
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class SAMEEncoder(nn.Module):
|
||||
def __init__(self,
|
||||
in_channels=2,
|
||||
channels=128,
|
||||
latent_dim=32,
|
||||
c_mults = [1, 2, 4, 8],
|
||||
strides = [2, 4, 8, 8],
|
||||
transformer_depths = [3,3,3,3],
|
||||
sliding_window = None,
|
||||
checkpointing = False,
|
||||
conformer = False,
|
||||
layer_scale = False,
|
||||
causal = False,
|
||||
differential = True,
|
||||
variable_stride = False,
|
||||
mask_noise = 0.0,
|
||||
conv_mapping = False,
|
||||
freeze_backbone = False,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.strides = strides
|
||||
|
||||
channel_dims = [c * channels for c in c_mults]
|
||||
channel_dims = [in_channels] + channel_dims
|
||||
|
||||
self.depth = len(c_mults)
|
||||
|
||||
layers = []
|
||||
|
||||
for i in range(self.depth):
|
||||
layers += [TransformerResamplingBlock(in_channels=channel_dims[i], out_channels=channel_dims[i+1], stride=strides[i], transformer_depth = transformer_depths[i],
|
||||
sliding_window = sliding_window, checkpointing = checkpointing, conformer = conformer, layer_scale = layer_scale, causal = causal,
|
||||
differential = differential, variable_stride = variable_stride, mask_noise = mask_noise, conv_mapping = conv_mapping,
|
||||
freeze_backbone = freeze_backbone, **kwargs)]
|
||||
|
||||
layers += [Transpose(), nn.Linear(channel_dims[-1], latent_dim), Transpose()]
|
||||
self.layers = nn.ModuleList(layers)
|
||||
|
||||
if freeze_backbone:
|
||||
for param in self.layers[-2].parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, x, override_stride = None, return_features = False, **kwargs):
|
||||
if override_stride != None:
|
||||
assert isinstance(override_stride, list), "override_stride must be a list"
|
||||
assert len(override_stride) == self.depth, "override_stride must be a list containing strides for every layer"
|
||||
for i, layer in enumerate(self.layers):
|
||||
if isinstance(layer, TransformerResamplingBlock):
|
||||
if override_stride != None:
|
||||
stride = override_stride[i]
|
||||
else:
|
||||
stride = None
|
||||
if return_features:
|
||||
x, features = layer(x, stride = stride, return_features = True)
|
||||
else:
|
||||
x = layer(x, stride = stride)
|
||||
else:
|
||||
x = layer(x)
|
||||
if return_features:
|
||||
return x, features
|
||||
else:
|
||||
return x
|
||||
|
||||
class SAMEDecoder(nn.Module):
|
||||
def __init__(self,
|
||||
out_channels=2,
|
||||
channels=128,
|
||||
latent_dim=32,
|
||||
c_mults = [1, 2, 4, 8],
|
||||
strides = [2, 4, 8, 8],
|
||||
transformer_depths = [3,3,3,3],
|
||||
sliding_window = None,
|
||||
checkpointing = False,
|
||||
conformer = False,
|
||||
layer_scale = False,
|
||||
causal = False,
|
||||
differential = True,
|
||||
variable_stride = False,
|
||||
sinusoidal_blocks = [0,0,0,0],
|
||||
mask_noise = 0.0,
|
||||
conv_mapping = False,
|
||||
freeze_backbone = False,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
channel_dims = [c * channels for c in c_mults]
|
||||
channel_dims = [out_channels] + channel_dims
|
||||
|
||||
self.depth = len(c_mults)
|
||||
|
||||
layers = [Transpose(), nn.Linear(latent_dim, channel_dims[-1]), Transpose()]
|
||||
|
||||
for i in range(self.depth, 0, -1):
|
||||
layers += [TransformerResamplingBlock(in_channels=channel_dims[i], out_channels=channel_dims[i-1], stride=strides[i-1], type = 'decoder', transformer_depth = transformer_depths[i-1],
|
||||
sliding_window = sliding_window, checkpointing = checkpointing, conformer = conformer, layer_scale = layer_scale, causal = causal, differential = differential,
|
||||
variable_stride = variable_stride, sinusoidal_blocks = sinusoidal_blocks[i-1], mask_noise = mask_noise, conv_mapping = conv_mapping,
|
||||
freeze_backbone = freeze_backbone, **kwargs)]
|
||||
|
||||
self.layers = nn.ModuleList(layers)
|
||||
|
||||
if freeze_backbone:
|
||||
for param in self.layers[1].parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, x, override_stride = None, **kwargs):
|
||||
if override_stride != None:
|
||||
assert isinstance(override_stride, list), "override_stride must be a list"
|
||||
assert len(override_stride) == self.depth, "override_stride must be a list containing strides for every layer"
|
||||
|
||||
transformer_layer_index = 0
|
||||
for i, layer in enumerate(self.layers):
|
||||
if isinstance(layer, TransformerResamplingBlock):
|
||||
if override_stride != None:
|
||||
stride = override_stride[transformer_layer_index]
|
||||
else:
|
||||
stride = None
|
||||
x = layer(x, stride = stride)
|
||||
transformer_layer_index += 1
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
class AudioAutoencoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
encoder,
|
||||
decoder,
|
||||
latent_dim,
|
||||
downsampling_ratio,
|
||||
sample_rate,
|
||||
io_channels=2,
|
||||
bottleneck: nn.Module = None,
|
||||
pretransform: nn.Module = None,
|
||||
in_channels = None,
|
||||
out_channels = None,
|
||||
soft_clip = False,
|
||||
freeze_pretransform = False
|
||||
):
|
||||
super().__init__()
|
||||
self.downsampling_ratio = downsampling_ratio
|
||||
self.sample_rate = sample_rate
|
||||
|
||||
self.latent_dim = latent_dim
|
||||
self.io_channels = io_channels
|
||||
self.in_channels = io_channels
|
||||
self.out_channels = io_channels
|
||||
|
||||
self.min_length = self.downsampling_ratio
|
||||
|
||||
if in_channels is not None:
|
||||
self.in_channels = in_channels
|
||||
|
||||
if out_channels is not None:
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.bottleneck = bottleneck
|
||||
|
||||
self.encoder = encoder
|
||||
|
||||
self.decoder = decoder
|
||||
|
||||
self.pretransform = pretransform
|
||||
|
||||
self.freeze_pretransform = freeze_pretransform
|
||||
if self.pretransform is not None:
|
||||
if self.freeze_pretransform:
|
||||
for p in self.pretransform.parameters():
|
||||
p.requires_grad = False
|
||||
else:
|
||||
for p in self.pretransform.parameters():
|
||||
p.requires_grad = True
|
||||
|
||||
self.soft_clip = soft_clip
|
||||
|
||||
self.is_discrete = False
|
||||
|
||||
def encode(self, audio, return_info=False, skip_pretransform=False, iterate_batch=False, return_pretransform = False, **kwargs):
|
||||
|
||||
info = {}
|
||||
if self.pretransform is not None and not skip_pretransform:
|
||||
if self.pretransform.enable_grad:
|
||||
if iterate_batch:
|
||||
audios = []
|
||||
for i in range(audio.shape[0]):
|
||||
audios.append(self.pretransform.encode(audio[i:i+1]))
|
||||
audio = torch.cat(audios, dim=0)
|
||||
else:
|
||||
audio = self.pretransform.encode(audio)
|
||||
else:
|
||||
with torch.no_grad():
|
||||
if iterate_batch:
|
||||
audios = []
|
||||
for i in range(audio.shape[0]):
|
||||
audios.append(self.pretransform.encode(audio[i:i+1]))
|
||||
audio = torch.cat(audios, dim=0)
|
||||
else:
|
||||
audio = self.pretransform.encode(audio)
|
||||
if self.encoder is not None:
|
||||
if iterate_batch:
|
||||
latents = []
|
||||
for i in range(audio.shape[0]):
|
||||
latents.append(self.encoder(audio[i:i+1], **kwargs))
|
||||
latents = torch.cat(latents, dim=0)
|
||||
else:
|
||||
latents = self.encoder(audio, **kwargs)
|
||||
else:
|
||||
latents = audio
|
||||
if self.bottleneck is not None:
|
||||
# TODO: Add iterate batch logic, needs to merge the info dicts
|
||||
latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs)
|
||||
|
||||
info.update(bottleneck_info)
|
||||
|
||||
if return_info and return_pretransform:
|
||||
return latents, info, audio
|
||||
elif return_info:
|
||||
return latents, info
|
||||
elif return_pretransform:
|
||||
return latents, audio
|
||||
else:
|
||||
return latents
|
||||
|
||||
def decode(self, latents, iterate_batch=False, return_loss = False, **kwargs):
|
||||
if self.bottleneck is not None:
|
||||
if iterate_batch:
|
||||
decoded = []
|
||||
for i in range(latents.shape[0]):
|
||||
decoded.append(self.bottleneck.decode(latents[i:i+1]))
|
||||
latents = torch.cat(decoded, dim=0)
|
||||
else:
|
||||
latents = self.bottleneck.decode(latents)
|
||||
if iterate_batch:
|
||||
decoded = []
|
||||
for i in range(latents.shape[0]):
|
||||
decoded.append(self.decoder(latents[i:i+1], **kwargs))
|
||||
decoded = torch.cat(decoded, dim=0)
|
||||
else:
|
||||
if return_loss:
|
||||
decoded, loss = self.decoder(latents, **kwargs)
|
||||
else:
|
||||
decoded = self.decoder(latents, **kwargs)
|
||||
if self.pretransform is not None:
|
||||
if self.pretransform.enable_grad:
|
||||
if iterate_batch:
|
||||
decodeds = []
|
||||
for i in range(decoded.shape[0]):
|
||||
decodeds.append(self.pretransform.decode(decoded[i:i+1]))
|
||||
decoded = torch.cat(decodeds, dim=0)
|
||||
else:
|
||||
decoded = self.pretransform.decode(decoded)
|
||||
else:
|
||||
with torch.no_grad():
|
||||
if iterate_batch:
|
||||
decodeds = []
|
||||
for i in range(latents.shape[0]):
|
||||
decodeds.append(self.pretransform.decode(decoded[i:i+1]))
|
||||
decoded = torch.cat(decodeds, dim=0)
|
||||
else:
|
||||
decoded = self.pretransform.decode(decoded)
|
||||
|
||||
if self.soft_clip:
|
||||
decoded = torch.tanh(decoded)
|
||||
if return_loss:
|
||||
return decoded, loss
|
||||
else:
|
||||
return decoded
|
||||
|
||||
|
||||
def preprocess_audio_for_encoder(self, audio, in_sr):
|
||||
'''
|
||||
Preprocess single audio tensor (Channels x Length) to be compatible with the encoder.
|
||||
If the model is mono, stereo audio will be converted to mono.
|
||||
Audio will be silence-padded to be a multiple of the model's downsampling ratio.
|
||||
Audio will be resampled to the model's sample rate.
|
||||
The output will have batch size 1 and be shape (1 x Channels x Length)
|
||||
'''
|
||||
return self.preprocess_audio_list_for_encoder([audio], [in_sr])
|
||||
|
||||
def preprocess_audio_list_for_encoder(self, audio_list, in_sr_list):
|
||||
'''
|
||||
Preprocess a [list] of audio (Channels x Length) into a batch tensor to be compatable with the encoder.
|
||||
The audio in that list can be of different lengths and channels.
|
||||
in_sr can be an integer or list. If it's an integer it will be assumed it is the input sample_rate for every audio.
|
||||
All audio will be resampled to the model's sample rate.
|
||||
Audio will be silence-padded to the longest length, and further padded to be a multiple of the model's downsampling ratio.
|
||||
If the model is mono, all audio will be converted to mono.
|
||||
The output will be a tensor of shape (Batch x Channels x Length)
|
||||
'''
|
||||
batch_size = len(audio_list)
|
||||
if isinstance(in_sr_list, int):
|
||||
in_sr_list = [in_sr_list]*batch_size
|
||||
assert len(in_sr_list) == batch_size, "list of sample rates must be the same length of audio_list"
|
||||
new_audio = []
|
||||
max_length = 0
|
||||
# resample & find the max length
|
||||
for i in range(batch_size):
|
||||
audio = audio_list[i]
|
||||
in_sr = in_sr_list[i]
|
||||
if len(audio.shape) == 3 and audio.shape[0] == 1:
|
||||
# batchsize 1 was given by accident. Just squeeze it.
|
||||
audio = audio.squeeze(0)
|
||||
elif len(audio.shape) == 1:
|
||||
# Mono signal, channel dimension is missing, unsqueeze it in
|
||||
audio = audio.unsqueeze(0)
|
||||
assert len(audio.shape)==2, "Audio should be shape (Channels x Length) with no batch dimension"
|
||||
# Resample audio
|
||||
if in_sr != self.sample_rate:
|
||||
resample_tf = T.Resample(in_sr, self.sample_rate).to(audio.device)
|
||||
audio = resample_tf(audio)
|
||||
new_audio.append(audio)
|
||||
if audio.shape[-1] > max_length:
|
||||
max_length = audio.shape[-1]
|
||||
# Pad every audio to the same length, multiple of model's downsampling ratio
|
||||
padded_audio_length = max_length + (self.min_length - (max_length % self.min_length)) % self.min_length
|
||||
for i in range(batch_size):
|
||||
# Pad it & if necessary, mixdown/duplicate stereo/mono channels to support model
|
||||
new_audio[i] = prepare_audio(new_audio[i], in_sr=in_sr, target_sr=in_sr, target_length=padded_audio_length,
|
||||
target_channels=self.in_channels, device=new_audio[i].device).squeeze(0)
|
||||
# convert to tensor
|
||||
return torch.stack(new_audio)
|
||||
|
||||
def encode_audio(self, audio, chunked=False, overlap=32, chunk_size=128, **kwargs):
|
||||
'''
|
||||
Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder.
|
||||
If chunked is True, split the audio into chunks of a given maximum size chunk_size, with given overlap.
|
||||
Overlap and chunk_size params are both measured in number of latents (not audio samples)
|
||||
# and therefore you likely could use the same values with decode_audio.
|
||||
A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size.
|
||||
Every autoencoder will have a different receptive field size, and thus ideal overlap.
|
||||
You can determine it empirically by diffing unchunked vs chunked output and looking at maximum diff.
|
||||
The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
|
||||
Smaller chunk_size uses less memory, but more compute.
|
||||
The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
|
||||
For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
|
||||
'''
|
||||
samples_per_latent = int(self.downsampling_ratio)
|
||||
if not chunked or audio.shape[-1] < chunk_size * samples_per_latent:
|
||||
return self.encode(audio, **kwargs)
|
||||
|
||||
# chunk_size/overlap are in latent units; scale to samples for slicing.
|
||||
chunk_size_samples = chunk_size * samples_per_latent
|
||||
hop_samples = (chunk_size - overlap) * samples_per_latent
|
||||
total_samples = audio.shape[-1]
|
||||
|
||||
# Anchor the final chunk to the signal end (may overlap more than `overlap`).
|
||||
chunk_starts = list(range(0, total_samples - chunk_size_samples + 1, hop_samples))
|
||||
if chunk_starts[-1] != total_samples - chunk_size_samples:
|
||||
chunk_starts.append(total_samples - chunk_size_samples)
|
||||
|
||||
encoded_chunks = [self.encode(audio[..., s:s + chunk_size_samples]) for s in chunk_starts]
|
||||
|
||||
total_latents = total_samples // samples_per_latent
|
||||
half_overlap_latents = overlap // 2
|
||||
output = audio.new_zeros(*encoded_chunks[0].shape[:-1], total_latents)
|
||||
num_chunks = len(chunk_starts)
|
||||
|
||||
# Trim half the overlap off inner edges (edges facing a neighbour chunk).
|
||||
for i, (start_sample, chunk) in enumerate(zip(chunk_starts, encoded_chunks)):
|
||||
is_first = i == 0
|
||||
is_last = i == num_chunks - 1
|
||||
out_start = (total_latents - chunk_size) if is_last else (start_sample // samples_per_latent)
|
||||
left = 0 if is_first else half_overlap_latents
|
||||
right = chunk_size if is_last else chunk_size - half_overlap_latents
|
||||
output[..., out_start + left : out_start + right] = chunk[..., left:right]
|
||||
|
||||
return output
|
||||
|
||||
def decode_audio(self, latents, chunked=False, overlap=32, chunk_size=128, **kwargs):
|
||||
'''
|
||||
Decode latents to audio.
|
||||
If chunked is True, split the latents into chunks of a given maximum size chunk_size, with given overlap, both of which are measured in number of latents.
|
||||
A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size.
|
||||
Every autoencoder will have a different receptive field size, and thus ideal overlap.
|
||||
You can determine it empirically by diffing unchunked vs chunked audio and looking at maximum diff.
|
||||
The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
|
||||
Smaller chunk_size uses less memory, but more compute.
|
||||
The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
|
||||
For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
|
||||
'''
|
||||
if not chunked or latents.shape[-1] < chunk_size:
|
||||
return self.decode(latents, **kwargs)
|
||||
|
||||
# chunk_size/overlap are in latent units — same as `latents`, no scaling.
|
||||
samples_per_latent = int(self.downsampling_ratio)
|
||||
hop_latents = chunk_size - overlap
|
||||
total_latents = latents.shape[-1]
|
||||
|
||||
# Anchor the final chunk to the signal end (may overlap more than `overlap`).
|
||||
chunk_starts = list(range(0, total_latents - chunk_size + 1, hop_latents))
|
||||
if chunk_starts[-1] != total_latents - chunk_size:
|
||||
chunk_starts.append(total_latents - chunk_size)
|
||||
|
||||
decoded_chunks = [self.decode(latents[..., s:s + chunk_size]) for s in chunk_starts]
|
||||
|
||||
total_samples = total_latents * samples_per_latent
|
||||
chunk_size_samples = chunk_size * samples_per_latent
|
||||
half_overlap_samples = (overlap // 2) * samples_per_latent
|
||||
output = latents.new_zeros(*decoded_chunks[0].shape[:-1], total_samples)
|
||||
num_chunks = len(chunk_starts)
|
||||
|
||||
# Trim half the overlap off inner edges (edges facing a neighbour chunk).
|
||||
for i, (start_latent, chunk) in enumerate(zip(chunk_starts, decoded_chunks)):
|
||||
is_first = i == 0
|
||||
is_last = i == num_chunks - 1
|
||||
out_start = (total_samples - chunk_size_samples) if is_last else (start_latent * samples_per_latent)
|
||||
left = 0 if is_first else half_overlap_samples
|
||||
right = chunk_size_samples if is_last else chunk_size_samples - half_overlap_samples
|
||||
output[..., out_start + left : out_start + right] = chunk[..., left:right]
|
||||
|
||||
return output
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn.utils import weight_norm
|
||||
|
||||
|
||||
def get_activation(activation, channels=None) -> nn.Module:
|
||||
if activation == "elu":
|
||||
return nn.ELU()
|
||||
elif activation == "none":
|
||||
return nn.Identity()
|
||||
else:
|
||||
raise ValueError(f"Unknown activation {activation}")
|
||||
|
||||
def WNConv1d(*args, **kwargs):
|
||||
return weight_norm(nn.Conv1d(*args, **kwargs))
|
||||
|
||||
class ResidualUnit(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, dilation, depthwise=False, bias=True):
|
||||
super().__init__()
|
||||
|
||||
self.dilation = dilation
|
||||
|
||||
padding = (dilation * (7-1)) // 2
|
||||
|
||||
self.layers = nn.Sequential(
|
||||
get_activation("elu"),
|
||||
WNConv1d(in_channels=in_channels, out_channels=out_channels,
|
||||
kernel_size=7, dilation=dilation, padding=padding, groups=1 if not depthwise else out_channels, bias=bias),
|
||||
get_activation("elu"),
|
||||
WNConv1d(in_channels=out_channels, out_channels=out_channels,
|
||||
kernel_size=1, bias=bias)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
res = x
|
||||
x = self.layers(x)
|
||||
return x + res
|
||||
|
||||
class FourierFeatures(nn.Module):
|
||||
def __init__(self, in_features, out_features, std=16.):
|
||||
super().__init__()
|
||||
assert out_features % 2 == 0
|
||||
self.register_buffer('weight', torch.randn([out_features // 2, in_features]) * std)
|
||||
|
||||
def forward(self, input):
|
||||
f = 2 * math.pi * input @ self.weight.T
|
||||
return torch.cat([f.cos(), f.sin()], dim=-1)
|
||||
|
||||
class ExpoFourierFeatures(nn.Module):
|
||||
def __init__(self, dim, min_freq=0.5, max_freq=10000.0):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.min_freq = min_freq
|
||||
self.max_freq = max_freq
|
||||
|
||||
@torch.amp.autocast("cuda",enabled=False)
|
||||
def forward(self, t):
|
||||
"""
|
||||
t: [B] tensor.
|
||||
"""
|
||||
in_dtype = t.dtype
|
||||
t = t.float()
|
||||
|
||||
if t.dim() == 1:
|
||||
t = t.unsqueeze(-1)
|
||||
|
||||
half_dim = self.dim // 2
|
||||
|
||||
# Calculate frequencies (safely in FP32)
|
||||
ramp = torch.linspace(0, 1, half_dim, device=t.device, dtype=torch.float32)
|
||||
log_min = math.log(self.min_freq)
|
||||
log_max = math.log(self.max_freq)
|
||||
|
||||
freqs = torch.exp(ramp * (log_max - log_min) + log_min)
|
||||
|
||||
# Calculate arguments (safely in FP32)
|
||||
args = t * freqs * 2 * math.pi
|
||||
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
|
||||
return embedding.to(in_dtype)
|
||||
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
class SoftNormBottleneck(nn.Module):
|
||||
def __init__(self, dim = 32, noise_augment_dim=0, noise_regularize = False, auto_scale = False, freeze = False):
|
||||
super().__init__()
|
||||
|
||||
self.noise_augment_dim = noise_augment_dim
|
||||
self.scaling_factor = nn.Parameter(torch.ones(1,dim,1))
|
||||
self.bias = nn.Parameter(torch.zeros(1,dim,1))
|
||||
self.noise_scaling_factor = nn.Parameter(torch.ones(1,noise_augment_dim,1))
|
||||
self.noise_regularize = noise_regularize
|
||||
self.freeze = freeze
|
||||
if self.freeze:
|
||||
self.scaling_factor.requires_grad = False
|
||||
self.bias.requires_grad = False
|
||||
self.noise_scaling_factor.requires_grad = False
|
||||
if auto_scale:
|
||||
running_std = torch.ones(1)
|
||||
self.register_parameter("running_std", nn.Parameter(running_std, requires_grad=False))
|
||||
|
||||
def encode(self, x, return_info=False, **kwargs):
|
||||
info = {}
|
||||
|
||||
x = x * self.scaling_factor + self.bias
|
||||
|
||||
if self.training and hasattr(self,"running_std") and not self.freeze:
|
||||
# Update running std
|
||||
self.running_std.data = (self.running_std.data * 0.999 + x.std().detach() * 0.001).clamp(min = 1e-4)
|
||||
|
||||
if hasattr(self, "running_std"):
|
||||
x = x / self.running_std
|
||||
|
||||
if self.training and return_info:
|
||||
var = (x.std(dim=-1) ** 2).clip(min = 1e-4)
|
||||
logvar = torch.log(var)
|
||||
mean = x.mean(dim=-1)
|
||||
loss = (mean * mean + var - logvar - 1).mean()
|
||||
var = (x.std(dim=-2) ** 2).clip(min = 1e-4)
|
||||
logvar = torch.log(var)
|
||||
mean = x.mean(dim=-2)
|
||||
loss = loss + 0.4 * (mean * mean + var - logvar - 1).mean()
|
||||
info["softnorm_loss"] = loss
|
||||
|
||||
if return_info:
|
||||
return x, info
|
||||
|
||||
return x
|
||||
|
||||
def decode(self, x, **kwargs):
|
||||
if hasattr(self, "running_std"):
|
||||
x = x * self.running_std
|
||||
|
||||
if self.noise_regularize:
|
||||
if hasattr(self, "running_std"):
|
||||
scaling = self.running_std
|
||||
else:
|
||||
scaling = x.std(dim = -1).unsqueeze(-1)
|
||||
if self.training:
|
||||
scale = 5e-2
|
||||
else:
|
||||
scale = 1e-3
|
||||
noise = torch.randn_like(x) * scaling * scale
|
||||
x = x + noise
|
||||
|
||||
if self.noise_augment_dim > 0:
|
||||
noise = self.noise_scaling_factor * torch.randn(x.shape[0], self.noise_augment_dim,
|
||||
x.shape[-1]).type_as(x)
|
||||
x = torch.cat([x, noise], dim=1)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,318 @@
|
||||
#Heavily influenced by https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/modules/conditioners.py
|
||||
|
||||
import torch
|
||||
import logging, warnings
|
||||
import typing as tp
|
||||
from enum import Enum
|
||||
from math import pi
|
||||
from typing import List, Union
|
||||
from torch import Tensor, nn
|
||||
from einops import rearrange
|
||||
from .blocks import ExpoFourierFeatures
|
||||
from .utils import enable_torch_compile
|
||||
import os
|
||||
|
||||
class PaddingMode(str, Enum):
|
||||
"""Enum for handling padding in text conditioner embeddings."""
|
||||
NONE = "none" # No padding handling (raw embeddings with pad token)
|
||||
ZERO = "zero" # Zero out padding positions (default)
|
||||
LEARNED = "learned" # Use learned padding embedding
|
||||
|
||||
class Conditioner(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
output_dim: int,
|
||||
project_out: bool = False,
|
||||
padding_mode: str = "zero"
|
||||
):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.output_dim = output_dim
|
||||
self.padding_mode = padding_mode
|
||||
self.proj_out = nn.Linear(dim, output_dim) if (dim != output_dim or project_out) else nn.Identity()
|
||||
|
||||
# Learned padding embedding (only created if needed)
|
||||
if padding_mode == "learned" or padding_mode == PaddingMode.LEARNED:
|
||||
self.padding_embedding = nn.Parameter(torch.randn(output_dim) * 0.02)
|
||||
|
||||
def apply_padding(self, embeddings: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Apply padding handling based on padding_mode.
|
||||
|
||||
Args:
|
||||
embeddings: [batch, seq_len, dim] - the embeddings to process
|
||||
attention_mask: [batch, seq_len] bool/int, True/1 = valid token
|
||||
|
||||
Returns:
|
||||
embeddings with padding handled according to mode
|
||||
"""
|
||||
mode = self.padding_mode
|
||||
if isinstance(mode, str):
|
||||
mode = PaddingMode(mode)
|
||||
|
||||
if mode == PaddingMode.NONE:
|
||||
return embeddings
|
||||
elif mode == PaddingMode.ZERO:
|
||||
return embeddings * attention_mask.unsqueeze(-1).float()
|
||||
elif mode == PaddingMode.LEARNED:
|
||||
mask_expanded = attention_mask.unsqueeze(-1).bool()
|
||||
return torch.where(
|
||||
mask_expanded,
|
||||
embeddings,
|
||||
self.padding_embedding.unsqueeze(0).unsqueeze(0).expand_as(embeddings)
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown padding mode: {mode}")
|
||||
|
||||
def forward(self, x: tp.Any) -> tp.Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
class LearnedPositionalEmbedding(nn.Module):
|
||||
def __init__(self, dim: int, std: float = 16.0):
|
||||
super().__init__()
|
||||
assert (dim % 2) == 0
|
||||
half_dim = dim // 2
|
||||
self.weights = nn.Parameter(torch.randn(half_dim) * std)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = rearrange(x, "b -> b 1")
|
||||
freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * pi
|
||||
fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
|
||||
fouriered = torch.cat((x, fouriered), dim=-1)
|
||||
return fouriered
|
||||
|
||||
|
||||
def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module:
|
||||
return nn.Sequential(
|
||||
LearnedPositionalEmbedding(dim),
|
||||
nn.Linear(in_features=dim + 1, out_features=out_features),
|
||||
)
|
||||
|
||||
|
||||
class NumberEmbedder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
features: int,
|
||||
dim: int = 256,
|
||||
fourier_features_type: tp.Literal["learned", "expo"] = "learned"
|
||||
):
|
||||
super().__init__()
|
||||
self.features = features
|
||||
if fourier_features_type == "expo":
|
||||
self.embedding = nn.Sequential(ExpoFourierFeatures(dim=dim), nn.Linear(in_features=dim, out_features=features))
|
||||
else:
|
||||
self.embedding = TimePositionalEmbedding(dim=dim, out_features=features)
|
||||
|
||||
def forward(self, x: Union[List[float], Tensor]) -> Tensor:
|
||||
if not torch.is_tensor(x):
|
||||
device = next(self.embedding.parameters()).device
|
||||
x = torch.tensor(x, device=device)
|
||||
assert isinstance(x, Tensor)
|
||||
shape = x.shape
|
||||
x = rearrange(x, "... -> (...)")
|
||||
embedding = self.embedding(x)
|
||||
x = embedding.view(*shape, self.features)
|
||||
return x # type: ignore
|
||||
|
||||
|
||||
class NumberConditioner(Conditioner):
|
||||
'''
|
||||
Conditioner that takes a list of floats, normalizes them for a given range, and returns a list of embeddings
|
||||
'''
|
||||
def __init__(self,
|
||||
output_dim: int,
|
||||
min_val: float=0,
|
||||
max_val: float=1,
|
||||
fourier_features_type : tp.Literal["learned", "expo"] = "learned"
|
||||
):
|
||||
super().__init__(output_dim, output_dim)
|
||||
|
||||
self.min_val = min_val
|
||||
self.max_val = max_val
|
||||
|
||||
self.embedder = NumberEmbedder(features=output_dim, fourier_features_type=fourier_features_type)
|
||||
|
||||
def forward(self, floats: tp.List[float], device=None) -> tp.Any:
|
||||
self.embedder.to(device)
|
||||
# Cast the inputs to floats
|
||||
floats = [float(x) for x in floats]
|
||||
|
||||
floats = torch.tensor(floats).to(device)
|
||||
|
||||
floats = floats.clamp(self.min_val, self.max_val)
|
||||
|
||||
normalized_floats = (floats - self.min_val) / (self.max_val - self.min_val)
|
||||
|
||||
# Cast floats to same type as embedder
|
||||
embedder_dtype = next(self.embedder.parameters()).dtype
|
||||
normalized_floats = normalized_floats.to(embedder_dtype)
|
||||
|
||||
float_embeds = self.embedder(normalized_floats).unsqueeze(1)
|
||||
|
||||
return [float_embeds, torch.ones(float_embeds.shape[0], 1).to(device)]
|
||||
|
||||
class T5GemmaConditioner(Conditioner):
|
||||
|
||||
T5GEMMA_MODELS = ["google/t5gemma-b-b-ul2"]
|
||||
|
||||
T5GEMMA_MODEL_DIMS = {
|
||||
"google/t5gemma-b-b-ul2": 768,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
output_dim: int,
|
||||
model_name: str = "google/t5gemma-b-b-ul2",
|
||||
max_length: str = 128,
|
||||
enable_grad: bool = False,
|
||||
project_out: bool = False,
|
||||
padding_mode: str = "zero",
|
||||
model_path: str = None,
|
||||
repo_id: str = None,
|
||||
subfolder: str = None,
|
||||
):
|
||||
assert model_name in self.T5GEMMA_MODELS, f"Unknown T5 model name: {model_name}"
|
||||
super().__init__(self.T5GEMMA_MODEL_DIMS[model_name], output_dim, project_out=project_out, padding_mode=padding_mode)
|
||||
|
||||
load_from = model_path or repo_id or model_name
|
||||
|
||||
self.max_length = max_length
|
||||
self.enable_grad = enable_grad
|
||||
|
||||
# Set environment variables to disable progress bars BEFORE importing transformers
|
||||
# This is the most reliable way to suppress HuggingFace progress bars
|
||||
prev_hf_hub = os.environ.get("HF_HUB_DISABLE_PROGRESS_BARS")
|
||||
prev_transformers = os.environ.get("TRANSFORMERS_VERBOSITY")
|
||||
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
|
||||
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
|
||||
|
||||
# Suppress logging from transformers
|
||||
previous_level = logging.root.manager.disable
|
||||
logging.disable(logging.ERROR)
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore")
|
||||
try:
|
||||
from transformers import T5GemmaEncoderModel, AutoTokenizer, AutoConfig
|
||||
logging.info(f"Loading T5Gemma tokenizer and model from: {load_from}")
|
||||
hf_kwargs = {"subfolder": subfolder} if subfolder else {}
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(load_from, **hf_kwargs)
|
||||
config = AutoConfig.from_pretrained(load_from, **hf_kwargs)
|
||||
config.is_encoder_decoder = False
|
||||
model = T5GemmaEncoderModel.from_pretrained(load_from, config=config, **hf_kwargs).train(enable_grad).requires_grad_(enable_grad)
|
||||
|
||||
finally:
|
||||
logging.disable(previous_level)
|
||||
# Restore environment variables
|
||||
if prev_hf_hub is None:
|
||||
os.environ.pop("HF_HUB_DISABLE_PROGRESS_BARS", None)
|
||||
else:
|
||||
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = prev_hf_hub
|
||||
if prev_transformers is None:
|
||||
os.environ.pop("TRANSFORMERS_VERBOSITY", None)
|
||||
else:
|
||||
os.environ["TRANSFORMERS_VERBOSITY"] = prev_transformers
|
||||
|
||||
# Compile the model to reduce CPU-GPU kernel launch overhead,
|
||||
# which is sensitive to CPU contention from DataLoader workers
|
||||
if enable_torch_compile:
|
||||
model = torch.compile(model)
|
||||
|
||||
if self.enable_grad:
|
||||
self.model = model
|
||||
else:
|
||||
self.__dict__["model"] = model
|
||||
|
||||
self._device_initialized = False
|
||||
|
||||
def forward(self, inputs: tp.Union[tp.List[str], tp.List[tp.Dict[str, torch.Tensor]]], device: tp.Union[torch.device, str]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
|
||||
|
||||
# Only move to device once (avoid overhead on every forward call)
|
||||
if not self._device_initialized:
|
||||
self.model.to(device)
|
||||
self.proj_out.to(device)
|
||||
self.model.eval()
|
||||
self._device_initialized = True
|
||||
|
||||
# Handle pre-tokenized inputs (dicts with input_ids/attention_mask from DataLoader workers)
|
||||
# or raw strings (from demo generation / inference)
|
||||
if isinstance(inputs[0], dict):
|
||||
input_ids = torch.stack([x["input_ids"] for x in inputs]).to(device, non_blocking=True)
|
||||
attention_mask = torch.stack([x["attention_mask"] for x in inputs]).to(device, non_blocking=True).to(torch.bool)
|
||||
else:
|
||||
encoded = self.tokenizer(
|
||||
inputs,
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
padding="max_length",
|
||||
return_tensors="pt",
|
||||
)
|
||||
input_ids = encoded["input_ids"].to(device, non_blocking=True)
|
||||
attention_mask = encoded["attention_mask"].to(device, non_blocking=True).to(torch.bool)
|
||||
|
||||
with torch.no_grad():
|
||||
embeddings = self.model(
|
||||
input_ids=input_ids, attention_mask=attention_mask
|
||||
)["last_hidden_state"]
|
||||
|
||||
# Cast embeddings to same type as proj_out, unless proj_out is Identity
|
||||
if not isinstance(self.proj_out, nn.Identity):
|
||||
proj_out_dtype = next(self.proj_out.parameters()).dtype
|
||||
embeddings = embeddings.to(proj_out_dtype)
|
||||
|
||||
embeddings = self.proj_out(embeddings)
|
||||
embeddings = self.apply_padding(embeddings, attention_mask)
|
||||
|
||||
return embeddings, attention_mask
|
||||
|
||||
|
||||
class MultiConditioner(nn.Module):
|
||||
"""
|
||||
A module that applies multiple conditioners to an input dictionary based on the keys
|
||||
|
||||
Args:
|
||||
conditioners: a dictionary of conditioners with keys corresponding to the keys of the conditioning input dictionary (e.g. "prompt")
|
||||
default_keys: a dictionary of default keys to use if the key is not in the input dictionary (e.g. {"prompt_t5": "prompt"})
|
||||
"""
|
||||
def __init__(self, conditioners: tp.Dict[str, Conditioner], default_keys: tp.Dict[str, str] = {}, pre_encoded_keys: tp.List[str] = []):
|
||||
super().__init__()
|
||||
|
||||
self.conditioners = nn.ModuleDict(conditioners)
|
||||
self.default_keys = default_keys
|
||||
self.pre_encoded_keys = pre_encoded_keys
|
||||
|
||||
def forward(self, batch_metadata: tp.List[tp.Dict[str, tp.Any]], device: tp.Union[torch.device, str]) -> tp.Dict[str, tp.Any]:
|
||||
output = {}
|
||||
|
||||
for key, conditioner in self.conditioners.items():
|
||||
condition_key = key
|
||||
|
||||
conditioner_inputs = []
|
||||
|
||||
for x in batch_metadata:
|
||||
|
||||
if condition_key not in x:
|
||||
if condition_key in self.default_keys:
|
||||
condition_key = self.default_keys[condition_key]
|
||||
else:
|
||||
raise ValueError(f"Conditioner key {condition_key} not found in batch metadata")
|
||||
|
||||
#Unwrap the condition info if it's a single-element list or tuple, this is to support collation functions that wrap everything in a list
|
||||
if isinstance(x[condition_key], list) or isinstance(x[condition_key], tuple) and len(x[condition_key]) == 1:
|
||||
conditioner_input = x[condition_key][0]
|
||||
|
||||
else:
|
||||
conditioner_input = x[condition_key]
|
||||
|
||||
conditioner_inputs.append(conditioner_input)
|
||||
|
||||
if key in self.pre_encoded_keys:
|
||||
output[key] = [torch.stack(conditioner_inputs, dim=0).to(device), None]
|
||||
else:
|
||||
output[key] = conditioner(conditioner_inputs, device)
|
||||
|
||||
return output
|
||||
|
||||
@@ -0,0 +1,253 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import typing as tp
|
||||
|
||||
from .conditioners import MultiConditioner
|
||||
from .dit import DiffusionTransformer
|
||||
from stable_audio_3.inference.distribution_shift import FluxDistributionShift, DistributionShift, LogSNRShift, IdentityDistributionShift
|
||||
|
||||
from time import time
|
||||
|
||||
class Profiler:
|
||||
|
||||
def __init__(self):
|
||||
self.ticks = [[time(), None]]
|
||||
|
||||
def tick(self, msg):
|
||||
self.ticks.append([time(), msg])
|
||||
|
||||
def __repr__(self):
|
||||
rep = 80 * "=" + "\n"
|
||||
for i in range(1, len(self.ticks)):
|
||||
msg = self.ticks[i][1]
|
||||
ellapsed = self.ticks[i][0] - self.ticks[i - 1][0]
|
||||
rep += msg + f": {ellapsed*1000:.2f}ms\n"
|
||||
rep += 80 * "=" + "\n\n\n"
|
||||
return rep
|
||||
|
||||
class ConditionedDiffusionModelWrapper(nn.Module):
|
||||
"""
|
||||
A diffusion model that takes in conditioning
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
conditioner: MultiConditioner,
|
||||
io_channels,
|
||||
sample_rate,
|
||||
min_input_length: int,
|
||||
diffusion_objective: tp.Literal["v", "rectified_flow", "rf_denoiser"] = "v",
|
||||
distribution_shift_options = None,
|
||||
sampling_distribution_shift_options = None,
|
||||
mask_padding_attention: bool = False,
|
||||
use_effective_length_for_schedule: bool = False,
|
||||
pretransform: tp.Optional[nn.Module] = None,
|
||||
cross_attn_cond_ids: tp.List[str] = [],
|
||||
global_cond_ids: tp.List[str] = [],
|
||||
input_concat_ids: tp.List[str] = [],
|
||||
local_add_cond_ids: tp.List[str] = [],
|
||||
modular_local_cond_ids: tp.List[str] = [],
|
||||
prepend_cond_ids: tp.List[str] = [],
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.model = model
|
||||
self.conditioner = conditioner
|
||||
self.io_channels = io_channels
|
||||
self.sample_rate = sample_rate
|
||||
self.diffusion_objective = diffusion_objective
|
||||
self.pretransform = pretransform
|
||||
self.cross_attn_cond_ids = cross_attn_cond_ids
|
||||
self.global_cond_ids = global_cond_ids
|
||||
self.input_concat_ids = input_concat_ids
|
||||
self.local_add_cond_ids = local_add_cond_ids
|
||||
self.modular_local_cond_ids = modular_local_cond_ids
|
||||
self.prepend_cond_ids = prepend_cond_ids
|
||||
self.min_input_length = min_input_length
|
||||
self.mask_padding_attention = mask_padding_attention
|
||||
self.use_effective_length_for_schedule = use_effective_length_for_schedule
|
||||
|
||||
self.dist_shift = None
|
||||
if distribution_shift_options is not None:
|
||||
self.dist_shift = self._create_dist_shift(distribution_shift_options)
|
||||
|
||||
# Sampling dist_shift: separate config for inference-time schedule
|
||||
if sampling_distribution_shift_options is not None:
|
||||
self.sampling_dist_shift = self._create_dist_shift(sampling_distribution_shift_options)
|
||||
else:
|
||||
# Default: seq_len-invariant LogSNR shift matching legacy log_snr_sampling=True
|
||||
self.sampling_dist_shift = LogSNRShift(rate=0, anchor_logsnr=-6.2, logsnr_end=2.0)
|
||||
|
||||
@staticmethod
|
||||
def _create_dist_shift(options: dict):
|
||||
"""Create a distribution shift object from config options."""
|
||||
dist_shift_type = options.get("type", "full")
|
||||
dist_shift_kwargs = {k: v for k, v in options.items() if k != "type"}
|
||||
if dist_shift_type == "none":
|
||||
return IdentityDistributionShift()
|
||||
elif dist_shift_type == "flux":
|
||||
return FluxDistributionShift(**dist_shift_kwargs)
|
||||
elif dist_shift_type == "full":
|
||||
return DistributionShift(**dist_shift_kwargs)
|
||||
elif dist_shift_type == "logsnr":
|
||||
return LogSNRShift(**dist_shift_kwargs)
|
||||
else:
|
||||
raise ValueError(f"Unknown distribution shift type: {dist_shift_type}. Expected 'none', 'flux', 'full', or 'logsnr'.")
|
||||
|
||||
def get_conditioning_inputs(self, conditioning_tensors: tp.Dict[str, tp.Any], negative=False):
|
||||
cross_attention_input = None
|
||||
cross_attention_masks = None
|
||||
global_cond = None
|
||||
input_concat_cond = None
|
||||
prepend_cond = None
|
||||
prepend_cond_mask = None
|
||||
local_add_cond = None
|
||||
modular_local_cond = None
|
||||
|
||||
if len(self.cross_attn_cond_ids) > 0:
|
||||
# Concatenate all cross-attention inputs over the sequence dimension
|
||||
# Assumes that the cross-attention inputs are of shape (batch, seq, channels)
|
||||
cross_attention_input = []
|
||||
cross_attention_masks = []
|
||||
|
||||
for key in self.cross_attn_cond_ids:
|
||||
cross_attn_in, cross_attn_mask = conditioning_tensors[key]
|
||||
|
||||
# Add sequence dimension if it's not there
|
||||
if len(cross_attn_in.shape) == 2:
|
||||
cross_attn_in = cross_attn_in.unsqueeze(1)
|
||||
cross_attn_mask = cross_attn_mask.unsqueeze(1)
|
||||
|
||||
cross_attention_input.append(cross_attn_in)
|
||||
cross_attention_masks.append(cross_attn_mask)
|
||||
|
||||
cross_attention_input = torch.cat(cross_attention_input, dim=1)
|
||||
cross_attention_masks = torch.cat(cross_attention_masks, dim=1)
|
||||
|
||||
if len(self.global_cond_ids) > 0:
|
||||
# Concatenate all global conditioning inputs over the channel dimension
|
||||
# Assumes that the global conditioning inputs are of shape (batch, channels)
|
||||
global_conds = []
|
||||
for key in self.global_cond_ids:
|
||||
global_cond_input = conditioning_tensors[key][0]
|
||||
|
||||
global_conds.append(global_cond_input)
|
||||
|
||||
# Concatenate over the channel dimension
|
||||
global_cond = torch.cat(global_conds, dim=-1)
|
||||
|
||||
if len(global_cond.shape) == 3:
|
||||
global_cond = global_cond.squeeze(1)
|
||||
|
||||
if len(self.input_concat_ids) > 0:
|
||||
# Concatenate all input concat conditioning inputs over the channel dimension
|
||||
# Assumes that the input concat conditioning inputs are of shape (batch, channels, seq)
|
||||
input_concat_cond = torch.cat([conditioning_tensors[key][0] for key in self.input_concat_ids], dim=1)
|
||||
|
||||
if len(self.local_add_cond_ids) > 0:
|
||||
# Concatenate all local conditioning inputs over the channel dimension
|
||||
# Assumes that the local conditioning inputs are of shape (batch, channels, seq)
|
||||
local_add_cond = torch.cat([conditioning_tensors[key][0] for key in self.local_add_cond_ids], dim=1)
|
||||
|
||||
if len(self.modular_local_cond_ids) > 0:
|
||||
# Keep modular local conditioning as a dict of tensors (not concatenated)
|
||||
# Each tensor is of shape (batch, channels, seq)
|
||||
modular_local_cond = {}
|
||||
for key in self.modular_local_cond_ids:
|
||||
if key in conditioning_tensors:
|
||||
modular_local_cond[key] = conditioning_tensors[key][0]
|
||||
# Only set if we have any conditioning
|
||||
if len(modular_local_cond) == 0:
|
||||
modular_local_cond = None
|
||||
|
||||
if len(self.prepend_cond_ids) > 0:
|
||||
# Concatenate all prepend conditioning inputs over the sequence dimension
|
||||
# Assumes that the prepend conditioning inputs are of shape (batch, seq, channels)
|
||||
prepend_conds = []
|
||||
prepend_cond_masks = []
|
||||
|
||||
for key in self.prepend_cond_ids:
|
||||
prepend_cond_input, prepend_cond_mask = conditioning_tensors[key]
|
||||
prepend_conds.append(prepend_cond_input)
|
||||
prepend_cond_masks.append(prepend_cond_mask)
|
||||
|
||||
prepend_cond = torch.cat(prepend_conds, dim=1)
|
||||
prepend_cond_mask = torch.cat(prepend_cond_masks, dim=1)
|
||||
|
||||
if negative:
|
||||
return {
|
||||
"negative_cross_attn_cond": cross_attention_input,
|
||||
"negative_cross_attn_mask": cross_attention_masks,
|
||||
"negative_global_cond": global_cond,
|
||||
"negative_input_concat_cond": input_concat_cond
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"cross_attn_cond": cross_attention_input,
|
||||
"cross_attn_mask": cross_attention_masks,
|
||||
"global_cond": global_cond,
|
||||
"input_concat_cond": input_concat_cond,
|
||||
"local_add_cond": local_add_cond,
|
||||
"modular_local_cond": modular_local_cond,
|
||||
"prepend_cond": prepend_cond,
|
||||
"prepend_cond_mask": prepend_cond_mask
|
||||
}
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor, cond: tp.Dict[str, tp.Any], **kwargs):
|
||||
return self.model(x, t, **self.get_conditioning_inputs(cond), **kwargs)
|
||||
|
||||
|
||||
class DiTWrapper(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
diffusion_objective: str,
|
||||
*args,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.diffusion_objective = diffusion_objective
|
||||
|
||||
self.model = DiffusionTransformer(diffusion_objective=diffusion_objective, *args, **kwargs)
|
||||
|
||||
def forward(self,
|
||||
x,
|
||||
t,
|
||||
cross_attn_cond=None,
|
||||
cross_attn_mask=None,
|
||||
negative_cross_attn_cond=None,
|
||||
negative_cross_attn_mask=None,
|
||||
input_concat_cond=None,
|
||||
local_add_cond=None,
|
||||
negative_input_concat_cond=None,
|
||||
global_cond=None,
|
||||
negative_global_cond=None,
|
||||
prepend_cond=None,
|
||||
prepend_cond_mask=None,
|
||||
cfg_scale=1.0,
|
||||
cfg_dropout_prob: float = 0.0,
|
||||
batch_cfg: bool = True,
|
||||
rescale_cfg: bool = False,
|
||||
scale_phi: float = 0.0,
|
||||
**kwargs):
|
||||
|
||||
assert batch_cfg, "batch_cfg must be True for DiTWrapper"
|
||||
#assert negative_input_concat_cond is None, "negative_input_concat_cond is not supported for DiTWrapper"
|
||||
|
||||
return self.model(
|
||||
x,
|
||||
t,
|
||||
cross_attn_cond=cross_attn_cond,
|
||||
cross_attn_cond_mask=cross_attn_mask,
|
||||
negative_cross_attn_cond=negative_cross_attn_cond,
|
||||
negative_cross_attn_mask=negative_cross_attn_mask,
|
||||
input_concat_cond=input_concat_cond,
|
||||
prepend_cond=prepend_cond,
|
||||
prepend_cond_mask=prepend_cond_mask,
|
||||
cfg_scale=cfg_scale,
|
||||
cfg_dropout_prob=cfg_dropout_prob,
|
||||
scale_phi=scale_phi,
|
||||
global_embed=global_cond,
|
||||
local_add_cond=local_add_cond,
|
||||
**kwargs)
|
||||
|
||||
@@ -0,0 +1,643 @@
|
||||
import typing as tp
|
||||
import math
|
||||
import torch
|
||||
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from .blocks import FourierFeatures, ExpoFourierFeatures
|
||||
from .transformer import ContinuousTransformer
|
||||
from .lora import LoRAParametrization, set_lora_strength, has_lora, enable_lora, disable_lora, filter_lora_layers
|
||||
|
||||
class DiffusionTransformer(nn.Module):
|
||||
def __init__(self,
|
||||
io_channels=32,
|
||||
patch_size=1,
|
||||
embed_dim=768,
|
||||
cond_token_dim=0,
|
||||
project_cond_tokens=True,
|
||||
global_cond_dim=0,
|
||||
project_global_cond=True,
|
||||
input_concat_dim=0,
|
||||
prepend_cond_dim=0,
|
||||
depth=12,
|
||||
num_heads=8,
|
||||
transformer_type: tp.Literal["continuous_transformer", "mm_transformer"] = "continuous_transformer",
|
||||
global_cond_type: tp.Literal["prepend", "adaLN"] = "prepend",
|
||||
timestep_cond_type: tp.Literal["global", "input_concat"] = "global",
|
||||
timestep_embed_dim=None,
|
||||
diffusion_objective: tp.Literal["v", "rectified_flow", "rf_denoiser"] = "v",
|
||||
timestep_features_type: tp.Literal["learned", "expo"] = "learned",
|
||||
timestep_features_dim = 256,
|
||||
timestep_features_logsnr: bool = False,
|
||||
modular_local_cond_configs = None,
|
||||
**kwargs):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.cond_token_dim = cond_token_dim
|
||||
|
||||
# Timestep embeddings
|
||||
self.timestep_cond_type = timestep_cond_type
|
||||
self.timestep_features_logsnr = timestep_features_logsnr
|
||||
|
||||
timestep_features_dim = timestep_features_dim
|
||||
|
||||
if timestep_features_type == "expo":
|
||||
self.timestep_features = ExpoFourierFeatures(timestep_features_dim, 0.5, 10000.0)
|
||||
else:
|
||||
self.timestep_features = FourierFeatures(1, timestep_features_dim)
|
||||
|
||||
if timestep_cond_type == "global":
|
||||
timestep_embed_dim = embed_dim
|
||||
elif timestep_cond_type == "input_concat":
|
||||
assert timestep_embed_dim is not None, "timestep_embed_dim must be specified if timestep_cond_type is input_concat"
|
||||
input_concat_dim += timestep_embed_dim
|
||||
|
||||
self.to_timestep_embed = nn.Sequential(
|
||||
nn.Linear(timestep_features_dim, timestep_embed_dim, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(timestep_embed_dim, timestep_embed_dim, bias=True),
|
||||
)
|
||||
|
||||
self.diffusion_objective = diffusion_objective
|
||||
|
||||
if cond_token_dim > 0:
|
||||
# Conditioning tokens
|
||||
|
||||
cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim
|
||||
self.to_cond_embed = nn.Sequential(
|
||||
nn.Linear(cond_token_dim, cond_embed_dim, bias=False),
|
||||
nn.SiLU(),
|
||||
nn.Linear(cond_embed_dim, cond_embed_dim, bias=False)
|
||||
)
|
||||
else:
|
||||
cond_embed_dim = 0
|
||||
|
||||
if global_cond_dim > 0:
|
||||
# Global conditioning
|
||||
global_embed_dim = global_cond_dim if not project_global_cond else embed_dim
|
||||
self.to_global_embed = nn.Sequential(
|
||||
nn.Linear(global_cond_dim, global_embed_dim, bias=False),
|
||||
nn.SiLU(),
|
||||
nn.Linear(global_embed_dim, global_embed_dim, bias=False)
|
||||
)
|
||||
|
||||
if prepend_cond_dim > 0:
|
||||
# Prepend conditioning
|
||||
self.to_prepend_embed = nn.Sequential(
|
||||
nn.Linear(prepend_cond_dim, embed_dim, bias=False),
|
||||
nn.SiLU(),
|
||||
nn.Linear(embed_dim, embed_dim, bias=False)
|
||||
)
|
||||
|
||||
self.input_concat_dim = input_concat_dim
|
||||
|
||||
dim_in = io_channels + self.input_concat_dim
|
||||
|
||||
self.patch_size = patch_size
|
||||
|
||||
# Transformer
|
||||
|
||||
self.transformer_type = transformer_type
|
||||
|
||||
self.global_cond_type = global_cond_type
|
||||
|
||||
transformer_dim_out = io_channels * patch_size
|
||||
|
||||
if self.transformer_type == "continuous_transformer":
|
||||
|
||||
global_dim = None
|
||||
|
||||
if self.global_cond_type == "adaLN":
|
||||
# The global conditioning is projected to the embed_dim already at this point
|
||||
global_dim = embed_dim
|
||||
|
||||
self.transformer = ContinuousTransformer(
|
||||
dim=embed_dim,
|
||||
depth=depth,
|
||||
dim_heads=embed_dim // num_heads,
|
||||
dim_in=dim_in * patch_size,
|
||||
dim_out=transformer_dim_out,
|
||||
cross_attend = cond_token_dim > 0,
|
||||
cond_token_dim = cond_embed_dim,
|
||||
global_cond_dim=global_dim,
|
||||
modular_local_cond_configs=modular_local_cond_configs,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown transformer type: {self.transformer_type}")
|
||||
|
||||
self.preprocess_conv = nn.Conv1d(dim_in, dim_in, 1, bias=False)
|
||||
nn.init.zeros_(self.preprocess_conv.weight)
|
||||
self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
|
||||
nn.init.zeros_(self.postprocess_conv.weight)
|
||||
|
||||
# Fixed logsnr normalization range: maps logsnr to [0, 1] preserving direction (t=0→0, t=1→1)
|
||||
_LOGSNR_MIN = -12.0
|
||||
_LOGSNR_MAX = 5.0
|
||||
_LOGSNR_RANGE = _LOGSNR_MAX - _LOGSNR_MIN
|
||||
|
||||
def _t_to_logsnr_cond(self, t: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert t to normalized logsnr in [0, 1] for timestep conditioning.
|
||||
|
||||
Maps t through logsnr = log((1-t)/t), clamps to fixed range,
|
||||
then normalizes to [0, 1] preserving direction (t=0→0, t=1→1).
|
||||
"""
|
||||
t_clamped = t.float().clamp(1e-7, 1 - 1e-7)
|
||||
logsnr = torch.log((1 - t_clamped) / t_clamped)
|
||||
logsnr = logsnr.clamp(self._LOGSNR_MIN, self._LOGSNR_MAX)
|
||||
return ((self._LOGSNR_MAX - logsnr) / self._LOGSNR_RANGE).to(t.dtype)
|
||||
|
||||
def _call_transformer(self, x, *, prepend_inputs=None, cross_attn_cond=None,
|
||||
mask=None, prepend_mask=None, return_info=False,
|
||||
exit_layer_ix=None, local_add_cond=None,
|
||||
modular_local_cond=None, padding_mask=None,
|
||||
extra_args=None, **kwargs):
|
||||
"""Helper method to call transformer and handle early exit logic."""
|
||||
|
||||
output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond,
|
||||
return_info=return_info, exit_layer_ix=exit_layer_ix,
|
||||
local_add_cond=local_add_cond, modular_local_cond=modular_local_cond,
|
||||
padding_mask=padding_mask,
|
||||
**(extra_args or {}), **kwargs)
|
||||
|
||||
if return_info:
|
||||
output, info = output
|
||||
|
||||
# Avoid postprocessing on early exit
|
||||
if exit_layer_ix is not None:
|
||||
if return_info:
|
||||
return output, info
|
||||
else:
|
||||
return output
|
||||
|
||||
return (output, info) if return_info and 'info' in locals() else output
|
||||
|
||||
def _forward(
|
||||
self,
|
||||
x,
|
||||
t,
|
||||
mask=None,
|
||||
cross_attn_cond=None,
|
||||
cross_attn_cond_mask=None,
|
||||
input_concat_cond=None,
|
||||
local_add_cond=None,
|
||||
modular_local_cond=None,
|
||||
global_embed=None,
|
||||
prepend_cond=None,
|
||||
prepend_cond_mask=None,
|
||||
padding_mask=None,
|
||||
return_info=False,
|
||||
exit_layer_ix=None,
|
||||
**kwargs):
|
||||
|
||||
if cross_attn_cond is not None:
|
||||
cross_attn_cond = self.to_cond_embed(cross_attn_cond)
|
||||
|
||||
if global_embed is not None:
|
||||
# Project the global conditioning to the embedding dimension
|
||||
global_embed = self.to_global_embed(global_embed)
|
||||
|
||||
prepend_inputs = None
|
||||
prepend_mask = None
|
||||
prepend_length = 0
|
||||
if prepend_cond is not None:
|
||||
# Project the prepend conditioning to the embedding dimension
|
||||
prepend_cond = self.to_prepend_embed(prepend_cond)
|
||||
|
||||
prepend_inputs = prepend_cond
|
||||
if prepend_cond_mask is not None:
|
||||
prepend_mask = prepend_cond_mask
|
||||
|
||||
prepend_length = prepend_cond.shape[1]
|
||||
|
||||
if input_concat_cond is not None:
|
||||
# Interpolate input_concat_cond to the same length as x
|
||||
if input_concat_cond.shape[2] != x.shape[2]:
|
||||
input_concat_cond = F.interpolate(input_concat_cond, (x.shape[2], ), mode='nearest')
|
||||
|
||||
x = torch.cat([x, input_concat_cond], dim=1)
|
||||
|
||||
if local_add_cond is not None:
|
||||
local_add_cond = rearrange(local_add_cond, "b c t -> b t c")
|
||||
|
||||
# Rearrange modular_local_cond tensors
|
||||
if modular_local_cond is not None:
|
||||
modular_local_cond = {
|
||||
k: rearrange(v, "b c t -> b t c")
|
||||
for k, v in modular_local_cond.items()
|
||||
}
|
||||
|
||||
# Get the batch of timestep embeddings
|
||||
t_cond = self._t_to_logsnr_cond(t) if self.timestep_features_logsnr else t
|
||||
# Convert to model dtype for linear layers (t itself is kept in float32 for precision)
|
||||
# x has already been converted to model dtype in the outer forward() method
|
||||
t_cond = t_cond.to(x.dtype)
|
||||
timestep_embed = self.to_timestep_embed(self.timestep_features(t_cond[:, None])) # (b, embed_dim)
|
||||
|
||||
# Timestep embedding is considered a global embedding. Add to the global conditioning if it exists
|
||||
|
||||
if self.timestep_cond_type == "global":
|
||||
if global_embed is not None:
|
||||
global_embed = global_embed + timestep_embed
|
||||
else:
|
||||
global_embed = timestep_embed
|
||||
elif self.timestep_cond_type == "input_concat":
|
||||
x = torch.cat([x, timestep_embed.unsqueeze(2).expand(-1, -1, x.shape[2])], dim=1)
|
||||
|
||||
# Add the global_embed to the prepend inputs if there is no global conditioning support in the transformer
|
||||
if self.global_cond_type == "prepend" and global_embed is not None:
|
||||
if prepend_inputs is None:
|
||||
# Prepend inputs are just the global embed, and the mask is all ones
|
||||
prepend_inputs = global_embed.unsqueeze(1)
|
||||
prepend_mask = torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)
|
||||
else:
|
||||
# Prepend inputs are the prepend conditioning + the global embed
|
||||
prepend_inputs = torch.cat([prepend_inputs, global_embed.unsqueeze(1)], dim=1)
|
||||
prepend_mask = torch.cat([prepend_mask, torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)], dim=1)
|
||||
|
||||
prepend_length = prepend_inputs.shape[1]
|
||||
|
||||
x = self.preprocess_conv(x) + x
|
||||
|
||||
x = rearrange(x, "b c t -> b t c")
|
||||
|
||||
extra_args = {}
|
||||
|
||||
if self.global_cond_type == "adaLN":
|
||||
extra_args["global_cond"] = global_embed
|
||||
|
||||
if self.patch_size > 1:
|
||||
x = rearrange(x, "b (t p) c -> b t (c p)", p=self.patch_size)
|
||||
|
||||
result = self._call_transformer(
|
||||
x,
|
||||
prepend_inputs=prepend_inputs,
|
||||
cross_attn_cond=cross_attn_cond,
|
||||
mask=mask,
|
||||
prepend_mask=prepend_mask,
|
||||
return_info=return_info,
|
||||
exit_layer_ix=exit_layer_ix,
|
||||
local_add_cond=local_add_cond,
|
||||
modular_local_cond=modular_local_cond,
|
||||
padding_mask=padding_mask,
|
||||
extra_args=extra_args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Handle early exit (result contains both output and info)
|
||||
if exit_layer_ix is not None:
|
||||
return result
|
||||
|
||||
output = result[0] if return_info else result
|
||||
if return_info:
|
||||
info = result[1]
|
||||
|
||||
output = rearrange(output, "b t c -> b c t")[:,:,prepend_length:]
|
||||
|
||||
if self.patch_size > 1:
|
||||
output = rearrange(output, "b (c p) t -> b c (t p)", p=self.patch_size)
|
||||
|
||||
output = self.postprocess_conv(output) + output
|
||||
|
||||
if return_info:
|
||||
return output, info
|
||||
|
||||
return output
|
||||
|
||||
def apg_project(self, v0, v1, padding_mask=None):
|
||||
"""
|
||||
Project v0 into components parallel and orthogonal to v1.
|
||||
|
||||
Args:
|
||||
v0: Tensor to project (B, C, T)
|
||||
v1: Reference direction (B, C, T)
|
||||
padding_mask: Optional mask (B, T) where True = valid, False = padding.
|
||||
If provided, only valid positions contribute to the projection.
|
||||
"""
|
||||
dtype = v0.dtype
|
||||
v0, v1 = v0.float(), v1.float()
|
||||
|
||||
if padding_mask is not None:
|
||||
# Expand mask to match tensor shape: (B, T) -> (B, 1, T)
|
||||
mask = padding_mask.unsqueeze(1).float()
|
||||
# Zero out padding positions for projection computation
|
||||
v0_masked = v0 * mask
|
||||
v1_masked = v1 * mask
|
||||
# Normalize only over valid positions
|
||||
v1_norm = v1_masked.norm(dim=[-1, -2], keepdim=True).clamp(min=1e-8)
|
||||
v1_normalized = v1_masked / v1_norm
|
||||
# Compute projection using masked values
|
||||
v0_parallel = (v0_masked * v1_normalized).sum(dim=[-1, -2], keepdim=True) * v1_normalized
|
||||
# Orthogonal component: subtract parallel from original (not masked) v0
|
||||
# but apply mask to ensure padding stays zero
|
||||
v0_orthogonal = (v0 - (v0 * v1_normalized).sum(dim=[-1, -2], keepdim=True) * v1_normalized) * mask
|
||||
else:
|
||||
v1 = torch.nn.functional.normalize(v1, dim=[-1, -2])
|
||||
v0_parallel = (v0 * v1).sum(dim=[-1, -2], keepdim=True) * v1
|
||||
v0_orthogonal = v0 - v0_parallel
|
||||
|
||||
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
t,
|
||||
cross_attn_cond=None,
|
||||
cross_attn_cond_mask=None,
|
||||
negative_cross_attn_cond=None,
|
||||
negative_cross_attn_mask=None,
|
||||
input_concat_cond=None,
|
||||
local_add_cond=None,
|
||||
modular_local_cond=None,
|
||||
global_embed=None,
|
||||
negative_global_embed=None,
|
||||
prepend_cond=None,
|
||||
prepend_cond_mask=None,
|
||||
padding_mask=None,
|
||||
cfg_scale=1.0,
|
||||
cfg_dropout_prob=0.0,
|
||||
cfg_interval = (0, 1),
|
||||
lora_interval = (0, 1),
|
||||
lora_layer_filter = "",
|
||||
lora_configs = None,
|
||||
causal=False,
|
||||
scale_phi=0.0,
|
||||
cfg_norm_threshold=0.0,
|
||||
apg_scale=1.0,
|
||||
mask=None,
|
||||
return_info=False,
|
||||
exit_layer_ix=None,
|
||||
**kwargs):
|
||||
|
||||
assert not causal, "Causal mode is not supported for DiffusionTransformer"
|
||||
|
||||
model_dtype = next(self.parameters()).dtype
|
||||
|
||||
x = x.to(model_dtype)
|
||||
|
||||
# Keep t in float32: the logsnr transform log((1-t)/t) amplifies bf16
|
||||
# quantization error ~380x near t=1, causing catastrophic conditioning errors.
|
||||
# t is a 1D batch-size tensor so float32 has zero memory impact.
|
||||
t = t.float()
|
||||
|
||||
if cross_attn_cond is not None:
|
||||
cross_attn_cond = cross_attn_cond.to(model_dtype)
|
||||
|
||||
if negative_cross_attn_cond is not None:
|
||||
negative_cross_attn_cond = negative_cross_attn_cond.to(model_dtype)
|
||||
|
||||
if input_concat_cond is not None:
|
||||
input_concat_cond = input_concat_cond.to(model_dtype)
|
||||
|
||||
if local_add_cond is not None:
|
||||
local_add_cond = local_add_cond.to(model_dtype)
|
||||
|
||||
if modular_local_cond is not None:
|
||||
modular_local_cond = {k: v.to(model_dtype) for k, v in modular_local_cond.items()}
|
||||
|
||||
if global_embed is not None:
|
||||
global_embed = global_embed.to(model_dtype)
|
||||
|
||||
if negative_global_embed is not None:
|
||||
negative_global_embed = negative_global_embed.to(model_dtype)
|
||||
|
||||
if prepend_cond is not None:
|
||||
prepend_cond = prepend_cond.to(model_dtype)
|
||||
|
||||
if cross_attn_cond_mask is not None:
|
||||
cross_attn_cond_mask = cross_attn_cond_mask.bool()
|
||||
|
||||
cross_attn_cond_mask = None # Temporarily disabling conditioning masks due to kernel issue for flash attention
|
||||
|
||||
if prepend_cond_mask is not None:
|
||||
prepend_cond_mask = prepend_cond_mask.bool()
|
||||
|
||||
# Early exit bypasses CFG processing
|
||||
if exit_layer_ix is not None:
|
||||
assert self.transformer_type == "continuous_transformer", "exit_layer_ix is only supported for continuous_transformer"
|
||||
return self._forward(
|
||||
x,
|
||||
t,
|
||||
cross_attn_cond=cross_attn_cond,
|
||||
cross_attn_cond_mask=cross_attn_cond_mask,
|
||||
input_concat_cond=input_concat_cond,
|
||||
local_add_cond=local_add_cond,
|
||||
modular_local_cond=modular_local_cond,
|
||||
global_embed=global_embed,
|
||||
prepend_cond=prepend_cond,
|
||||
prepend_cond_mask=prepend_cond_mask,
|
||||
padding_mask=padding_mask,
|
||||
mask=mask,
|
||||
return_info=return_info,
|
||||
exit_layer_ix=exit_layer_ix,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
# CFG dropout
|
||||
if cfg_dropout_prob > 0.0 and cfg_scale == 1.0:
|
||||
if cross_attn_cond is not None:
|
||||
null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)
|
||||
dropout_mask = torch.bernoulli(torch.full((cross_attn_cond.shape[0], 1, 1), cfg_dropout_prob, device=cross_attn_cond.device)).to(torch.bool)
|
||||
cross_attn_cond = torch.where(dropout_mask, null_embed, cross_attn_cond)
|
||||
|
||||
if prepend_cond is not None:
|
||||
null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)
|
||||
dropout_mask = torch.bernoulli(torch.full((prepend_cond.shape[0], 1, 1), cfg_dropout_prob, device=prepend_cond.device)).to(torch.bool)
|
||||
prepend_cond = torch.where(dropout_mask, null_embed, prepend_cond)
|
||||
|
||||
if self.diffusion_objective == "v":
|
||||
sigma = torch.sin(t * math.pi / 2)
|
||||
alpha = torch.cos(t * math.pi / 2)
|
||||
elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
|
||||
sigma = t
|
||||
|
||||
# LoRA interval
|
||||
if has_lora(self):
|
||||
if lora_configs is not None:
|
||||
# Multi-LoRA: per-LoRA interval and layer filter
|
||||
for lora_config in lora_configs:
|
||||
idx = lora_config["lora_index"]
|
||||
interval = lora_config.get("interval", (0, 1))
|
||||
layer_filter = lora_config.get("layer_filter", "")
|
||||
if interval[0] <= sigma[0] <= interval[1]:
|
||||
enable_lora(self, lora_index=idx)
|
||||
filter_lora_layers(self, layer_filter, lora_index=idx)
|
||||
else:
|
||||
disable_lora(self, lora_index=idx)
|
||||
else:
|
||||
# Legacy single-LoRA path
|
||||
if lora_interval[0] <= sigma[0] <= lora_interval[1]:
|
||||
enable_lora(self)
|
||||
filter_lora_layers(self, lora_layer_filter)
|
||||
else:
|
||||
disable_lora(self)
|
||||
|
||||
if cfg_scale != 1.0 and (cross_attn_cond is not None or prepend_cond is not None) and (cfg_interval[0] <= sigma[0] <= cfg_interval[1]):
|
||||
|
||||
# Classifier-free guidance
|
||||
# Concatenate conditioned and unconditioned inputs on the batch dimension
|
||||
batch_inputs = torch.cat([x, x], dim=0)
|
||||
batch_timestep = torch.cat([t, t], dim=0)
|
||||
|
||||
if global_embed is not None:
|
||||
batch_global_cond = torch.cat([global_embed, global_embed], dim=0)
|
||||
else:
|
||||
batch_global_cond = None
|
||||
|
||||
if input_concat_cond is not None:
|
||||
batch_input_concat_cond = torch.cat([input_concat_cond, input_concat_cond], dim=0)
|
||||
else:
|
||||
batch_input_concat_cond = None
|
||||
|
||||
if local_add_cond is not None:
|
||||
batch_local_add_cond = torch.cat([local_add_cond, local_add_cond], dim=0)
|
||||
else:
|
||||
batch_local_add_cond = None
|
||||
|
||||
if modular_local_cond is not None:
|
||||
batch_modular_local_cond = {k: torch.cat([v, v], dim=0) for k, v in modular_local_cond.items()}
|
||||
else:
|
||||
batch_modular_local_cond = None
|
||||
|
||||
batch_cond = None
|
||||
batch_cond_masks = None
|
||||
|
||||
# Handle CFG for cross-attention conditioning
|
||||
if cross_attn_cond is not None:
|
||||
|
||||
null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)
|
||||
|
||||
# For negative cross-attention conditioning, replace the null embed with the negative cross-attention conditioning
|
||||
if negative_cross_attn_cond is not None:
|
||||
|
||||
# If there's a negative cross-attention mask, set the masked tokens to the null embed
|
||||
if negative_cross_attn_mask is not None:
|
||||
negative_cross_attn_mask = negative_cross_attn_mask.to(torch.bool).unsqueeze(2)
|
||||
|
||||
negative_cross_attn_cond = torch.where(negative_cross_attn_mask, negative_cross_attn_cond, null_embed)
|
||||
|
||||
batch_cond = torch.cat([cross_attn_cond, negative_cross_attn_cond], dim=0)
|
||||
|
||||
else:
|
||||
batch_cond = torch.cat([cross_attn_cond, null_embed], dim=0)
|
||||
|
||||
if cross_attn_cond_mask is not None:
|
||||
batch_cond_masks = torch.cat([cross_attn_cond_mask, cross_attn_cond_mask], dim=0)
|
||||
|
||||
batch_prepend_cond = None
|
||||
batch_prepend_cond_mask = None
|
||||
|
||||
if prepend_cond is not None:
|
||||
|
||||
null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)
|
||||
|
||||
batch_prepend_cond = torch.cat([prepend_cond, null_embed], dim=0)
|
||||
|
||||
if prepend_cond_mask is not None:
|
||||
batch_prepend_cond_mask = torch.cat([prepend_cond_mask, prepend_cond_mask], dim=0)
|
||||
|
||||
|
||||
if mask is not None:
|
||||
batch_masks = torch.cat([mask, mask], dim=0)
|
||||
else:
|
||||
batch_masks = None
|
||||
|
||||
if padding_mask is not None:
|
||||
batch_padding_mask = torch.cat([padding_mask, padding_mask], dim=0)
|
||||
else:
|
||||
batch_padding_mask = None
|
||||
|
||||
batch_output = self._forward(
|
||||
batch_inputs,
|
||||
batch_timestep,
|
||||
cross_attn_cond=batch_cond,
|
||||
cross_attn_cond_mask=batch_cond_masks,
|
||||
mask = batch_masks,
|
||||
input_concat_cond = batch_input_concat_cond,
|
||||
local_add_cond = batch_local_add_cond,
|
||||
modular_local_cond=batch_modular_local_cond,
|
||||
global_embed = batch_global_cond,
|
||||
prepend_cond = batch_prepend_cond,
|
||||
prepend_cond_mask = batch_prepend_cond_mask,
|
||||
padding_mask = batch_padding_mask,
|
||||
return_info = return_info,
|
||||
**kwargs)
|
||||
|
||||
if return_info:
|
||||
batch_output, info = batch_output
|
||||
|
||||
cond_output, uncond_output = torch.chunk(batch_output, 2, dim=0)
|
||||
|
||||
if self.diffusion_objective == "v":
|
||||
cond_denoised = x * alpha[:, None, None] - cond_output * sigma[:, None, None]
|
||||
uncond_denoised = x * alpha[:, None, None] - uncond_output * sigma[:, None, None]
|
||||
|
||||
elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
|
||||
cond_denoised = x - cond_output * sigma[:, None, None]
|
||||
uncond_denoised = x - uncond_output * sigma[:, None, None]
|
||||
|
||||
diff = cond_denoised - uncond_denoised
|
||||
|
||||
if cfg_norm_threshold > 0:
|
||||
if padding_mask is not None:
|
||||
# Only compute norm over valid positions
|
||||
mask = padding_mask.unsqueeze(1).float() # (B, 1, T)
|
||||
diff_masked = diff * mask
|
||||
diff_norm = diff_masked.norm(p=2, dim=[-1, -2], keepdim=True)
|
||||
else:
|
||||
diff_norm = diff.norm(p=2, dim=[-1, -2], keepdim=True)
|
||||
scale_factor = torch.minimum(torch.ones_like(diff), cfg_norm_threshold / diff_norm)
|
||||
diff *= scale_factor
|
||||
|
||||
if apg_scale == 0.0:
|
||||
# Vanilla CFG: use full diff
|
||||
cfg_diff = diff
|
||||
elif apg_scale == 1.0:
|
||||
# Full APG: use only orthogonal component
|
||||
_, diff_orthogonal = self.apg_project(diff, cond_denoised, padding_mask=padding_mask)
|
||||
cfg_diff = diff_orthogonal
|
||||
else:
|
||||
# Blended APG: interpolate between full diff and orthogonal
|
||||
diff_parallel, diff_orthogonal = self.apg_project(diff, cond_denoised, padding_mask=padding_mask)
|
||||
cfg_diff = apg_scale * diff_orthogonal + (1 - apg_scale) * diff
|
||||
|
||||
cfg_denoised = cond_denoised + (cfg_scale - 1) * cfg_diff
|
||||
|
||||
if self.diffusion_objective == "v":
|
||||
output = (x * alpha[:, None, None] - cfg_denoised) / sigma[:, None, None]
|
||||
elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
|
||||
output = (x - cfg_denoised) / sigma[:, None, None]
|
||||
|
||||
# CFG Rescale
|
||||
if scale_phi != 0.0:
|
||||
cond_out_std = cond_output.std(dim=1, keepdim=True)
|
||||
out_cfg_std = output.std(dim=1, keepdim=True)
|
||||
output = scale_phi * (output * (cond_out_std/out_cfg_std)) + (1-scale_phi) * output
|
||||
|
||||
if return_info:
|
||||
info["uncond_output"] = uncond_output
|
||||
return output, info
|
||||
|
||||
return output
|
||||
|
||||
else:
|
||||
return self._forward(
|
||||
x,
|
||||
t,
|
||||
cross_attn_cond=cross_attn_cond,
|
||||
cross_attn_cond_mask=cross_attn_cond_mask,
|
||||
input_concat_cond=input_concat_cond,
|
||||
local_add_cond=local_add_cond,
|
||||
modular_local_cond=modular_local_cond,
|
||||
global_embed=global_embed,
|
||||
prepend_cond=prepend_cond,
|
||||
prepend_cond_mask=prepend_cond_mask,
|
||||
padding_mask=padding_mask,
|
||||
mask=mask,
|
||||
return_info=return_info,
|
||||
**kwargs
|
||||
)
|
||||
@@ -0,0 +1,242 @@
|
||||
import random
|
||||
import torch
|
||||
from enum import Enum
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
class MaskType(Enum):
|
||||
RANDOM_SEGMENTS = 0 # Legacy: uncontrolled ratio, overlapping segments
|
||||
FULL_MASK = 1
|
||||
CAUSAL_MASK = 2
|
||||
RANDOM_SPANS = 3 # New: explicit ratio control, non-overlapping spans
|
||||
|
||||
|
||||
def _generate_random_spans_mask(
|
||||
real_sequence_length: int,
|
||||
sequence_length: int,
|
||||
device: torch.device,
|
||||
max_spans: int = 4,
|
||||
mask_ratio_range: Tuple[float, float] = (0.2, 1.0),
|
||||
span_count_weights: Optional[List[float]] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a mask with 1-N non-overlapping contiguous spans covering a target ratio.
|
||||
|
||||
Samples a target mask ratio from U[mask_ratio_range], then allocates that budget
|
||||
across a variable number of spans with Dirichlet-like random allocation.
|
||||
|
||||
Args:
|
||||
real_sequence_length: Number of valid (non-padding) tokens.
|
||||
sequence_length: Total sequence length including padding.
|
||||
device: Torch device.
|
||||
max_spans: Maximum number of masked spans.
|
||||
mask_ratio_range: (lo, hi) range for uniform sampling of target mask ratio.
|
||||
span_count_weights: Weights for sampling number of spans [1..max_spans].
|
||||
Defaults to geometric decay [4, 2, 1, 1] (biased toward fewer spans).
|
||||
"""
|
||||
item_mask = torch.ones((1, 1, sequence_length), device=device, dtype=torch.float32)
|
||||
|
||||
if real_sequence_length == 0:
|
||||
return item_mask
|
||||
|
||||
# 1. Sample target mask ratio
|
||||
lo, hi = mask_ratio_range
|
||||
target_ratio = random.uniform(lo, hi)
|
||||
target_masked_tokens = max(1, int(target_ratio * real_sequence_length))
|
||||
|
||||
# 2. Sample number of spans (biased toward fewer spans)
|
||||
if span_count_weights is None:
|
||||
span_count_weights = [4, 2, 1, 1]
|
||||
|
||||
# Truncate/pad weights to max_spans
|
||||
weights = span_count_weights[:max_spans]
|
||||
while len(weights) < max_spans:
|
||||
weights.append(weights[-1])
|
||||
|
||||
num_spans = random.choices(range(1, max_spans + 1), weights=weights, k=1)[0]
|
||||
num_spans = min(num_spans, real_sequence_length)
|
||||
|
||||
# 3. Allocate target tokens across spans (Dirichlet-like)
|
||||
raw_weights = [random.random() for _ in range(num_spans)]
|
||||
total_weight = sum(raw_weights)
|
||||
span_lengths = [max(1, int(w / total_weight * target_masked_tokens)) for w in raw_weights]
|
||||
|
||||
# Adjust to hit exact target
|
||||
diff = target_masked_tokens - sum(span_lengths)
|
||||
if diff > 0:
|
||||
span_lengths[0] += diff
|
||||
elif diff < 0:
|
||||
for i in range(len(span_lengths)):
|
||||
reduce = min(-diff, span_lengths[i] - 1)
|
||||
span_lengths[i] -= reduce
|
||||
diff += reduce
|
||||
if diff >= 0:
|
||||
break
|
||||
|
||||
# 4. Place spans without overlap (largest first for better packing)
|
||||
span_lengths.sort(reverse=True)
|
||||
placed_spans = []
|
||||
|
||||
for length in span_lengths:
|
||||
length = min(length, real_sequence_length)
|
||||
if length <= 0:
|
||||
continue
|
||||
|
||||
max_start = real_sequence_length - length
|
||||
if max_start < 0:
|
||||
continue
|
||||
|
||||
# Try random placement avoiding overlaps
|
||||
placed = False
|
||||
for _ in range(50):
|
||||
start = random.randint(0, max_start)
|
||||
end = start + length
|
||||
if not any(s < end and start < e for s, e in placed_spans):
|
||||
placed_spans.append((start, end))
|
||||
item_mask[:, :, start:end] = 0
|
||||
placed = True
|
||||
break
|
||||
|
||||
# Fallback: scan for first valid position
|
||||
if not placed:
|
||||
for start in range(max_start + 1):
|
||||
end = start + length
|
||||
if not any(s < end and start < e for s, e in placed_spans):
|
||||
placed_spans.append((start, end))
|
||||
item_mask[:, :, start:end] = 0
|
||||
break
|
||||
|
||||
return item_mask
|
||||
|
||||
|
||||
def random_inpaint_mask(
|
||||
sequence: torch.Tensor,
|
||||
padding_masks: torch.Tensor,
|
||||
max_mask_segments: int = 10,
|
||||
mask_type_probabilities: Optional[List[float]] = None,
|
||||
mask_padding: bool = False,
|
||||
force_mask_type: Optional[MaskType] = None,
|
||||
# RANDOM_SPANS parameters
|
||||
max_spans: int = 4,
|
||||
mask_ratio_range: Tuple[float, float] = (0.2, 1.0),
|
||||
span_count_weights: Optional[List[float]] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Generates random inpainting masks for a batch of latent audio sequences.
|
||||
The output inpainting mask has 0 where data should be inpainted, and 1 where data is provided.
|
||||
|
||||
Supports two segment masking modes:
|
||||
- RANDOM_SEGMENTS (legacy): Uncontrolled mask ratio, overlapping segments allowed.
|
||||
- RANDOM_SPANS (new): Explicit target mask ratio sampled from mask_ratio_range,
|
||||
non-overlapping contiguous spans with controllable count distribution.
|
||||
|
||||
Args:
|
||||
sequence: The input sequence tensor of shape (b, c, sequence_length).
|
||||
padding_masks: A tensor of shape (b, sequence_length)
|
||||
where 1 indicates real data latents and 0 indicates latents encoding silence padding.
|
||||
max_mask_segments: (RANDOM_SEGMENTS only) Maximum number of segments.
|
||||
mask_type_probabilities: Probabilities for each mask type. Length must match the
|
||||
number of MaskType values used:
|
||||
- 3 elements: [P(RANDOM_SEGMENTS), P(FULL_MASK), P(CAUSAL_MASK)]
|
||||
- 4 elements: [P(RANDOM_SEGMENTS), P(FULL_MASK), P(CAUSAL_MASK), P(RANDOM_SPANS)]
|
||||
If None, defaults to [0.1, 0.8, 0.1] (legacy 3-type behavior).
|
||||
mask_padding: If True, zero out padding region in the inpaint mask for
|
||||
RANDOM_SEGMENTS/RANDOM_SPANS/CAUSAL_MASK. Use when training with
|
||||
attention masking for variable-length sequences.
|
||||
max_spans: (RANDOM_SPANS only) Maximum number of masked spans.
|
||||
mask_ratio_range: (RANDOM_SPANS only) (lo, hi) for uniform sampling of target mask ratio.
|
||||
span_count_weights: (RANDOM_SPANS only) Weights for sampling number of spans [1..max_spans].
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- masked_sequence: The sequence with masks applied (original sequence where mask is 1,
|
||||
and usually 0 or a placeholder where mask is 0).
|
||||
- inpaint_mask: The generated inpainting mask tensor (0 for inpaint, 1 for keep).
|
||||
"""
|
||||
b, _, sequence_length = sequence.size()
|
||||
|
||||
# Skip probability validation when forcing a specific mask type
|
||||
if force_mask_type is not None:
|
||||
mask_types_to_sample = None
|
||||
else:
|
||||
if mask_type_probabilities is None:
|
||||
mask_type_probabilities = [0.1, 0.8, 0.1]
|
||||
|
||||
num_probs = len(mask_type_probabilities)
|
||||
|
||||
# Determine which mask types are active based on probability list length
|
||||
if num_probs == 3:
|
||||
# Legacy: [RANDOM_SEGMENTS, FULL_MASK, CAUSAL_MASK]
|
||||
mask_types_to_sample = [MaskType.RANDOM_SEGMENTS.value, MaskType.FULL_MASK.value, MaskType.CAUSAL_MASK.value]
|
||||
elif num_probs == 4:
|
||||
# Extended: [RANDOM_SEGMENTS, FULL_MASK, CAUSAL_MASK, RANDOM_SPANS]
|
||||
mask_types_to_sample = [MaskType.RANDOM_SEGMENTS.value, MaskType.FULL_MASK.value, MaskType.CAUSAL_MASK.value, MaskType.RANDOM_SPANS.value]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"mask_type_probabilities must have 3 or 4 elements. Got {num_probs}."
|
||||
)
|
||||
|
||||
if not torch.isclose(torch.tensor(sum(mask_type_probabilities)), torch.tensor(1.0)):
|
||||
raise ValueError(
|
||||
f"mask_type_probabilities must sum to 1.0. "
|
||||
f"Current sum: {sum(mask_type_probabilities)}"
|
||||
)
|
||||
|
||||
output_masks_list = []
|
||||
|
||||
for i in range(b):
|
||||
padding_mask_single_item = padding_masks[i]
|
||||
real_sequence_length = (padding_mask_single_item == 1).sum().item()
|
||||
|
||||
item_mask = torch.ones((1, 1, sequence_length), device=sequence.device, dtype=torch.float32)
|
||||
|
||||
if force_mask_type is not None:
|
||||
current_mask_type = force_mask_type
|
||||
else:
|
||||
chosen_mask_value = random.choices(mask_types_to_sample, weights=mask_type_probabilities, k=1)[0]
|
||||
current_mask_type = MaskType(chosen_mask_value)
|
||||
|
||||
if current_mask_type == MaskType.FULL_MASK:
|
||||
item_mask = torch.zeros((1, 1, sequence_length), device=sequence.device, dtype=torch.float32)
|
||||
elif real_sequence_length == 0:
|
||||
pass # item_mask remains all ones
|
||||
else:
|
||||
if current_mask_type == MaskType.RANDOM_SEGMENTS:
|
||||
# Legacy behavior: uncontrolled ratio, overlapping segments
|
||||
num_segments = random.randint(1, max_mask_segments)
|
||||
max_len_per_segment_calc = max(1, real_sequence_length // num_segments)
|
||||
|
||||
for _ in range(num_segments):
|
||||
segment_length = random.randint(1, max_len_per_segment_calc)
|
||||
|
||||
if real_sequence_length - segment_length < 0:
|
||||
continue
|
||||
mask_start = random.randint(0, real_sequence_length - segment_length)
|
||||
item_mask[:, :, mask_start : mask_start + segment_length] = 0
|
||||
|
||||
elif current_mask_type == MaskType.RANDOM_SPANS:
|
||||
item_mask = _generate_random_spans_mask(
|
||||
real_sequence_length=real_sequence_length,
|
||||
sequence_length=sequence_length,
|
||||
device=sequence.device,
|
||||
max_spans=max_spans,
|
||||
mask_ratio_range=mask_ratio_range,
|
||||
span_count_weights=span_count_weights,
|
||||
)
|
||||
|
||||
elif current_mask_type == MaskType.CAUSAL_MASK:
|
||||
unmasked_prefix_len = random.randint(0, real_sequence_length)
|
||||
|
||||
if unmasked_prefix_len < real_sequence_length:
|
||||
item_mask[:, :, unmasked_prefix_len:real_sequence_length] = 0
|
||||
|
||||
# When using attention masking for variable-length training, zero out padding
|
||||
# so the model doesn't see it as provided context. Without attention masking,
|
||||
# padding is handled via mask_loss_weight and should remain as provided input.
|
||||
if mask_padding:
|
||||
item_mask[:, :, real_sequence_length:] = 0
|
||||
|
||||
output_masks_list.append(item_mask)
|
||||
|
||||
final_inpaint_mask = torch.cat(output_masks_list, dim=0).to(sequence.device)
|
||||
masked_sequence = sequence * final_inpaint_mask
|
||||
return masked_sequence, final_inpaint_mask
|
||||
@@ -0,0 +1,28 @@
|
||||
from .model import LoRAParametrization, add_lora, default_lora_config, merge_lora, remove_lora, remove_lora_by_index, set_lora_strength
|
||||
from .utils import (
|
||||
apply_to_lora,
|
||||
convert_lora_ckpt_to_safetensors,
|
||||
disable_lora,
|
||||
enable_lora,
|
||||
filter_lora_layers,
|
||||
get_lora_count,
|
||||
has_lora,
|
||||
get_bias_params,
|
||||
get_lora_params,
|
||||
get_lora_layers,
|
||||
get_lora_state_dict,
|
||||
load_lora_checkpoint,
|
||||
load_multiple_lora,
|
||||
name_is_lora,
|
||||
remap_lora_state_dict,
|
||||
save_lora_safetensors,
|
||||
select_lora,
|
||||
tie_weights,
|
||||
untie_weights,
|
||||
infer_global_rank,
|
||||
cast_base_to_precision,
|
||||
detect_dora_variant,
|
||||
prepare_dora_state_dict,
|
||||
resolve_adapter_type
|
||||
)
|
||||
from .loader import load_and_apply_loras
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,112 @@
|
||||
"""Helpers for loading pre-trained LoRA checkpoints onto a live model."""
|
||||
import os
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
|
||||
from .model import LoRAParametrization, add_lora
|
||||
from .utils import (
|
||||
get_lora_layers,
|
||||
infer_global_rank,
|
||||
load_lora_checkpoint,
|
||||
prepare_dora_state_dict,
|
||||
remap_lora_state_dict,
|
||||
resolve_adapter_type,
|
||||
)
|
||||
from ...verbose import vprint
|
||||
|
||||
|
||||
def load_and_apply_loras(model, lora_ckpt_paths, model_type, svd_bases_path=None):
|
||||
"""Load LoRA checkpoints from disk and attach them to `model`.
|
||||
|
||||
- Uses a two-pass approach: first resolves each LoRA's adapter type, then
|
||||
loads SVD bases (only if any LoRA is -XS), then applies each LoRA.
|
||||
- Handles legacy "dora" adapter_type via `resolve_adapter_type`.
|
||||
- Only passes svd_bases to -XS LoRAs to avoid spurious "no -XS layers"
|
||||
messages for non-XS LoRAs.
|
||||
- Sets `model.use_lora = True` and `model.lora_names = [...]` as a
|
||||
convenience for downstream UI code.
|
||||
|
||||
Args:
|
||||
model: The loaded model. For `diffusion_cond` / `diffusion_cond_inpaint`
|
||||
types, LoRAs are applied to `model.model` and `model.conditioner`;
|
||||
otherwise they are applied to `model` directly.
|
||||
lora_ckpt_paths: List of paths to `.ckpt` or `.safetensors` files.
|
||||
model_type: String from the model config, e.g. "diffusion_cond".
|
||||
svd_bases_path: Optional path to a precomputed SVD bases .pt file
|
||||
(used only by -XS adapter types).
|
||||
|
||||
Returns:
|
||||
List of display names (file stems) in the same order as the input paths.
|
||||
"""
|
||||
if not lora_ckpt_paths:
|
||||
model.use_lora = False
|
||||
model.lora_names = []
|
||||
return []
|
||||
|
||||
is_cond = model_type in ("diffusion_cond", "diffusion_cond_inpaint")
|
||||
|
||||
# Pass 1: load each checkpoint and resolve adapter type.
|
||||
lora_entries = [] # (path, state_dict, config_dict, adapter_type)
|
||||
for i, lora_path in enumerate(lora_ckpt_paths):
|
||||
print(f"Loading LoRA {i} from {lora_path}")
|
||||
state_dict, config_dict = load_lora_checkpoint(lora_path)
|
||||
adapter_type_raw = config_dict.get("adapter_type", "lora")
|
||||
adapter_type = resolve_adapter_type(adapter_type_raw, state_dict)
|
||||
if adapter_type != adapter_type_raw:
|
||||
vprint(f"Resolved legacy '{adapter_type_raw}' -> {adapter_type}")
|
||||
lora_entries.append((lora_path, state_dict, config_dict, adapter_type))
|
||||
|
||||
# Load SVD bases only if at least one LoRA is -XS.
|
||||
model_svd_bases = None
|
||||
cond_svd_bases = None
|
||||
any_xs = any(t.endswith("-xs") for _, _, _, t in lora_entries)
|
||||
if any_xs:
|
||||
if svd_bases_path is not None:
|
||||
vprint(f"Loading SVD bases from {svd_bases_path}")
|
||||
svd_bases_all = torch.load(svd_bases_path, map_location="cpu", weights_only=True)
|
||||
model_svd_bases = {k[len("model."):]: v for k, v in svd_bases_all.items() if k.startswith("model.")}
|
||||
cond_svd_bases = {k[len("conditioner."):]: v for k, v in svd_bases_all.items() if k.startswith("conditioner.")}
|
||||
vprint(f" model bases: {len(model_svd_bases)}, conditioner bases: {len(cond_svd_bases)}")
|
||||
else:
|
||||
print("WARNING: -XS adapter type present without svd_bases_path -- SVD will be computed on current device")
|
||||
|
||||
# Pass 2: apply each LoRA.
|
||||
lora_names = []
|
||||
for i, (lora_path, state_dict, config_dict, adapter_type) in enumerate(lora_entries):
|
||||
rank = config_dict.get("rank", infer_global_rank(state_dict))
|
||||
alpha = config_dict.get("alpha", rank)
|
||||
include = config_dict.get("include", None)
|
||||
exclude = config_dict.get("exclude", None)
|
||||
is_xs = adapter_type.endswith("-xs")
|
||||
|
||||
lora_config = {
|
||||
torch.nn.Linear: {
|
||||
"weight": partial(LoRAParametrization.from_linear, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
|
||||
},
|
||||
torch.nn.Conv1d: {
|
||||
"weight": partial(LoRAParametrization.from_conv1d, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
|
||||
},
|
||||
}
|
||||
if is_cond:
|
||||
add_lora(model.model, lora_config, include=include, exclude=exclude,
|
||||
svd_bases=model_svd_bases if is_xs else None)
|
||||
add_lora(model.conditioner, lora_config, include=include, exclude=exclude,
|
||||
svd_bases=cond_svd_bases if is_xs else None)
|
||||
else:
|
||||
add_lora(model, lora_config, include=include, exclude=exclude,
|
||||
svd_bases=model_svd_bases if is_xs else None)
|
||||
|
||||
prepare_dora_state_dict(state_dict)
|
||||
remapped_sd = remap_lora_state_dict(state_dict, i)
|
||||
if is_cond:
|
||||
model.model.load_state_dict(remapped_sd, strict=False)
|
||||
model.conditioner.load_state_dict(remapped_sd, strict=False)
|
||||
else:
|
||||
model.load_state_dict(remapped_sd, strict=False)
|
||||
lora_names.append(os.path.splitext(os.path.basename(lora_path))[0])
|
||||
|
||||
vprint("lora layers:", len(get_lora_layers(model)))
|
||||
model.use_lora = True
|
||||
model.lora_names = lora_names
|
||||
return lora_names
|
||||
@@ -0,0 +1,471 @@
|
||||
"""
|
||||
References:
|
||||
1) the official LoRA implementation released by Microsoft:
|
||||
https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
|
||||
"""
|
||||
|
||||
import math
|
||||
import re
|
||||
from functools import partial
|
||||
from itertools import product
|
||||
|
||||
import torch
|
||||
import torch.nn.utils.parametrize as parametrize
|
||||
from torch import nn
|
||||
|
||||
from ...verbose import vprint
|
||||
|
||||
|
||||
def _canonicalize_svd_signs(U, Vh):
|
||||
"""Enforce deterministic sign convention: largest-magnitude element of each U column is positive."""
|
||||
max_abs_idx = U.abs().argmax(dim=0)
|
||||
signs = U[max_abs_idx, torch.arange(U.shape[1])].sign()
|
||||
signs[signs == 0] = 1
|
||||
return U * signs.unsqueeze(0), Vh * signs.unsqueeze(1)
|
||||
|
||||
|
||||
class LoRAParametrization(nn.Module):
|
||||
def __init__(self, fan_in, fan_out, fan_in_fan_out=False, rank=4, lora_dropout_p=0.0, lora_alpha=1,
|
||||
adapter_type="lora", W0=None, lora_index=0, svd_bases=None
|
||||
):
|
||||
super().__init__()
|
||||
dtype = torch.get_default_dtype()
|
||||
device = W0.device if W0 is not None else None
|
||||
# if weight is stored as (fan_out, fan_in), the memory layout of A & B follows (W + BA)x
|
||||
# otherwise, it's x(W + AB). This allows us to tie the weights between linear layers and embeddings
|
||||
self.swap = (lambda x: (x[1], x[0])) if fan_in_fan_out else (lambda x: x)
|
||||
self.lora_alpha, self.rank = lora_alpha, rank
|
||||
self.scaling = lora_alpha / rank
|
||||
self.lora_dropout = nn.Dropout(p=lora_dropout_p) if lora_dropout_p > 0 else lambda x: x
|
||||
self.dropout_fn = self._dropout if lora_dropout_p > 0 else lambda x: x
|
||||
self.register_buffer("lora_dropout_mask", torch.ones(self.swap((1, fan_in)), dtype=dtype, device=device))
|
||||
self.register_buffer("lora_strength", torch.tensor(1.0, dtype=dtype, device=device), persistent=False)
|
||||
self.adapter_type = adapter_type
|
||||
self.lora_index = lora_index
|
||||
|
||||
if self.adapter_type == "lora":
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
self.forward_fn = self.lora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type in ("dora-rows", "dora"):
|
||||
# DoRA-rows: magnitude per row (per-output-neuron, dim=1) - paper-correct variant
|
||||
# "dora" maps here as the default for backward compat
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
self._norm_dim = 1
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "dora-cols":
|
||||
# DoRA-cols: magnitude per column (per-input-feature, dim=0) - legacy variant
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
self._norm_dim = 0
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "bora":
|
||||
# BoRA: bi-dimensional DoRA - magnitude on both row and column axes
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
self.magnitude_r = nn.Parameter(W0_2d.norm(dim=1).to(dtype))
|
||||
self.magnitude_c = nn.Parameter(W0_2d.norm(dim=0).to(dtype))
|
||||
else:
|
||||
self.magnitude_r = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.magnitude_c = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.bora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "lora-xs":
|
||||
# frozen bases U:(fan_out,r), V:(fan_in,r); tiny trainable core M:(r,r)
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
elif W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1).cpu().float()
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d, full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.forward_fn = self.lora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "dora-rows-xs":
|
||||
# DoRA-rows + LoRA-XS: frozen SVD bases + trainable M_xs core + row magnitude
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
self._norm_dim = 1
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
if svd_bases is None:
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "dora-cols-xs":
|
||||
# DoRA-cols + LoRA-XS: frozen SVD bases + trainable M_xs core + column magnitude
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
self._norm_dim = 0
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
if svd_bases is None:
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "bora-xs":
|
||||
# BoRA + LoRA-XS: frozen SVD bases + trainable M_xs core + row & column magnitudes
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
if svd_bases is None:
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.magnitude_r = nn.Parameter(W0_2d.norm(dim=1).to(dtype))
|
||||
self.magnitude_c = nn.Parameter(W0_2d.norm(dim=0).to(dtype))
|
||||
else:
|
||||
self.magnitude_r = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.magnitude_c = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.bora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
|
||||
def lora_forward(self, W):
|
||||
delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A)))).view(W.shape)
|
||||
delta = self.scaling * self.lora_strength * delta
|
||||
return (W + delta.to(W.dtype))
|
||||
|
||||
def dora_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
# work in 2D so norm/magnitude ops don't broadcast badly for Conv1d/Conv2d
|
||||
# Upcast W to LoRA param dtype for stable norm computation, cast result back.
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.lora_A.dtype)
|
||||
# low-rank update on the *direction*
|
||||
delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A))))
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
# normalize along _norm_dim, then scale by per-element magnitude
|
||||
V_hat = V / (V.norm(dim=self._norm_dim, keepdim=True) + 1e-12)
|
||||
return (V_hat * self.magnitude.unsqueeze(self._norm_dim)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def bora_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.lora_A.dtype)
|
||||
delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A))))
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
# Row-normalize, then scale by row magnitudes
|
||||
V_r = V / (V.norm(dim=1, keepdim=True) + 1e-12)
|
||||
intermediate = self.magnitude_r.unsqueeze(1) * V_r
|
||||
# Column-normalize, then scale by column magnitudes
|
||||
H_c = intermediate / (intermediate.norm(dim=0, keepdim=True) + 1e-12)
|
||||
return (H_c * self.magnitude_c.unsqueeze(0)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def lora_xs_forward(self, W):
|
||||
delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
|
||||
delta = delta.view(W.shape)
|
||||
return (W + (self.scaling * self.lora_strength * delta).to(W.dtype))
|
||||
|
||||
def dora_xs_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.U.dtype)
|
||||
delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
V_hat = V / (V.norm(dim=self._norm_dim, keepdim=True) + 1e-12)
|
||||
return (V_hat * self.magnitude.unsqueeze(self._norm_dim)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def bora_xs_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.U.dtype)
|
||||
delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
V_r = V / (V.norm(dim=1, keepdim=True) + 1e-12)
|
||||
intermediate = self.magnitude_r.unsqueeze(1) * V_r
|
||||
H_c = intermediate / (intermediate.norm(dim=0, keepdim=True) + 1e-12)
|
||||
return (H_c * self.magnitude_c.unsqueeze(0)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def forward(self, X):
|
||||
return self.forward_fn(X)
|
||||
|
||||
def _dropout(self, A):
|
||||
# to mimic the original implementation: A @ dropout(x), we do (A * dropout(ones)) @ x
|
||||
return A * self.lora_dropout(self.lora_dropout_mask)
|
||||
|
||||
def disable_lora(self):
|
||||
self.forward_fn = lambda x: x
|
||||
|
||||
def enable_lora(self):
|
||||
self.forward_fn = self._adapter_forward_fn
|
||||
|
||||
@classmethod
|
||||
def _get_original_weight(cls, layer, attr_name="weight"):
|
||||
"""Get the true original weight, even when the layer is already parametrized."""
|
||||
if hasattr(layer, 'parametrizations') and hasattr(layer.parametrizations, attr_name):
|
||||
return getattr(layer.parametrizations, attr_name).original.detach()
|
||||
return getattr(layer, attr_name).detach()
|
||||
|
||||
@classmethod
|
||||
def from_linear(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_out, fan_in = layer.weight.shape
|
||||
W0 = cls._get_original_weight(layer)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_conv2d(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_out, fan_in = layer.weight.view(layer.weight.shape[0], -1).shape
|
||||
W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_conv1d(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_out, fan_in = layer.weight.view(layer.weight.shape[0], -1).shape
|
||||
W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_embedding(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_in, fan_out = layer.weight.shape
|
||||
W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=True, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
|
||||
default_lora_config = { # specify which layers to add lora to, by default only add to linear layers
|
||||
nn.Linear: {
|
||||
"weight": partial(LoRAParametrization.from_linear, rank=4),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# --- layer name filtering helpers ---
|
||||
|
||||
def _expand(s):
|
||||
"""Expand bracket notation in layer filter strings, e.g. 'layers[0-5]' -> ['layers0', ..., 'layers5']."""
|
||||
t = re.split(r'\[(\d+)-(\d+)\]', s)
|
||||
if len(t) == 1:
|
||||
return [s]
|
||||
lits, starts, ends = t[0::3], t[1::3], t[2::3]
|
||||
pools = []
|
||||
for a, b in zip(starts, ends):
|
||||
ai, bi = int(a), int(b)
|
||||
step = 1 if bi >= ai else -1
|
||||
pools.append([str(n) for n in range(ai, bi + step, step)])
|
||||
out = []
|
||||
for combo in product(*pools):
|
||||
pieces = []
|
||||
for i, lit in enumerate(lits):
|
||||
pieces.append(lit)
|
||||
if i < len(combo):
|
||||
pieces.append(combo[i])
|
||||
out.append(''.join(pieces))
|
||||
return out
|
||||
|
||||
|
||||
def _matches_any(name, patterns):
|
||||
"""Check if name contains any pattern substring (with bracket expansion)."""
|
||||
if not patterns:
|
||||
return False
|
||||
for pattern in patterns:
|
||||
for expanded in _expand(pattern.strip()):
|
||||
if expanded in name:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# --- core LoRA application ---
|
||||
|
||||
def _match_layer_type(layer, lora_config):
|
||||
"""Find the matching lora_config key for a layer, using isinstance to handle ParametrizedLinear etc."""
|
||||
for layer_type in lora_config:
|
||||
if isinstance(layer, layer_type):
|
||||
return layer_type
|
||||
return None
|
||||
|
||||
|
||||
def apply_lora(layer, register=True, merge=False, lora_config=default_lora_config):
|
||||
"""add lora parametrization to a layer, designed to be used with model.apply"""
|
||||
if register:
|
||||
matched_type = _match_layer_type(layer, lora_config)
|
||||
if matched_type is not None:
|
||||
for attr_name, parametrization in lora_config[matched_type].items():
|
||||
parametrize.register_parametrization(layer, attr_name, parametrization(layer), unsafe=True)
|
||||
else: # this will remove all parametrizations, use with caution
|
||||
if hasattr(layer, "parametrizations"):
|
||||
for attr_name in list(layer.parametrizations.keys()):
|
||||
parametrize.remove_parametrizations(layer, attr_name, leave_parametrized=merge)
|
||||
|
||||
|
||||
def add_lora(model, lora_config=default_lora_config, include=None, exclude=None, svd_bases=None):
|
||||
"""Add LoRA parametrization to layers in a model.
|
||||
|
||||
Args:
|
||||
model: The model to add LoRA to.
|
||||
lora_config: Dict mapping nn.Module types to parametrization configs.
|
||||
include: Optional list of substrings. If provided, only modules whose
|
||||
name contains at least one pattern get LoRA. Supports bracket
|
||||
expansion (e.g. "layers[0-11]").
|
||||
exclude: Optional list of substrings. Modules matching any pattern are
|
||||
skipped even if they match an include pattern.
|
||||
svd_bases: Optional dict of pre-computed SVD bases keyed by
|
||||
"layer_name.weight". When provided, -XS adapter types use
|
||||
these instead of computing SVD.
|
||||
"""
|
||||
if svd_bases is not None or include is not None or exclude is not None:
|
||||
# Named-module path: iterate by name so we can look up SVD bases per layer
|
||||
n_cached = 0
|
||||
n_recomputed = 0
|
||||
for name, module in model.named_modules():
|
||||
matched_type = _match_layer_type(module, lora_config)
|
||||
if matched_type is None:
|
||||
continue
|
||||
|
||||
if include is not None and not _matches_any(name, include):
|
||||
continue
|
||||
|
||||
if exclude is not None and _matches_any(name, exclude):
|
||||
continue
|
||||
|
||||
for attr_name, parametrization_fn in lora_config[matched_type].items():
|
||||
layer_bases = None
|
||||
if svd_bases is not None:
|
||||
bases_key = f"{name}.{attr_name}"
|
||||
layer_bases = svd_bases.get(bases_key)
|
||||
param = parametrization_fn(module, svd_bases=layer_bases)
|
||||
parametrize.register_parametrization(module, attr_name, param, unsafe=True)
|
||||
if hasattr(param, 'U'): # -XS layer that needs SVD bases
|
||||
if layer_bases is not None:
|
||||
n_cached += 1
|
||||
else:
|
||||
n_recomputed += 1
|
||||
if n_cached > 0:
|
||||
vprint(f"SVD bases: {n_cached} layers loaded from cache")
|
||||
if n_recomputed > 0:
|
||||
vprint(f"SVD bases: {n_recomputed} layers RECOMPUTED (missing from cache)")
|
||||
if svd_bases is not None and n_cached == 0 and n_recomputed == 0:
|
||||
vprint("SVD bases: provided but no -XS layers found")
|
||||
else:
|
||||
# Fast path: original behavior, no name inspection needed
|
||||
model.apply(partial(apply_lora, lora_config=lora_config))
|
||||
|
||||
def add_lora_by_name(model, target_module_names, lora_config=default_lora_config, svd_bases=None):
|
||||
"""Add LoRA to specific layers by name. Convenience wrapper around add_lora()."""
|
||||
add_lora(model, lora_config=lora_config, include=target_module_names, svd_bases=svd_bases)
|
||||
|
||||
|
||||
def merge_lora(model):
|
||||
"""merge lora parametrization to all layers in a model. This will remove all parametrization"""
|
||||
model.apply(partial(apply_lora, register=False, merge=True))
|
||||
|
||||
def remove_lora(model):
|
||||
"""remove lora parametrization to all layers in a model. This will remove all parametrization"""
|
||||
model.apply(partial(apply_lora, register=False, merge=False))
|
||||
|
||||
def remove_lora_by_index(model, lora_index):
|
||||
"""Remove all LoRA parametrizations with a specific lora_index.
|
||||
|
||||
For modules where other LoRA entries remain, only the matching entry is deleted.
|
||||
For modules where the removed entry was the only one, the parametrization is
|
||||
removed entirely and the weight reverts to its original value.
|
||||
"""
|
||||
for name, mod in model.named_modules():
|
||||
if not hasattr(mod, 'parametrizations'):
|
||||
continue
|
||||
for attr in list(mod.parametrizations.keys()):
|
||||
params = mod.parametrizations[attr]
|
||||
to_remove = [i for i, p in enumerate(params)
|
||||
if isinstance(p, LoRAParametrization) and p.lora_index == lora_index]
|
||||
if not to_remove:
|
||||
continue
|
||||
if len(params) == len(to_remove):
|
||||
# All entries being removed — remove parametrization entirely
|
||||
parametrize.remove_parametrizations(mod, attr, leave_parametrized=False)
|
||||
else:
|
||||
for i in reversed(to_remove):
|
||||
del params[i]
|
||||
|
||||
def set_lora_strength(model, strength: float, lora_index=None):
|
||||
"""Set lora strength. If lora_index is None, sets all LoRAs. If specified, sets only that index."""
|
||||
strength = float(strength)
|
||||
if lora_index is None:
|
||||
for name, buf in model.named_buffers(recurse=True):
|
||||
if name.endswith("lora_strength"):
|
||||
buf.fill_(strength)
|
||||
else:
|
||||
for p in _iter_lora_params(model):
|
||||
if p.lora_index == lora_index:
|
||||
p.lora_strength.fill_(strength)
|
||||
|
||||
def _iter_lora_params(model):
|
||||
for _, mod in model.named_modules():
|
||||
plist = getattr(getattr(mod, "parametrizations", None), "weight", None)
|
||||
if plist is None:
|
||||
continue
|
||||
for p in plist:
|
||||
if isinstance(p, LoRAParametrization):
|
||||
yield p
|
||||
@@ -0,0 +1,420 @@
|
||||
import json
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file as _st_save_file
|
||||
from torch import nn
|
||||
|
||||
from .model import LoRAParametrization, _expand
|
||||
from ...verbose import vprint
|
||||
|
||||
def apply_to_lora(fn):
|
||||
"""apply a function to LoRAParametrization layers, designed to be used with model.apply"""
|
||||
|
||||
def apply_fn(layer):
|
||||
if isinstance(layer, LoRAParametrization):
|
||||
fn(layer)
|
||||
|
||||
return apply_fn
|
||||
|
||||
def enable_lora(model, lora_index=None):
|
||||
"""Enable LoRA layers. If lora_index is None, enables all. If specified, enables only that index."""
|
||||
def _enable(layer):
|
||||
if isinstance(layer, LoRAParametrization):
|
||||
if lora_index is None or layer.lora_index == lora_index:
|
||||
layer.enable_lora()
|
||||
model.apply(_enable)
|
||||
|
||||
def disable_lora(model, lora_index=None):
|
||||
"""Disable LoRA layers. If lora_index is None, disables all. If specified, disables only that index."""
|
||||
def _disable(layer):
|
||||
if isinstance(layer, LoRAParametrization):
|
||||
if lora_index is None or layer.lora_index == lora_index:
|
||||
layer.disable_lora()
|
||||
model.apply(_disable)
|
||||
|
||||
def get_lora_layers(model):
|
||||
layers = []
|
||||
for name, m in model.named_modules():
|
||||
plist = getattr(getattr(m, "parametrizations", None), "weight", None)
|
||||
if plist is None:
|
||||
continue
|
||||
for p in plist:
|
||||
if isinstance(p, LoRAParametrization):
|
||||
layers.append((name, p))
|
||||
return layers
|
||||
|
||||
def has_lora(model):
|
||||
"""Return True if the model has at least one LoRAParametrization on any weight."""
|
||||
return len(get_lora_layers(model))>=1
|
||||
|
||||
def filter_lora_layers(model, lora_layer_filter, lora_index=None):
|
||||
"""Enable/disable LoRA layers by name filter. If lora_index is specified, only affects that index."""
|
||||
lora_layer_filter = (lora_layer_filter or "").strip()
|
||||
# lora_layer_filter: logical OR on comma separated substrings
|
||||
# e.g. ".transformer.layers, .to_global_embed."
|
||||
def is_filtered(layer_name):
|
||||
filters = [x.lower().strip() for x in lora_layer_filter.split(",") if len(x)]
|
||||
for f in filters:
|
||||
for _f in _expand(f):
|
||||
if _f in layer_name:
|
||||
return True
|
||||
return False
|
||||
for name, p in get_lora_layers(model):
|
||||
if lora_index is not None and p.lora_index != lora_index:
|
||||
continue
|
||||
if is_filtered(name):
|
||||
p.disable_lora()
|
||||
else:
|
||||
p.enable_lora()
|
||||
|
||||
|
||||
|
||||
# ------------------- helper function for collecting parameters for training/saving -------------------
|
||||
|
||||
|
||||
def name_is_lora(name):
|
||||
return (
|
||||
len(name.split(".")) >= 4
|
||||
and (name.split(".")[-4]) == "parametrizations"
|
||||
and name.split(".")[-1] in ["lora_A", "lora_B", "M_xs", "magnitude", "magnitude_r", "magnitude_c"]
|
||||
)
|
||||
|
||||
|
||||
def name_is_bias(name):
|
||||
return name.split(".")[-1] == "bias"
|
||||
|
||||
|
||||
def get_params_by_name(model, print_shapes=False, name_filter=None):
|
||||
for n, p in model.named_parameters():
|
||||
if name_filter is None or name_filter(n):
|
||||
if print_shapes:
|
||||
print(n, p.shape)
|
||||
yield p
|
||||
|
||||
|
||||
def get_lora_params(model, print_shapes=False):
|
||||
return get_params_by_name(model, print_shapes=print_shapes, name_filter=name_is_lora)
|
||||
|
||||
|
||||
def get_bias_params(model, print_shapes=False):
|
||||
return get_params_by_name(model, print_shapes=print_shapes, name_filter=name_is_bias)
|
||||
|
||||
|
||||
def get_lora_state_dict(model):
|
||||
return {k: v for k, v in model.state_dict().items() if name_is_lora(k)}
|
||||
|
||||
|
||||
def cast_base_to_precision(model, precision):
|
||||
"""Cast frozen base weights to lower precision, keeping LoRA params in fp32.
|
||||
|
||||
Args:
|
||||
model: Model with LoRA parametrizations applied.
|
||||
precision: One of "bf16", "bfloat16", "fp16", "float16".
|
||||
"""
|
||||
import torch
|
||||
|
||||
dtype_map = {"bf16": torch.bfloat16, "bfloat16": torch.bfloat16,
|
||||
"fp16": torch.float16, "float16": torch.float16}
|
||||
target_dtype = dtype_map.get(precision)
|
||||
if target_dtype is None:
|
||||
print(f"Unknown base_precision={precision!r}, keeping fp32")
|
||||
return
|
||||
vprint(f"Casting frozen base weights to {target_dtype}")
|
||||
model.to(target_dtype)
|
||||
# Restore LoRA params to fp32 — optimizer needs full precision
|
||||
n_restored = 0
|
||||
for p in get_lora_params(model):
|
||||
p.data = p.data.to(torch.float32)
|
||||
n_restored += 1
|
||||
vprint(f" Restored {n_restored} LoRA params to fp32")
|
||||
|
||||
|
||||
# ------------------- helper function for loading LoRa at inference -------------------
|
||||
|
||||
def infer_global_rank(lora_sd: dict) -> int:
|
||||
candidates = []
|
||||
for k, v in lora_sd.items():
|
||||
if k.endswith(".lora_A"):
|
||||
candidates.append(v.shape[0])
|
||||
elif k.endswith(".lora_B"):
|
||||
candidates.append(v.shape[1])
|
||||
elif k.endswith(".M_xs"):
|
||||
candidates.append(v.shape[0])
|
||||
if not candidates:
|
||||
raise ValueError("No LoRA/LoRA-XS tensors found to infer rank")
|
||||
r = candidates[0]
|
||||
# sanity check they all match (most checkpoints use one global r)
|
||||
if any(c != r for c in candidates):
|
||||
print(f"Multiple ranks detected in ckpt: {sorted(set(candidates))}. "
|
||||
"You’ll need per-layer ranks.")
|
||||
return r
|
||||
|
||||
|
||||
def detect_dora_variant(state_dict):
|
||||
"""Detect DoRA variant from a state dict.
|
||||
|
||||
Returns "dora-cols" (dim=0, per-input-feature) or "dora-rows" (dim=1,
|
||||
per-output-neuron) based on the shape of the first .magnitude tensor
|
||||
found. Returns None if no magnitude tensors exist or axis is ambiguous.
|
||||
"""
|
||||
for key, val in state_dict.items():
|
||||
if not key.endswith(".magnitude"):
|
||||
continue
|
||||
if val.dim() == 2:
|
||||
if val.shape[0] == 1:
|
||||
return "dora-cols" # old format: (1, fan_in)
|
||||
elif val.shape[1] == 1:
|
||||
return "dora-rows" # (fan_out, 1)
|
||||
# 1D magnitude — can't determine axis without model context
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def prepare_dora_state_dict(state_dict):
|
||||
"""Squeeze 2D DoRA magnitude tensors to 1D for loading. Modifies in-place."""
|
||||
for key in list(state_dict.keys()):
|
||||
if key.endswith(".magnitude") and state_dict[key].dim() == 2:
|
||||
state_dict[key] = state_dict[key].squeeze()
|
||||
|
||||
|
||||
def resolve_adapter_type(adapter_type, state_dict=None):
|
||||
"""Resolve legacy "dora" adapter_type to the correct variant.
|
||||
|
||||
If adapter_type is "dora" and a checkpoint state_dict is provided,
|
||||
attempts to detect dora-rows vs dora-cols from magnitude tensor shapes.
|
||||
Falls back to "dora-rows" (paper-correct default).
|
||||
All other adapter types are returned unchanged.
|
||||
"""
|
||||
if adapter_type != "dora":
|
||||
return adapter_type
|
||||
if state_dict is not None:
|
||||
detected = detect_dora_variant(state_dict)
|
||||
if detected:
|
||||
return detected
|
||||
return "dora-rows"
|
||||
|
||||
|
||||
# ------------------- helper functions for safetensors LoRA checkpoints -------------------
|
||||
|
||||
|
||||
def save_lora_safetensors(state_dict, lora_config, path):
|
||||
"""Save a LoRA checkpoint in safetensors format with config as metadata.
|
||||
|
||||
The lora_config dict is JSON-serialized and stored under the metadata key
|
||||
``"lora_config"``, which safetensors stores in the file header as a string.
|
||||
|
||||
Args:
|
||||
state_dict: Dict of LoRA tensors (from get_lora_state_dict).
|
||||
lora_config: Dict with keys like rank, alpha, adapter_type, include, exclude.
|
||||
path: Output file path (should end in .safetensors).
|
||||
"""
|
||||
metadata = {"lora_config": json.dumps(lora_config)}
|
||||
fp16_dict = {k: v.half() if v.is_floating_point() else v for k, v in state_dict.items()}
|
||||
_st_save_file(fp16_dict, str(path), metadata=metadata)
|
||||
|
||||
|
||||
def load_lora_checkpoint(path):
|
||||
"""Load a LoRA checkpoint from either .safetensors or .ckpt format.
|
||||
|
||||
Returns:
|
||||
Tuple of (state_dict, lora_config) where lora_config may be an empty
|
||||
dict if the checkpoint has no embedded config.
|
||||
"""
|
||||
path = str(path)
|
||||
if path.endswith(".safetensors"):
|
||||
with safe_open(path, framework="pt", device="cpu") as f:
|
||||
state_dict = {key: f.get_tensor(key) for key in f.keys()}
|
||||
metadata = f.metadata()
|
||||
lora_config = {}
|
||||
if metadata and "lora_config" in metadata:
|
||||
lora_config = json.loads(metadata["lora_config"])
|
||||
return state_dict, lora_config
|
||||
else:
|
||||
ckpt = torch.load(path, map_location="cpu", weights_only=False)
|
||||
state_dict = ckpt["state_dict"]
|
||||
lora_config = ckpt.get("lora_config", {})
|
||||
return state_dict, lora_config
|
||||
|
||||
|
||||
def convert_lora_ckpt_to_safetensors(ckpt_path, output_path=None):
|
||||
"""Convert a .ckpt LoRA checkpoint to .safetensors format.
|
||||
|
||||
Args:
|
||||
ckpt_path: Path to the input .ckpt file.
|
||||
output_path: Path for the output .safetensors file.
|
||||
If None, replaces .ckpt extension with .safetensors.
|
||||
|
||||
Returns:
|
||||
The output path.
|
||||
"""
|
||||
ckpt_path = str(ckpt_path)
|
||||
if output_path is None:
|
||||
if ckpt_path.endswith(".ckpt"):
|
||||
output_path = ckpt_path[:-5] + ".safetensors"
|
||||
else:
|
||||
output_path = ckpt_path + ".safetensors"
|
||||
state_dict, lora_config = load_lora_checkpoint(ckpt_path)
|
||||
save_lora_safetensors(state_dict, lora_config, output_path)
|
||||
return output_path
|
||||
|
||||
|
||||
# ------------------- helper functions for multi-LoRA stacking -------------------
|
||||
|
||||
|
||||
def get_lora_count(model):
|
||||
"""Return the number of stacked LoRAs (distinct lora_index values)."""
|
||||
indices = set()
|
||||
for _, p in get_lora_layers(model):
|
||||
indices.add(p.lora_index)
|
||||
return len(indices)
|
||||
|
||||
|
||||
def remap_lora_state_dict(state_dict, target_index):
|
||||
"""Remap LoRA state dict keys from parametrization index 0 to target_index.
|
||||
|
||||
Checkpoints always save keys with index 0 (e.g. *.parametrizations.weight.0.lora_A).
|
||||
When loading the Nth stacked LoRA, we need to remap to index N.
|
||||
"""
|
||||
if target_index == 0:
|
||||
return state_dict
|
||||
remapped = {}
|
||||
for k, v in state_dict.items():
|
||||
new_key = k.replace(
|
||||
".parametrizations.weight.0.",
|
||||
".parametrizations.weight.{}.".format(target_index)
|
||||
)
|
||||
remapped[new_key] = v
|
||||
return remapped
|
||||
|
||||
|
||||
# ------------------- helper function for inferencing with multiple lora -------------------
|
||||
|
||||
|
||||
def _get_adapter_param_names(lora_layer):
|
||||
"""Return the list of trainable parameter names for this adapter type."""
|
||||
if lora_layer.adapter_type == "lora":
|
||||
return ["lora_A", "lora_B"]
|
||||
elif lora_layer.adapter_type in ("dora", "dora-rows", "dora-cols"):
|
||||
return ["lora_A", "lora_B", "magnitude"]
|
||||
elif lora_layer.adapter_type == "bora":
|
||||
return ["lora_A", "lora_B", "magnitude_r", "magnitude_c"]
|
||||
elif lora_layer.adapter_type == "lora-xs":
|
||||
return ["M_xs"]
|
||||
elif lora_layer.adapter_type in ("dora-rows-xs", "dora-cols-xs"):
|
||||
return ["M_xs", "magnitude"]
|
||||
elif lora_layer.adapter_type == "bora-xs":
|
||||
return ["M_xs", "magnitude_r", "magnitude_c"]
|
||||
return []
|
||||
|
||||
|
||||
def _prepare_for_multiple_lora(lora_layer):
|
||||
lora_layer._multi_lora_slots = []
|
||||
|
||||
|
||||
def _append_lora(lora_layer):
|
||||
slot = {}
|
||||
for name in _get_adapter_param_names(lora_layer):
|
||||
slot[name] = nn.Parameter(getattr(lora_layer, name).clone())
|
||||
lora_layer._multi_lora_slots.append(slot)
|
||||
|
||||
|
||||
def load_multiple_lora(model, lora_state_dicts):
|
||||
model.apply(apply_to_lora(_prepare_for_multiple_lora))
|
||||
for state_dict in lora_state_dicts:
|
||||
_ = model.load_state_dict(state_dict, strict=False)
|
||||
model.apply(apply_to_lora(_append_lora))
|
||||
return model
|
||||
|
||||
|
||||
def _select_lora(lora_layer, index):
|
||||
slot = lora_layer._multi_lora_slots[index]
|
||||
for name, param in slot.items():
|
||||
setattr(lora_layer, name, param)
|
||||
|
||||
|
||||
def select_lora(model, index):
|
||||
model.apply(apply_to_lora(lambda x: _select_lora(x, index)))
|
||||
return model
|
||||
|
||||
# ------------------- helper function for merging multiple LoRAs -------------------
|
||||
|
||||
def merge_loras_into_base_model(model, lora_configurations):
|
||||
"""
|
||||
Merges multiple LoRAs with specified application weights directly into the base model's weights.
|
||||
Works with all adapter types (LoRA, DoRA, LoRA-XS) by using the parametrization's own forward.
|
||||
After merging, the LoRA parametrization on each layer is disabled.
|
||||
|
||||
Args:
|
||||
model: The PyTorch model with LoRAParametrization already applied (via add_lora).
|
||||
lora_configurations: A list of dicts, each specifying a LoRA to merge:
|
||||
[
|
||||
{
|
||||
'name': 'style1',
|
||||
'state_dict': loaded_torch_state_dict,
|
||||
'application_weight': 0.7
|
||||
},
|
||||
...
|
||||
]
|
||||
"""
|
||||
import torch
|
||||
|
||||
if not lora_configurations:
|
||||
return
|
||||
|
||||
# Collect all parametrized layers with their LoRA instances
|
||||
parametrized_layers = []
|
||||
for name, module in model.named_modules():
|
||||
if not hasattr(module, "parametrizations"):
|
||||
continue
|
||||
for attr_name in list(module.parametrizations.keys()):
|
||||
plist = module.parametrizations[attr_name]
|
||||
for p in plist:
|
||||
if isinstance(p, LoRAParametrization):
|
||||
parametrized_layers.append((name, module, attr_name, p, plist.original))
|
||||
break
|
||||
|
||||
# Initialize per-layer delta accumulators
|
||||
deltas = {id(orig): torch.zeros_like(orig.data) for _, _, _, _, orig in parametrized_layers}
|
||||
|
||||
for lora_cfg in lora_configurations:
|
||||
lora_sd = lora_cfg['state_dict']
|
||||
app_weight = lora_cfg['application_weight']
|
||||
|
||||
# Load this LoRA's adapter weights into the model's parametrizations
|
||||
model.load_state_dict(lora_sd, strict=False)
|
||||
|
||||
# Compute each layer's delta using the parametrization's own forward
|
||||
with torch.no_grad():
|
||||
for name, module, attr_name, lora_p, original in parametrized_layers:
|
||||
merged = lora_p(original)
|
||||
deltas[id(original)] += app_weight * (merged - original)
|
||||
|
||||
# Apply accumulated deltas to original weights and disable LoRA
|
||||
with torch.no_grad():
|
||||
for name, module, attr_name, lora_p, original in parametrized_layers:
|
||||
original.data.add_(deltas[id(original)])
|
||||
lora_p.disable_lora()
|
||||
|
||||
# ------------------- helper function for tying and untieing weights -------------------
|
||||
|
||||
|
||||
def tie_weights(linear: nn.Linear, embedding: nn.Embedding):
|
||||
"""tie the weights of the linear layer and the embedding layer both with the same lora"""
|
||||
lora_p = embedding.parametrizations.weight[0]
|
||||
if lora_p.adapter_type != "lora":
|
||||
raise NotImplementedError(f"tie_weights only supports standard LoRA, got adapter_type='{lora_p.adapter_type}'")
|
||||
# this line below is optional if the original is already tied
|
||||
embedding.parametrizations.weight.original = linear.parametrizations.weight.original
|
||||
embedding.parametrizations.weight[0].lora_A = linear.parametrizations.weight[0].lora_B
|
||||
embedding.parametrizations.weight[0].lora_B = linear.parametrizations.weight[0].lora_A
|
||||
|
||||
|
||||
def untie_weights(linear: nn.Linear, embedding: nn.Embedding):
|
||||
"""untie the weights of the linear layer and the embedding layer"""
|
||||
lora_p = embedding.parametrizations.weight[0]
|
||||
if lora_p.adapter_type != "lora":
|
||||
raise NotImplementedError(f"untie_weights only supports standard LoRA, got adapter_type='{lora_p.adapter_type}'")
|
||||
embedding.parametrizations.weight.original = nn.Parameter(embedding.weight.original.clone())
|
||||
embedding.parametrizations.weight[0].lora_A = nn.Parameter(embedding.parametrizations.weight[0].lora_A.clone())
|
||||
embedding.parametrizations.weight[0].lora_B = nn.Parameter(embedding.parametrizations.weight[0].lora_B.clone())
|
||||
@@ -0,0 +1,84 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from torchaudio.transforms import Resample
|
||||
from .blocks import ResidualUnit, WNConv1d
|
||||
|
||||
|
||||
class AutoencoderPretransform(nn.Module):
|
||||
def __init__(self, model, scale=1.0, iterate_batch=False, chunked=False):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.model.requires_grad_(False).eval()
|
||||
self.scale = scale
|
||||
self.downsampling_ratio = model.downsampling_ratio
|
||||
self.io_channels = model.io_channels
|
||||
self.enable_grad = False
|
||||
self.iterate_batch = iterate_batch
|
||||
self.chunked = chunked
|
||||
|
||||
def encode(self, x, **kwargs):
|
||||
return self.model.encode_audio(x, chunked=self.chunked, iterate_batch=self.iterate_batch, **kwargs) / self.scale
|
||||
|
||||
def decode(self, z, chunked=None, **kwargs):
|
||||
chunked = self.chunked if chunked is None else chunked
|
||||
return self.model.decode_audio(z * self.scale, chunked=chunked, iterate_batch=self.iterate_batch, **kwargs)
|
||||
|
||||
def fold_channels_into_batch(x):
|
||||
x = rearrange(x, 'b c ... -> (b c) ...')
|
||||
return x
|
||||
|
||||
def unfold_channels_from_batch(x, channels):
|
||||
if channels == 1:
|
||||
return x.unsqueeze(1)
|
||||
x = rearrange(x, '(b c) ... -> b c ...', c = channels)
|
||||
return x
|
||||
|
||||
|
||||
class PatchedPretransform(nn.Module):
|
||||
def __init__(self, channels, patch_size, oversampling = 1, postfilter_channels = 0, **kwargs):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.patch_size = patch_size
|
||||
self.oversampling = oversampling
|
||||
|
||||
self.downsampling_ratio = patch_size
|
||||
self.io_channels = channels
|
||||
self.encoded_channels = channels * patch_size
|
||||
self.enable_grad = False
|
||||
|
||||
if self.oversampling > 1:
|
||||
self.input_upsampler = Resample(1, self.oversampling)
|
||||
self.output_downsampler = Resample(self.oversampling, 1)
|
||||
|
||||
if postfilter_channels > 0:
|
||||
self.postfilter = nn.Sequential(
|
||||
WNConv1d(in_channels=channels, out_channels=postfilter_channels, kernel_size=7, padding=3, bias=True),
|
||||
ResidualUnit(in_channels=postfilter_channels, out_channels=postfilter_channels,
|
||||
dilation=1, use_snake=True, bias = True),
|
||||
ResidualUnit(in_channels=postfilter_channels, out_channels=postfilter_channels,
|
||||
dilation=3, use_snake=True, bias = True),
|
||||
ResidualUnit(in_channels=postfilter_channels, out_channels=postfilter_channels,
|
||||
dilation=9, use_snake=True, bias = True),
|
||||
WNConv1d(in_channels=postfilter_channels, out_channels=channels, kernel_size=7, padding=3, bias=False))
|
||||
|
||||
def _pad(self, x):
|
||||
seq_len = x.shape[-1]
|
||||
pad_len = (self.patch_size - (seq_len % self.patch_size)) % self.patch_size
|
||||
if pad_len > 0:
|
||||
x = torch.cat([x, torch.zeros_like(x[:, :, :pad_len])], dim=-1)
|
||||
return x
|
||||
|
||||
def encode(self, x):
|
||||
if self.oversampling > 1:
|
||||
x = self.input_upsampler(x)
|
||||
x = self._pad(x)
|
||||
x = rearrange(x, "b c (l h) -> b (c h) l", h=self.patch_size)
|
||||
return x
|
||||
def decode(self, x):
|
||||
x = rearrange(x, "b (c h) l -> b c (l h)", h=self.patch_size)
|
||||
if hasattr(self, 'postfilter'):
|
||||
x = self.postfilter(x)
|
||||
if self.oversampling > 1:
|
||||
x = self.output_downsampler(x)
|
||||
return x
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,19 @@
|
||||
import torch
|
||||
import os
|
||||
|
||||
torch._dynamo.config.cache_size_limit = max(64, torch._dynamo.config.cache_size_limit)
|
||||
torch._dynamo.config.suppress_errors = True
|
||||
|
||||
# Get torch.compile flag from environment variable ENABLE_TORCH_COMPILE
|
||||
|
||||
enable_torch_compile = os.environ.get("ENABLE_TORCH_COMPILE", "0") == "1"
|
||||
|
||||
def compile(function, *args, **kwargs):
|
||||
|
||||
if enable_torch_compile:
|
||||
try:
|
||||
return torch.compile(function, *args, **kwargs)
|
||||
except RuntimeError:
|
||||
return function
|
||||
|
||||
return function
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user