diff --git a/docs/asset-pipeline.md b/docs/asset-pipeline.md index 5972c944..7564107d 100644 --- a/docs/asset-pipeline.md +++ b/docs/asset-pipeline.md @@ -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 \ diff --git a/docs/audio/asset-manifest.md b/docs/audio/asset-manifest.md index 79332275..7ffb2e7d 100644 --- a/docs/audio/asset-manifest.md +++ b/docs/audio/asset-manifest.md @@ -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） diff --git a/docs/design/15-radio.md b/docs/design/15-radio.md index a07d0697..e4459452 100644 --- a/docs/design/15-radio.md +++ b/docs/design/15-radio.md @@ -442,6 +442,11 @@ ENDING: e*_a → e*_b（按 rep×knot 四选一）→ ep_lead → ep_×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_×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 收尾时需协同） diff --git a/docs/setup.md b/docs/setup.md index 040a7e8a..216db9bd 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -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。 +
旧记录（2026-08-29，已被上条实测推翻） + +镜像曾把带 token 的请求 308 重定向回 huggingface.co，本机直连官方又不可达时会直接失败， +因此当时 `env.d/sfx.sh` 故意不设 `HF_ENDPOINT`。Qwen3-TTS（公开）与 Stable Audio 3 现均走镜像。 +
## 8. Mureka BGM API（可选，需付费 key） diff --git a/env.d/sfx.md b/env.d/sfx.md index c4242ffa..4416568a 100644 --- a/env.d/sfx.md +++ b/env.d/sfx.md @@ -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） diff --git a/research/new8/batch1-assets-report.md b/research/new8/batch1-assets-report.md index 1116ab2d..68c1df02 100644 --- a/research/new8/batch1-assets-report.md +++ b/research/new8/batch1-assets-report.md @@ -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//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//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 服务。 diff --git a/research/new8/radio/round2/live-player.md b/research/new8/radio/round2/live-player.md new file mode 100644 index 00000000..82e263ff --- /dev/null +++ b/research/new8/radio/round2/live-player.md @@ -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 处） diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/atk_hit.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/atk_hit.m4a new file mode 100644 index 00000000..de251231 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/atk_hit.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/atk_hit.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/atk_hit.ogg new file mode 100644 index 00000000..b539b58b Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/atk_hit.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_1.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_1.m4a new file mode 100644 index 00000000..57d7e9c1 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_1.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_1.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_1.ogg new file mode 100644 index 00000000..f08b32c2 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_1.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_2.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_2.m4a new file mode 100644 index 00000000..356307fc Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_2.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_2.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_2.ogg new file mode 100644 index 00000000..451e9133 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_2.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_3.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_3.m4a new file mode 100644 index 00000000..e8618e02 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_3.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_3.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_3.ogg new file mode 100644 index 00000000..5e73fdf7 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_3.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_4.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_4.m4a new file mode 100644 index 00000000..7b7eebaa Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_4.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_4.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_4.ogg new file mode 100644 index 00000000..5c10ff1a Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/echo_4.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/enemy_die.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/enemy_die.m4a new file mode 100644 index 00000000..e845f9a6 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/enemy_die.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/enemy_die.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/enemy_die.ogg new file mode 100644 index 00000000..9bd0c41b Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/enemy_die.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug.m4a new file mode 100644 index 00000000..8f62d422 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug.ogg new file mode 100644 index 00000000..b015e115 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_c.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_c.m4a new file mode 100644 index 00000000..5b1b4d7b Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_c.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_c.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_c.ogg new file mode 100644 index 00000000..0948cfd4 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_c.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_l.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_l.m4a new file mode 100644 index 00000000..7cc36e9b Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_l.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_l.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_l.ogg new file mode 100644 index 00000000..27a65094 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_l.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_r.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_r.m4a new file mode 100644 index 00000000..bc0bdd7c Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_r.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_r.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_r.ogg new file mode 100644 index 00000000..de1dc63a Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_bug_r.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy.m4a new file mode 100644 index 00000000..b3be4d00 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy.ogg new file mode 100644 index 00000000..92666773 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_c.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_c.m4a new file mode 100644 index 00000000..b960688b Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_c.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_c.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_c.ogg new file mode 100644 index 00000000..cc184115 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_c.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_l.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_l.m4a new file mode 100644 index 00000000..a7331692 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_l.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_l.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_l.ogg new file mode 100644 index 00000000..715c4463 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_l.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_r.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_r.m4a new file mode 100644 index 00000000..b4be3597 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_r.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_r.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_r.ogg new file mode 100644 index 00000000..084e1f89 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_heavy_r.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small.m4a new file mode 100644 index 00000000..acb1a162 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small.ogg new file mode 100644 index 00000000..8e63ad0b Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_c.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_c.m4a new file mode 100644 index 00000000..702609e0 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_c.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_c.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_c.ogg new file mode 100644 index 00000000..ee27d04a Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_c.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_l.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_l.m4a new file mode 100644 index 00000000..5777ee6e Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_l.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_l.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_l.ogg new file mode 100644 index 00000000..0ad9989c Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_l.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_r.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_r.m4a new file mode 100644 index 00000000..4395ef41 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_r.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_r.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_r.ogg new file mode 100644 index 00000000..1afd1878 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/mob_step_small_r.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/stairs_down.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/stairs_down.m4a new file mode 100644 index 00000000..94018ecf Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/stairs_down.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/stairs_down.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/stairs_down.ogg new file mode 100644 index 00000000..3f56a7eb Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/stairs_down.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug.m4a new file mode 100644 index 00000000..ee927960 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug.ogg new file mode 100644 index 00000000..7c285338 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_l.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_l.m4a new file mode 100644 index 00000000..176052df Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_l.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_l.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_l.ogg new file mode 100644 index 00000000..7d900219 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_l.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_r.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_r.m4a new file mode 100644 index 00000000..0ee84772 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_r.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_r.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_r.ogg new file mode 100644 index 00000000..281e2a4d Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_bug_r.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy.m4a new file mode 100644 index 00000000..f5731d76 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy.ogg new file mode 100644 index 00000000..87e5703f Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_l.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_l.m4a new file mode 100644 index 00000000..a24604e8 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_l.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_l.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_l.ogg new file mode 100644 index 00000000..4138371e Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_l.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_r.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_r.m4a new file mode 100644 index 00000000..55324487 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_r.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_r.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_r.ogg new file mode 100644 index 00000000..622bc4b3 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_heavy_r.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small.m4a new file mode 100644 index 00000000..f213b4ba Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small.ogg new file mode 100644 index 00000000..127eb961 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_l.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_l.m4a new file mode 100644 index 00000000..b9bbf298 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_l.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_l.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_l.ogg new file mode 100644 index 00000000..0d500406 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_l.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_r.m4a b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_r.m4a new file mode 100644 index 00000000..8ad706e8 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_r.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_r.ogg b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_r.ogg new file mode 100644 index 00000000..0a073b0f Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/dungeon/tel_small_r.ogg differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/radio/phone_ring.m4a b/tools/asset_fix_backup/sa3_ab_20260902/radio/phone_ring.m4a new file mode 100644 index 00000000..1ee12c8a Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/radio/phone_ring.m4a differ diff --git a/tools/asset_fix_backup/sa3_ab_20260902/radio/phone_ring.ogg b/tools/asset_fix_backup/sa3_ab_20260902/radio/phone_ring.ogg new file mode 100644 index 00000000..219fb484 Binary files /dev/null and b/tools/asset_fix_backup/sa3_ab_20260902/radio/phone_ring.ogg differ diff --git a/tools/stable-audio-3/.gitignore b/tools/stable-audio-3/.gitignore index bfc8ae19..e3808244 100644 --- a/tools/stable-audio-3/.gitignore +++ b/tools/stable-audio-3/.gitignore @@ -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/** diff --git a/tools/stable-audio-3/stable_audio_3/models/__init__.py b/tools/stable-audio-3/stable_audio_3/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/__init__.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 00000000..0d42c04e Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/__init__.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/autoencoders.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/autoencoders.cpython-310.pyc new file mode 100644 index 00000000..53944669 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/autoencoders.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/blocks.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/blocks.cpython-310.pyc new file mode 100644 index 00000000..cfa3bf7b Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/blocks.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/bottleneck.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/bottleneck.cpython-310.pyc new file mode 100644 index 00000000..e7259585 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/bottleneck.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/conditioners.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/conditioners.cpython-310.pyc new file mode 100644 index 00000000..5fbe4b98 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/conditioners.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/diffusion.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/diffusion.cpython-310.pyc new file mode 100644 index 00000000..30c30f46 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/diffusion.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/dit.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/dit.cpython-310.pyc new file mode 100644 index 00000000..7899c951 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/dit.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/pretransforms.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/pretransforms.cpython-310.pyc new file mode 100644 index 00000000..9836dec8 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/pretransforms.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/transformer.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 00000000..d19bd560 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/transformer.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/__pycache__/utils.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/__pycache__/utils.cpython-310.pyc new file mode 100644 index 00000000..71b7c25f Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/__pycache__/utils.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/autoencoders.py b/tools/stable-audio-3/stable_audio_3/models/autoencoders.py new file mode 100644 index 00000000..b822af8c --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/autoencoders.py @@ -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 + \ No newline at end of file diff --git a/tools/stable-audio-3/stable_audio_3/models/blocks.py b/tools/stable-audio-3/stable_audio_3/models/blocks.py new file mode 100644 index 00000000..32b64411 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/blocks.py @@ -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) + + diff --git a/tools/stable-audio-3/stable_audio_3/models/bottleneck.py b/tools/stable-audio-3/stable_audio_3/models/bottleneck.py new file mode 100644 index 00000000..3aa068c7 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/bottleneck.py @@ -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 diff --git a/tools/stable-audio-3/stable_audio_3/models/conditioners.py b/tools/stable-audio-3/stable_audio_3/models/conditioners.py new file mode 100644 index 00000000..c501e293 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/conditioners.py @@ -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 + diff --git a/tools/stable-audio-3/stable_audio_3/models/diffusion.py b/tools/stable-audio-3/stable_audio_3/models/diffusion.py new file mode 100644 index 00000000..0b313e4c --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/diffusion.py @@ -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) + diff --git a/tools/stable-audio-3/stable_audio_3/models/dit.py b/tools/stable-audio-3/stable_audio_3/models/dit.py new file mode 100644 index 00000000..5d4d7745 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/dit.py @@ -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 + ) \ No newline at end of file diff --git a/tools/stable-audio-3/stable_audio_3/models/inpainting.py b/tools/stable-audio-3/stable_audio_3/models/inpainting.py new file mode 100644 index 00000000..54f878e7 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/inpainting.py @@ -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 \ No newline at end of file diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/__init__.py b/tools/stable-audio-3/stable_audio_3/models/lora/__init__.py new file mode 100644 index 00000000..a5841fac --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/lora/__init__.py @@ -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 \ No newline at end of file diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/__init__.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 00000000..77d6221d Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/__init__.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/loader.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/loader.cpython-310.pyc new file mode 100644 index 00000000..0539fcd9 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/loader.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/model.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/model.cpython-310.pyc new file mode 100644 index 00000000..2b1f481e Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/model.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/utils.cpython-310.pyc b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/utils.cpython-310.pyc new file mode 100644 index 00000000..0bbd82f0 Binary files /dev/null and b/tools/stable-audio-3/stable_audio_3/models/lora/__pycache__/utils.cpython-310.pyc differ diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/loader.py b/tools/stable-audio-3/stable_audio_3/models/lora/loader.py new file mode 100644 index 00000000..5b7685d3 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/lora/loader.py @@ -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 \ No newline at end of file diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/model.py b/tools/stable-audio-3/stable_audio_3/models/lora/model.py new file mode 100644 index 00000000..ed330f88 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/lora/model.py @@ -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 \ No newline at end of file diff --git a/tools/stable-audio-3/stable_audio_3/models/lora/utils.py b/tools/stable-audio-3/stable_audio_3/models/lora/utils.py new file mode 100644 index 00000000..b3c379a1 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/lora/utils.py @@ -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()) \ No newline at end of file diff --git a/tools/stable-audio-3/stable_audio_3/models/pretransforms.py b/tools/stable-audio-3/stable_audio_3/models/pretransforms.py new file mode 100644 index 00000000..b8b41400 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/pretransforms.py @@ -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 diff --git a/tools/stable-audio-3/stable_audio_3/models/transformer.py b/tools/stable-audio-3/stable_audio_3/models/transformer.py new file mode 100644 index 00000000..72989d66 --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/transformer.py @@ -0,0 +1,1272 @@ +from functools import reduce, partial +from packaging import version +import logging +import math + +from einops import rearrange, repeat +from einops.layers.torch import Rearrange +import torch +import torch.nn.functional as F +from torch import nn, einsum +from torch.amp import autocast +from torch.nn.utils.parametrizations import weight_norm +from typing import Callable, Literal, Optional +try: + from torch.nn.attention.flex_attention import flex_attention, create_block_mask + flex_attention_available = True +except ImportError: + flex_attention = None + create_block_mask = None + flex_attention_available = False + +try: + from flash_attn import flash_attn_func, flash_attn_kvpacked_func +except ImportError as e: + print(e) + print('flash_attn not installed, disabling Flash Attention') + flash_attn_kvpacked_func = None + flash_attn_func = None + +try: + from flash_attn import flash_attn_varlen_func + from flash_attn.bert_padding import pad_input, unpad_input, index_first_axis +except ImportError as e: + print(e) + print('flash_attn varlen/bert_padding not available, disabling varlen attention') + flash_attn_varlen_func = None + pad_input = None + unpad_input = None + index_first_axis = None + + +def precompute_varlen_metadata(padding_mask: torch.Tensor): + """ + Precompute varlen attention metadata once to avoid recomputation in every attention layer. + + Args: + padding_mask: Boolean tensor of shape (batch, seq_len) where True = valid + + Returns: + Dict with cu_seqlens, max_seqlen, indices, batch_size, seq_len for use in attention + """ + if padding_mask is None or unpad_input is None: + return None + + batch_size, seq_len = padding_mask.shape + + # Compute cumulative sequence lengths (same for all of q, k, v) + seqlens = padding_mask.sum(dim=-1, dtype=torch.int32) + cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + max_seqlen = seqlens.max().item() + + # Compute indices for gathering valid tokens + # indices maps from packed position -> original (batch, seq) position + indices = torch.nonzero(padding_mask.flatten(), as_tuple=False).flatten() + + return { + "cu_seqlens": cu_seqlens, + "max_seqlen": max_seqlen, + "indices": indices, + "batch_size": batch_size, + "seq_len": seq_len, + } + +from .utils import compile + + +def _left_pad_to_match(emb, target_len): + """Left-pad or right-trim emb along seq dim to match target_len. + + Used for local conditioning embeddings that need to align with x + without affecting prepended tokens (memory tokens, global cond, etc.). + """ + emb_len = emb.shape[-2] + if emb_len < target_len: + return F.pad(emb, (0, 0, target_len - emb_len, 0), value=0.) + elif emb_len > target_len: + return emb[:, -target_len:, :] + return emb + +if flex_attention_available: + try: + torch._dynamo.config.cache_size_limit = 5000 + flex_attention_compiled = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + except Exception as e: + logging.debug(f"Could not compile flex_attention, using uncompiled version: {e}") + flex_attention_compiled = flex_attention +else: + flex_attention_compiled = None + + +# Cache band block_masks for sliding-window attention fallback (flex_attention path). +# Keyed by (seq_q, seq_k, w_left, w_right, device). create_block_mask is expensive +# but the result is reused across all transformer layers and forward passes. +_SLIDING_WINDOW_BLOCK_MASK_CACHE = {} + +def _get_sliding_window_block_mask(seq_q, seq_k, w_left, w_right, device): + key = (seq_q, seq_k, int(w_left), int(w_right), str(device)) + bm = _SLIDING_WINDOW_BLOCK_MASK_CACHE.get(key) + if bm is None: + wl, wr = int(w_left), int(w_right) + def _band_mod(b, h, q_idx, kv_idx): + delta = kv_idx - q_idx + return (delta >= -wl) & (delta <= wr) + bm = create_block_mask(_band_mod, B=None, H=None, Q_LEN=seq_q, KV_LEN=seq_k, device=device) + _SLIDING_WINDOW_BLOCK_MASK_CACHE[key] = bm + return bm + +def _sliding_window_additive_mask(seq_q, seq_k, w_left, w_right, device, dtype): + """Build a (seq_q, seq_k) additive mask for masked SDPA fallback. + 0 inside the band [i - w_left, i + w_right], -inf outside. + """ + ii = torch.arange(seq_q, device=device) + jj = torch.arange(seq_k, device=device) + delta = jj[None, :] - ii[:, None] + in_band = (delta >= -int(w_left)) & (delta <= int(w_right)) + mask = torch.zeros((seq_q, seq_k), dtype=dtype, device=device) + return mask.masked_fill(~in_band, float('-inf')) + + +# Chunked-halo SDPA fallback. Math-equivalent to masked SDPA with a band +# mask, but processes queries in non-overlapping chunks with a (w_left, +# w_right) halo of keys/values on each side — every query stays inside its +# chunk's softmax. Avoids materializing the O(N^2) mask. +# +# At realistic SAME-L decoder shapes (N=69632, W=17, packed sequence is +# latent_length * (stride+1)): ~34x faster than full masked SDPA, and +# ~140x less peak mask memory (~1 MB per chunk vs 9.7 GB for one N x N mask). +# Chunk size is a tunable; 1024 is a good default at typical pretransform +# decoder shapes. Larger chunks waste more compute on out-of-band tiles; +# smaller chunks suffer from launch overhead. +_SLIDING_WINDOW_CHUNK_SIZE = 1024 + +def _sliding_window_chunked_halo_sdpa(q, k, v, w_left, w_right, chunk_size=_SLIDING_WINDOW_CHUNK_SIZE): + B, H, N, D = q.shape + outs = [] + for q_start in range(0, N, chunk_size): + q_end = min(q_start + chunk_size, N) + k_start = max(0, q_start - int(w_left)) + k_end = min(N, q_end + int(w_right)) + q_c = q[..., q_start:q_end, :] + k_c = k[..., k_start:k_end, :] + v_c = v[..., k_start:k_end, :] + q_idx = torch.arange(q_start, q_end, device=q.device) + k_idx = torch.arange(k_start, k_end, device=q.device) + delta = k_idx[None, :] - q_idx[:, None] + in_band = (delta >= -int(w_left)) & (delta <= int(w_right)) + mask = torch.zeros(delta.shape, dtype=q.dtype, device=q.device).masked_fill(~in_band, float('-inf')) + outs.append(F.scaled_dot_product_attention(q_c, k_c, v_c, attn_mask=mask, is_causal=False)) + return torch.cat(outs, dim=-2) + + +def checkpoint(function, *args, **kwargs): + kwargs.setdefault("use_reentrant", False) + # Preserve autocast context during recomputation to avoid dtype mismatches + if "context_fn" not in kwargs: + from torch.amp import autocast + import functools + # Get current autocast state + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype('cuda') + def get_contexts(): + return ( + autocast('cuda', dtype=dtype), + autocast('cuda', dtype=dtype), + ) + kwargs["context_fn"] = get_contexts + return torch.utils.checkpoint.checkpoint(function, *args, **kwargs) + + +# Copied and modified from https://github.com/lucidrains/x-transformers/blob/main/x_transformers/attend.py under MIT License +# License can be found in LICENSES/LICENSE_XTRANSFORMERS.txt + +def create_causal_mask(i, j, device): + return torch.ones((i, j), device = device, dtype = torch.bool).triu(j - i + 1) + +def or_reduce(masks): + head, *body = masks + for rest in body: + head = head | rest + return head + +# positional embeddings + +class AbsolutePositionalEmbedding(nn.Module): + def __init__(self, dim, max_seq_len): + super().__init__() + self.scale = dim ** -0.5 + self.max_seq_len = max_seq_len + self.emb = nn.Embedding(max_seq_len, dim) + + def forward(self, x, pos = None, seq_start_pos = None): + seq_len, device = x.shape[1], x.device + assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}' + + if pos is None: + pos = torch.arange(seq_len, device = device) + + if seq_start_pos is not None: + pos = (pos - seq_start_pos[..., None]).clamp(min = 0) + + pos_emb = self.emb(pos) + pos_emb = pos_emb * self.scale + return pos_emb + +class ScaledSinusoidalEmbedding(nn.Module): + def __init__(self, dim, theta = 10000): + super().__init__() + assert (dim % 2) == 0, 'dimension must be divisible by 2' + self.scale = nn.Parameter(torch.ones(1) * dim ** -0.5) + + half_dim = dim // 2 + freq_seq = torch.arange(half_dim).float() / half_dim + inv_freq = theta ** -freq_seq + self.register_buffer('inv_freq', inv_freq, persistent = False) + + def forward(self, x, pos = None, seq_start_pos = None): + seq_len, device = x.shape[1], x.device + + if pos is None: + pos = torch.arange(seq_len, device = device) + + if seq_start_pos is not None: + pos = pos - seq_start_pos[..., None] + + emb = einsum('i, j -> i j', pos, self.inv_freq) + emb = torch.cat((emb.sin(), emb.cos()), dim = -1) + return emb * self.scale + +class RotaryEmbedding(nn.Module): + def __init__( + self, + dim, + use_xpos = False, + scale_base = 512, + interpolation_factor = 1., + base = 10000, + base_rescale_factor = 1. + ): + super().__init__() + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ + base *= base_rescale_factor ** (dim / (dim - 2)) + + inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer('inv_freq', inv_freq) + + assert interpolation_factor >= 1. + self.interpolation_factor = interpolation_factor + + if not use_xpos: + self.register_buffer('scale', None) + return + + scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) + + self.scale_base = scale_base + self.register_buffer('scale', scale) + + def forward_from_seq_len(self, seq_len): + device = self.inv_freq.device + + t = torch.arange(seq_len, device = device) + return self.forward(t) + + @autocast("cuda", enabled = False) + def forward(self, t): + device = self.inv_freq.device + + t = t.to(torch.float32) + + t = t / self.interpolation_factor + + freqs = torch.einsum('i , j -> i j', t, self.inv_freq) + freqs = torch.cat((freqs, freqs), dim = -1) + + if self.scale is None: + return freqs, 1. + + power = (torch.arange(seq_len, device = device) - (seq_len // 2)) / self.scale_base + scale = self.scale ** rearrange(power, 'n -> n 1') + scale = torch.cat((scale, scale), dim = -1) + + return freqs, scale + +def rotate_half(x): + x = rearrange(x, '... (j d) -> ... j d', j = 2) + x1, x2 = x.unbind(dim = -2) + return torch.cat((-x2, x1), dim = -1) + + +@autocast("cuda", enabled = False) +def apply_rotary_pos_emb(t, freqs, scale = 1): + out_dtype = t.dtype + + # cast to float32 if necessary for numerical stability + dtype = reduce(torch.promote_types, (t.dtype, freqs.dtype, torch.float32)) + rot_dim, seq_len = freqs.shape[-1], t.shape[-2] + freqs, t = freqs.to(dtype), t.to(dtype) + freqs = freqs[-seq_len:, :] + + if t.ndim == 4 and freqs.ndim == 3: + freqs = rearrange(freqs, 'b n d -> b 1 n d') + + # partial rotary embeddings, Wang et al. GPT-J + t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:] + + t = (t * freqs.cos() * scale ) + (rotate_half(t) * freqs.sin() * scale) + + t, t_unrotated = t.to(out_dtype), t_unrotated.to(out_dtype) + + return torch.cat((t, t_unrotated), dim = -1) + +# norms +class DynamicTanh(nn.Module): + def __init__(self, dim, init_alpha=4.0, **kwargs): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1) * init_alpha) + self.gamma = nn.Parameter(torch.ones(dim)) + self.beta = nn.Parameter(torch.zeros(dim)) + + def forward(self, x): + x = F.tanh(self.alpha * x) + return self.gamma * x + self.beta + +class RunningInstanceNorm(nn.Module): + def __init__(self, dim, momentum = 0.99, eps = 1e-4, saturate = True, trainable_gain = True): + super().__init__() + self.register_buffer("running_mean", torch.zeros(1,1,dim)) + self.register_buffer("running_std", torch.ones(1,1,dim)) + self.saturate = saturate + self.eps = eps + self.momentum = momentum + self.dim = dim + self.trainable_gain = trainable_gain + if self.trainable_gain: + self.gain = nn.Parameter(torch.ones(1)) + + def _update_stats(self, x): + self.running_mean = self.running_mean * self.momentum + x.detach().mean(dim = [0,1]).view(1, 1, self.dim) * (1 - self.momentum) + self.running_std = (self.running_std * self.momentum + x.detach().std(dim = [0,1]).view(1, 1, self.dim) * (1 - self.momentum)).clip(min = self.eps) + + def forward(self, x): + if self.training: + self._update_stats(x) + x = (x - self.running_mean) / self.running_std + if self.saturate: + x = torch.asinh(x) + if self.trainable_gain: + x = x * self.gain + return x + +class LayerNorm(nn.Module): + def __init__(self, dim, bias=False, fix_scale=False, force_fp32=False, eps=1e-5): + """ + bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less + """ + super().__init__() + + if fix_scale: + self.register_buffer("gamma", torch.ones(dim)) + else: + self.gamma = nn.Parameter(torch.ones(dim)) + + if bias: + self.beta = nn.Parameter(torch.zeros(dim)) + else: + self.register_buffer("beta", torch.zeros(dim)) + + self.eps = eps + + self.force_fp32 = force_fp32 + + #@autocast("cuda", enabled = False) + def forward(self, x): + if not self.force_fp32: + return F.layer_norm(x, x.shape[-1:], weight=self.gamma, bias=self.beta, eps=self.eps) + else: + output = F.layer_norm(x.float(), x.shape[-1:], weight=self.gamma.float(), bias=self.beta.float(), eps=self.eps) + return output.to(x.dtype) + +class RMSNorm(nn.Module): + def __init__(self, dim, fix_scale=False, force_fp32=False, eps=1e-5): + super().__init__() + + if fix_scale: + self.register_buffer("gamma", torch.ones(dim)) + else: + self.gamma = nn.Parameter(torch.ones(dim)) + + self.eps = eps + + self.force_fp32 = force_fp32 + + def forward(self, x): + if not self.force_fp32: + return F.rms_norm(x, x.shape[-1:], weight=self.gamma, eps=self.eps) + else: + output = F.rms_norm(x.float(), x.shape[-1:], weight=self.gamma.float(), eps=self.eps) + return output.to(x.dtype) + +class LayerScale(nn.Module): + def __init__(self, dim, init_val = 1e-5): + super().__init__() + self.scale = nn.Parameter(torch.full([dim], init_val)) + def forward(self, x): + return x * self.scale + +# feedforward + +class GLU(nn.Module): + def __init__( + self, + dim_in, + dim_out, + activation: Callable, + use_conv = False, + conv_kernel_size = 3, + ): + super().__init__() + self.act = activation + self.proj = nn.Linear(dim_in, dim_out * 2) if not use_conv else nn.Conv1d(dim_in, dim_out * 2, conv_kernel_size, padding = (conv_kernel_size // 2)) + self.use_conv = use_conv + + def forward(self, x): + if self.use_conv: + x = rearrange(x, 'b n d -> b d n') + x = self.proj(x) + x = rearrange(x, 'b d n -> b n d') + else: + x = self.proj(x) + + x, gate = x.chunk(2, dim = -1) + return x * self.act(gate) + +class Sin(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.sin(3.14159265359 * x) + +class FeedForward(nn.Module): + def __init__( + self, + dim, + dim_out = None, + mult = 4, + no_bias = False, + glu = True, + use_conv = False, + conv_kernel_size = 3, + zero_init_output = True, + sinusoidal = False + ): + super().__init__() + inner_dim = int(dim * mult) + + # Default to SwiGLU + + activation = nn.SiLU() if not sinusoidal else Sin() + + dim_out = dim if dim_out is None else dim_out + + if glu: + linear_in = GLU(dim, inner_dim, activation) + else: + linear_in = nn.Sequential( + Rearrange('b n d -> b d n') if use_conv else nn.Identity(), + nn.Linear(dim, inner_dim, bias = not no_bias) if not use_conv else nn.Conv1d(dim, inner_dim, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias), + Rearrange('b n d -> b d n') if use_conv else nn.Identity(), + activation + ) + + linear_out = nn.Linear(inner_dim, dim_out, bias = not no_bias) if not use_conv else nn.Conv1d(inner_dim, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias) + + # init last linear layer to 0 + if zero_init_output: + nn.init.zeros_(linear_out.weight) + if not no_bias: + nn.init.zeros_(linear_out.bias) + + + self.ff = nn.Sequential( + linear_in, + Rearrange('b d n -> b n d') if use_conv else nn.Identity(), + linear_out, + Rearrange('b n d -> b d n') if use_conv else nn.Identity(), + ) + + #@compile + def forward(self, x, varlen_metadata=None): + if varlen_metadata is not None and index_first_axis is not None and pad_input is not None: + # Pack valid tokens for efficient FFN computation (skip padding tokens) + # Padding positions become zeros after unpack, which is fine since FFN output + # is added to residual, preserving values at padding positions + batch_size = varlen_metadata["batch_size"] + seq_len = varlen_metadata["seq_len"] + indices = varlen_metadata["indices"] + dim = x.shape[-1] + + # Pack to (N_valid, D) + x_packed = index_first_axis(x.reshape(-1, dim), indices) + + # FFN on packed representation with pseudo-batch dim + x_packed = self.ff(x_packed.unsqueeze(0)).squeeze(0) + + # Unpack back to (B, T, D) + return pad_input(x_packed, indices, batch_size, seq_len) + else: + return self.ff(x) + +class Attention(nn.Module): + def __init__( + self, + dim, + dim_heads = 64, + dim_context = None, + causal = False, + zero_init_output=True, + qk_norm_eps = 1e-6, + qk_norm: Literal['l2', 'ln', 'rms', 'dyt', 'none'] = 'none', + differential = False, + feat_scale = False + ): + super().__init__() + self.dim = dim + self.dim_heads = dim_heads + + self.differential = differential + + dim_kv = dim_context if dim_context is not None else dim + + self.num_heads = dim // dim_heads + self.kv_heads = dim_kv // dim_heads + + if dim_context is not None: + if differential: + self.to_q = nn.Linear(dim, dim * 2, bias=False) + self.to_kv = nn.Linear(dim_kv, dim_kv * 3, bias=False) + else: + self.to_q = nn.Linear(dim, dim, bias=False) + self.to_kv = nn.Linear(dim_kv, dim_kv * 2, bias=False) + else: + if differential: + self.to_qkv = nn.Linear(dim, dim * 5, bias=False) + else: + self.to_qkv = nn.Linear(dim, dim * 3, bias=False) + + self.to_out = nn.Linear(dim, dim, bias=False) + + if zero_init_output: + nn.init.zeros_(self.to_out.weight) + + if qk_norm not in ['l2', 'ln', 'rms', 'dyt','none']: + raise ValueError(f'qk_norm must be one of ["l2", "ln", "rms" ,"dyt", "none"], got {qk_norm}') + + self.qk_norm = qk_norm + self.qk_norm_eps = qk_norm_eps + + if self.qk_norm == "ln": + self.q_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=qk_norm_eps) + self.k_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=qk_norm_eps) + elif self.qk_norm == "rms": + self.q_norm = RMSNorm(dim_heads, eps=qk_norm_eps) + self.k_norm = RMSNorm(dim_heads, eps=qk_norm_eps) + elif self.qk_norm == 'dyt': + self.q_norm = DynamicTanh(dim_heads) + self.k_norm = DynamicTanh(dim_heads) + + self.feat_scale = feat_scale + + if self.feat_scale: + self.lambda_dc = nn.Parameter(torch.zeros(dim)) + self.lambda_hf = nn.Parameter(torch.zeros(dim)) + + self.causal = causal + + @compile + def apply_qk_layernorm(self, q, k): + q_type = q.dtype + k_type = k.dtype + q = self.q_norm(q).to(q_type) + k = self.k_norm(k).to(k_type) + return q, k + + + def apply_attn(self, q, k, v, causal = None, flex_attention_block_mask = None, flex_attention_score_mod = None, flash_attn_sliding_window = None, padding_mask = None, varlen_metadata = None): + + if self.num_heads != self.kv_heads: + # Repeat interleave kv_heads to match q_heads for grouped query attention + heads_per_kv_head = self.num_heads // self.kv_heads + k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v)) + + flash_attn_available = flash_attn_func is not None + flash_attn_varlen_available = flash_attn_varlen_func is not None and index_first_axis is not None + + if causal and (flex_attention_block_mask is not None or flex_attention_score_mod is not None): + flex_attention_block_mask = None + flex_attention_score_mod = None + + if flex_attention_block_mask is not None or flex_attention_score_mod is not None: + # Flex attention path - use V-zeroing for padding mask + if padding_mask is not None: + mask_expanded = padding_mask.unsqueeze(1).unsqueeze(-1).to(v.dtype) + v = v * mask_expanded + out = flex_attention_compiled(q,k,v, + block_mask = flex_attention_block_mask, + score_mod = flex_attention_score_mod) + elif flash_attn_available and varlen_metadata is not None and flash_attn_varlen_available: + # Flash attention with varlen using precomputed metadata (fast path) + batch_size = varlen_metadata["batch_size"] + seq_len = varlen_metadata["seq_len"] + cu_seqlens = varlen_metadata["cu_seqlens"] + max_seqlen = varlen_metadata["max_seqlen"] + indices = varlen_metadata["indices"] + + fa_dtype_in = q.dtype + # Rearrange to (B, T, H, D) for flash_attn + q, k, v = map(lambda t: rearrange(t, 'b h n d -> b n h d'), (q, k, v)) + + if fa_dtype_in != torch.float16 and fa_dtype_in != torch.bfloat16: + q, k, v = map(lambda t: t.to(torch.float16), (q, k, v)) + + # Pack q, k, v using precomputed indices (much faster than calling unpad_input 3x) + num_heads, head_dim = q.shape[2], q.shape[3] + q_unpad = index_first_axis(q.reshape(-1, num_heads, head_dim), indices) + k_unpad = index_first_axis(k.reshape(-1, num_heads, head_dim), indices) + v_unpad = index_first_axis(v.reshape(-1, num_heads, head_dim), indices) + + out_unpad = flash_attn_varlen_func( + q_unpad, k_unpad, v_unpad, + cu_seqlens, cu_seqlens, + max_seqlen, max_seqlen, + causal=causal if causal is not None else False, + window_size=flash_attn_sliding_window if flash_attn_sliding_window is not None else (-1, -1), + ) + + # Pad output back to original shape + out = pad_input(out_unpad, indices, batch_size, seq_len) + out = rearrange(out.to(fa_dtype_in), 'b n h d -> b h n d') + elif flash_attn_available: + # Standard flash attention (no padding mask, or varlen imports not available) + # Apply V-zeroing fallback if padding_mask provided but we couldn't use varlen + if padding_mask is not None: + mask_expanded = padding_mask.unsqueeze(1).unsqueeze(-1).to(v.dtype) + v = v * mask_expanded + fa_dtype_in = q.dtype + q, k, v = map(lambda t: rearrange(t, 'b h n d -> b n h d'), (q, k, v)) + + if fa_dtype_in != torch.float16 and fa_dtype_in != torch.bfloat16: + q, k, v = map(lambda t: t.to(torch.float16), (q, k, v)) + + out = flash_attn_func(q, k, v, causal = causal, window_size=flash_attn_sliding_window if (flash_attn_sliding_window is not None) else [-1,-1]) + + out = rearrange(out.to(fa_dtype_in), 'b n h d -> b h n d') + else: + # No flash-attn available. Sliding-window fallback cascade: + # Tier 2: flex_attention with band block_mask (best when torch.compile works) + # Tier 3: chunked-halo masked SDPA (math-equivalent, ~30x faster than tier 4) + # Tier 4: full masked SDPA (N x N mask) (last resort; high memory) + # For the no-sliding-window case, fall through to plain SDPA full attention. + # All apply V-zeroing for padding masks (cheap and equivalent to masking + # those positions out of attention output). + if padding_mask is not None: + mask_expanded = padding_mask.unsqueeze(1).unsqueeze(-1).to(v.dtype) + v = v * mask_expanded + if flash_attn_sliding_window is not None: + seq_q, seq_k = q.shape[2], k.shape[2] + wl, wr = flash_attn_sliding_window + handled = False + if flex_attention_available and flex_attention_compiled is not None: + try: + bm = _get_sliding_window_block_mask(seq_q, seq_k, wl, wr, q.device) + out = flex_attention_compiled(q, k, v, block_mask=bm) + handled = True + except Exception as _flex_err: + logging.debug(f"flex_attention failed, trying chunked-halo SDPA: {_flex_err}") + if not handled: + try: + out = _sliding_window_chunked_halo_sdpa(q, k, v, wl, wr) + handled = True + except Exception as _chunk_err: + logging.debug(f"chunked-halo SDPA failed, falling back to full masked SDPA: {_chunk_err}") + if not handled: + add_mask = _sliding_window_additive_mask(seq_q, seq_k, wl, wr, q.device, q.dtype) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=add_mask, is_causal=False) + else: + out = F.scaled_dot_product_attention(q, k, v, is_causal=causal if causal is not None else False) + return out + + + #@compile + def forward( + self, + x, + context = None, + rotary_pos_emb = None, + rotary_pos_emb_k = None, + causal = None, + flex_attention_block_mask = None, + flex_attention_score_mod = None, + flash_attn_sliding_window = None, + padding_mask = None, + varlen_metadata = None, + ): + h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None + + kv_input = context if has_context else x + + if hasattr(self, 'to_q'): + # Use separate linear projections for q and k/v + if self.differential: + q, q_diff = self.to_q(x).chunk(2, dim=-1) + q, q_diff = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, q_diff)) + q = torch.stack([q, q_diff], dim = 1) + k, k_diff, v = self.to_kv(kv_input).chunk(3, dim=-1) + k, k_diff, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, k_diff, v)) + k = torch.stack([k, k_diff], dim = 1) + else: + q = self.to_q(x) + q = rearrange(q, 'b n (h d) -> b h n d', h = h) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, v)) + else: + # Use fused linear projection + if self.differential: + q, k, v, q_diff, k_diff = self.to_qkv(x).chunk(5, dim=-1) + q, k, v, q_diff, k_diff = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v, q_diff, k_diff)) + q = torch.stack([q, q_diff], dim = 1) + k = torch.stack([k, k_diff], dim = 1) + else: + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v)) + + # Normalize q and k for cosine sim attention + if self.qk_norm == "l2": + q = F.normalize(q, dim=-1, eps=self.qk_norm_eps) + k = F.normalize(k, dim=-1, eps=self.qk_norm_eps) + elif self.qk_norm != "none": + q, k = self.apply_qk_layernorm(q, k) + # q = self.q_norm(q) + # k = self.k_norm(k) + + if rotary_pos_emb is not None: + freqs, _ = rotary_pos_emb + q_dtype = q.dtype + k_dtype = k.dtype + q = q.to(torch.float32) + k = k.to(torch.float32) + freqs = freqs.to(torch.float32) + + q_freqs = freqs + + if rotary_pos_emb_k is not None: + k_freqs, _ = rotary_pos_emb_k + k_freqs = k_freqs.to(torch.float32) + else: + k_freqs = q_freqs + + if q.shape[-2] >= k.shape[-2]: + ratio = q.shape[-2] / k.shape[-2] + q_freqs, k_freqs = freqs, ratio * freqs + else: + ratio = k.shape[-2] / q.shape[-2] + q_freqs, k_freqs = ratio * freqs, freqs + + q = apply_rotary_pos_emb(q, q_freqs) + k = apply_rotary_pos_emb(k, k_freqs) + q = q.to(v.dtype) + k = k.to(v.dtype) + + n, device = q.shape[-2], q.device + + causal = self.causal if causal is None else causal + + if n == 1 and causal: + causal = False + + if self.differential: + q, q_diff = q.unbind(dim = 1) + k, k_diff = k.unbind(dim = 1) + out = self.apply_attn(q, k, v, causal = causal, flex_attention_block_mask = flex_attention_block_mask, flex_attention_score_mod = flex_attention_score_mod, flash_attn_sliding_window = flash_attn_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata) + out_diff = self.apply_attn(q_diff, k_diff, v, causal = causal, flex_attention_block_mask = flex_attention_block_mask, flex_attention_score_mod = flex_attention_score_mod, flash_attn_sliding_window = flash_attn_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata) + out = out - out_diff + else: + out = self.apply_attn(q, k, v, causal = causal, flex_attention_block_mask = flex_attention_block_mask, flex_attention_score_mod = flex_attention_score_mod, flash_attn_sliding_window = flash_attn_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata) + # merge heads + out = rearrange(out, ' b h n d -> b n (h d)') + + # Communicate between heads + + # with autocast(enabled = False): + # out_dtype = out.dtype + # out = out.to(torch.float32) + # out = self.to_out(out).to(out_dtype) + out = self.to_out(out) + + if self.feat_scale: + if padding_mask is not None: + mask = padding_mask.unsqueeze(-1).to(out.dtype) # (b, n, 1) + out_dc = (out * mask).sum(dim=-2, keepdim=True) / mask.sum(dim=-2, keepdim=True).clamp(min=1) + out_hf = out - out_dc + out = out + (self.lambda_dc * out_dc + self.lambda_hf * out_hf) * mask + else: + out_dc = out.mean(dim=-2, keepdim=True) + out_hf = out - out_dc + out = out + self.lambda_dc * out_dc + self.lambda_hf * out_hf + + return out + +class ConformerModule(nn.Module): + def __init__( + self, + dim, + norm_kwargs = {}, + ): + + super().__init__() + + self.dim = dim + + self.in_norm = LayerNorm(dim, **norm_kwargs) + self.pointwise_conv = nn.Conv1d(dim, dim, kernel_size=1, bias=False) + self.glu = GLU(dim, dim, nn.SiLU()) + self.depthwise_conv = nn.Conv1d(dim, dim, kernel_size=17, groups=dim, padding=8, bias=False) + self.mid_norm = LayerNorm(dim, **norm_kwargs) # This is a batch norm in the original but I don't like batch norm + self.swish = nn.SiLU() + self.pointwise_conv_2 = nn.Conv1d(dim, dim, kernel_size=1, bias=False) + + #@compile + def forward(self, x): + x = self.in_norm(x) + x = rearrange(x, 'b n d -> b d n') + x = self.pointwise_conv(x) + x = rearrange(x, 'b d n -> b n d') + x = self.glu(x) + x = rearrange(x, 'b n d -> b d n') + x = self.depthwise_conv(x) + x = rearrange(x, 'b d n -> b n d') + x = self.mid_norm(x) + x = self.swish(x) + x = rearrange(x, 'b n d -> b d n') + x = self.pointwise_conv_2(x) + x = rearrange(x, 'b d n -> b n d') + + return x + +class TransformerBlock(nn.Module): + def __init__( + self, + dim, + dim_heads = 64, + cross_attend = False, + dim_context = None, + global_cond_dim = None, + local_add_cond_dim = None, + modular_local_cond_configs = None, + causal = False, + zero_init_branch_outputs = True, + conformer = False, + layer_ix = -1, + add_rope = False, + layer_scale = False, + norm_type = 'layer_norm', + attn_kwargs = {}, + ff_kwargs = {}, + norm_kwargs = {} + ): + + super().__init__() + self.dim = dim + self.dim_heads = min(dim_heads,dim) + self.cross_attend = cross_attend + self.dim_context = dim_context + self.causal = causal + + if layer_scale and zero_init_branch_outputs: + print('zero_init_branch_outputs is redundant with layer_scale, setting zero_init_branch_outputs to False') + zero_init_branch_outputs = False + + if norm_type not in ['layer_norm', 'rms_norm', 'dyt']: + raise ValueError(f'norm_type must be one of ["layer_norm", "rms_norm", "dyt"], got {norm_type}') + + norm_layer_map = { + 'layer_norm': LayerNorm, + 'rms_norm': RMSNorm, + 'dyt': DynamicTanh + } + norm_layer = norm_layer_map[norm_type] + + self.pre_norm = norm_layer(dim,**norm_kwargs) + self.add_rope = add_rope + + self.self_attn = Attention( + dim, + dim_heads = self.dim_heads, + causal = causal, + zero_init_output=zero_init_branch_outputs, + **attn_kwargs + ) + + self.self_attn_scale = LayerScale(dim) if layer_scale else nn.Identity() + + self.cross_attend = cross_attend + if cross_attend: + self.cross_attend_norm = norm_layer(dim, **norm_kwargs) + self.cross_attn = Attention( + dim, + dim_heads = self.dim_heads, + dim_context=dim_context, + causal = causal, + zero_init_output=zero_init_branch_outputs, + **attn_kwargs + ) + self.cross_attn_scale = LayerScale(dim) if layer_scale else nn.Identity() + + self.ff_norm = norm_layer(dim, **norm_kwargs) + self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, **ff_kwargs) + self.ff_scale = LayerScale(dim) if layer_scale else nn.Identity() + + self.layer_ix = layer_ix + + self.conformer = None + if conformer: + self.conformer = ConformerModule(dim, norm_kwargs=norm_kwargs) + self.conformer_scale = LayerScale(dim) if layer_scale else nn.Identity() + + self.global_cond_dim = global_cond_dim + + if global_cond_dim is not None: + self.to_scale_shift_gate = nn.Parameter(torch.randn(6*dim)/dim**0.5) + + self.local_add_cond_dim = local_add_cond_dim + + if local_add_cond_dim is not None: + self.to_local_embed = nn.Sequential( + nn.Linear(local_add_cond_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + + nn.init.zeros_(self.to_local_embed[-1].weight) + nn.init.zeros_(self.to_local_embed[-1].bias) + + else: + self.to_local_embed = None + + # Modular local conditioning - independent projections per conditioning ID + self.modular_local_cond_configs = modular_local_cond_configs or [] + self.modular_local_embeds = nn.ModuleDict() + + for config in self.modular_local_cond_configs: + cond_id = config["id"] + cond_dim = config["dim"] + proj = nn.Sequential( + nn.Linear(cond_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + # Zero-init output layer so new conditioning doesn't affect model initially + nn.init.zeros_(proj[-1].weight) + nn.init.zeros_(proj[-1].bias) + self.modular_local_embeds[cond_id] = proj + + self.rope = RotaryEmbedding(self.dim_heads // 2) if add_rope else None + + def _apply_local_conditioning(self, x, local_add_cond, modular_local_cond): + """Apply local additive and modular local conditioning to x.""" + if local_add_cond is not None and self.to_local_embed is not None: + local_emb = self.to_local_embed(local_add_cond) + x = x + _left_pad_to_match(local_emb, x.shape[-2]) + + if modular_local_cond is not None and len(self.modular_local_embeds) > 0: + modular_sum = None + for cond_id, proj in self.modular_local_embeds.items(): + if cond_id in modular_local_cond: + local_emb = proj(modular_local_cond[cond_id]) + local_emb = _left_pad_to_match(local_emb, x.shape[-2]) + modular_sum = local_emb if modular_sum is None else modular_sum + local_emb + if modular_sum is not None: + x = x + modular_sum + + return x + + @compile + def forward( + self, + x, + context = None, + global_cond=None, + local_add_cond=None, + modular_local_cond=None, + rotary_pos_emb = None, + cross_attn_rotary_pos_emb = None, + self_attention_block_mask = None, + self_attention_score_mod = None, + cross_attention_block_mask = None, + cross_attention_score_mod = None, + self_attention_flash_sliding_window = None, + cross_attention_flash_sliding_window = None, + padding_mask = None, + varlen_metadata = None, + ): + if rotary_pos_emb is None and self.add_rope: + rotary_pos_emb = self.rope.forward_from_seq_len(x.shape[-2]) + + if self.global_cond_dim is not None and self.global_cond_dim > 0 and global_cond is not None: + + scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = (self.to_scale_shift_gate + global_cond).unsqueeze(1).chunk(6, dim=-1) + + # self-attention with adaLN + residual = x + x = self.pre_norm(x) + x = x * (1 + scale_self) + shift_self + x = self.self_attn(x, rotary_pos_emb = rotary_pos_emb, flex_attention_block_mask = self_attention_block_mask, flex_attention_score_mod = self_attention_score_mod, flash_attn_sliding_window = self_attention_flash_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata) + x = x * torch.sigmoid(1 - gate_self) + x = self.self_attn_scale(x) + x = x + residual + + if context is not None and self.cross_attend: + if cross_attn_rotary_pos_emb is not None: + x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), rotary_pos_emb = rotary_pos_emb, rotary_pos_emb_k = cross_attn_rotary_pos_emb, context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window)) + else: + x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window)) + + if self.conformer is not None: + x = x + self.conformer_scale(self.conformer(x)) + + x = self._apply_local_conditioning(x, local_add_cond, modular_local_cond) + + # feedforward with adaLN + residual = x + x = self.ff_norm(x) + x = x * (1 + scale_ff) + shift_ff + x = self.ff(x, varlen_metadata=varlen_metadata) + x = x * torch.sigmoid(1 - gate_ff) + x = self.ff_scale(x) + x = x + residual + + else: + x = x + self.self_attn_scale(self.self_attn(self.pre_norm(x), rotary_pos_emb = rotary_pos_emb, flex_attention_block_mask = self_attention_block_mask, flex_attention_score_mod = self_attention_score_mod, flash_attn_sliding_window = self_attention_flash_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata)) + + if context is not None and self.cross_attend: + if cross_attn_rotary_pos_emb is not None: + x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), rotary_pos_emb = rotary_pos_emb, rotary_pos_emb_k = cross_attn_rotary_pos_emb, context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window)) + else: + x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window)) + + if self.conformer is not None: + x = x + self.conformer_scale(self.conformer(x)) + + x = self._apply_local_conditioning(x, local_add_cond, modular_local_cond) + + x = x + self.ff_scale(self.ff(self.ff_norm(x), varlen_metadata=varlen_metadata)) + + + return x + +class ContinuousTransformer(nn.Module): + def __init__( + self, + dim, + depth, + *, + dim_in = None, + dim_out = None, + dim_heads = 64, + cross_attend=False, + cond_token_dim=None, + final_cross_attn_ix=-1, + global_cond_dim=None, + local_add_cond_dim=None, + modular_local_cond_configs=None, + causal=False, + rotary_pos_emb=True, + cross_attn_rotary_pos_emb=False, + zero_init_branch_outputs=True, + conformer=False, + use_sinusoidal_emb=False, + use_abs_pos_emb=False, + abs_pos_emb_max_length=10000, + num_memory_tokens=0, + sliding_window=None, + **kwargs + ): + + super().__init__() + + self.dim = dim + self.depth = depth + self.causal = causal + self.layers = nn.ModuleList([]) + + self.project_in = nn.Linear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity() + self.project_out = nn.Linear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity() + + if rotary_pos_emb: + self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32)) + else: + self.rotary_pos_emb = None + + if cross_attn_rotary_pos_emb: + self.cross_attn_rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32)) + else: + self.cross_attn_rotary_pos_emb = None + + self.num_memory_tokens = num_memory_tokens + if num_memory_tokens > 0: + self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim)) + + self.use_sinusoidal_emb = use_sinusoidal_emb + if use_sinusoidal_emb: + self.pos_emb = ScaledSinusoidalEmbedding(dim) + + self.use_abs_pos_emb = use_abs_pos_emb + if use_abs_pos_emb: + self.pos_emb = AbsolutePositionalEmbedding(dim, abs_pos_emb_max_length + self.num_memory_tokens) + + self.global_cond_embedder = None + if global_cond_dim is not None: + self.global_cond_embedder = nn.Sequential( + nn.Linear(global_cond_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim * 6) + ) + + self.final_cross_attn_ix = final_cross_attn_ix + + self.sliding_window = sliding_window + + for i in range(depth): + should_cross_attend = cross_attend and (self.final_cross_attn_ix == -1 or i <= (self.final_cross_attn_ix)) + self.layers.append( + TransformerBlock( + dim, + dim_heads = dim_heads, + cross_attend = should_cross_attend, + dim_context = cond_token_dim, + global_cond_dim = global_cond_dim, + local_add_cond_dim = local_add_cond_dim, + modular_local_cond_configs = modular_local_cond_configs, + causal = causal, + zero_init_branch_outputs = zero_init_branch_outputs, + conformer=conformer, + layer_ix=i, + **kwargs + ) + ) + + def forward( + self, + x, + context = None, + prepend_embeds = None, + global_cond = None, + local_add_cond = None, + modular_local_cond = None, + return_info = False, + use_checkpointing = True, + exit_layer_ix = None, + padding_mask: Optional[torch.Tensor] = None, + **kwargs + ): + batch, seq, device = *x.shape[:2], x.device + + model_dtype = next(self.parameters()).dtype + x = x.to(model_dtype) + + info = { + "hidden_states": [], + } + + x = self.project_in(x) + + if prepend_embeds is not None: + prepend_length, prepend_dim = prepend_embeds.shape[1:] + + assert prepend_dim == x.shape[-1], 'prepend dimension must match sequence dimension' + + x = torch.cat((prepend_embeds, x), dim = -2) + + if self.num_memory_tokens > 0: + memory_tokens = self.memory_tokens.expand(batch, -1, -1) + x = torch.cat((memory_tokens, x), dim=1) + + if self.rotary_pos_emb is not None: + rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1]) + else: + rotary_pos_emb = None + + if self.cross_attn_rotary_pos_emb is not None: + cross_attn_rotary_pos_emb = self.cross_attn_rotary_pos_emb.forward_from_seq_len(context.shape[-1]) + else: + cross_attn_rotary_pos_emb = None + + if self.use_sinusoidal_emb or self.use_abs_pos_emb: + x = x + self.pos_emb(x) + + if global_cond is not None and self.global_cond_embedder is not None: + global_cond = self.global_cond_embedder(global_cond) + + # Extend padding mask for prepended tokens if provided + extended_padding_mask = None + varlen_metadata = None + if padding_mask is not None: + # Compute total prepend length (memory tokens + prepend_embeds) + prepend_length = self.num_memory_tokens + if prepend_embeds is not None: + prepend_length += prepend_embeds.shape[1] + + # Prepend tokens are always valid for attention + if prepend_length > 0: + prepend_valid = torch.ones(batch, prepend_length, device=device, dtype=torch.bool) + extended_padding_mask = torch.cat([prepend_valid, padding_mask], dim=-1) + else: + extended_padding_mask = padding_mask + + # Precompute varlen metadata once for all layers (major performance optimization) + # Only compute if varlen attention is actually available + if flash_attn_varlen_func is not None and index_first_axis is not None: + varlen_metadata = precompute_varlen_metadata(extended_padding_mask) + + # Iterate over the transformer layers + for layer_ix, layer in enumerate(self.layers): + + layer_kwargs = { + "context": context, + "rotary_pos_emb": rotary_pos_emb, + "cross_attn_rotary_pos_emb": cross_attn_rotary_pos_emb, + "global_cond": global_cond, + "local_add_cond": local_add_cond, + "modular_local_cond": modular_local_cond, + "self_attention_flash_sliding_window": self.sliding_window, + "padding_mask": extended_padding_mask, + "varlen_metadata": varlen_metadata + } + + if use_checkpointing: + x = checkpoint(layer, x, **layer_kwargs, **kwargs) + else: + x = layer(x, **layer_kwargs, **kwargs) + + if return_info: + info["hidden_states"].append(x) + + if exit_layer_ix is not None and layer_ix == exit_layer_ix: + x = x[:, self.num_memory_tokens:, :] + + if return_info: + return x, info + + return x + + x = x[:, self.num_memory_tokens:, :] + + x = self.project_out(x) + + if return_info: + return x, info + + return x diff --git a/tools/stable-audio-3/stable_audio_3/models/utils.py b/tools/stable-audio-3/stable_audio_3/models/utils.py new file mode 100644 index 00000000..f3d8f36f --- /dev/null +++ b/tools/stable-audio-3/stable_audio_3/models/utils.py @@ -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 diff --git a/web/dungeon/audio/sfx/atk_hit.m4a b/web/dungeon/audio/sfx/atk_hit.m4a index de251231..42c8f538 100644 Binary files a/web/dungeon/audio/sfx/atk_hit.m4a and b/web/dungeon/audio/sfx/atk_hit.m4a differ diff --git a/web/dungeon/audio/sfx/atk_hit.ogg b/web/dungeon/audio/sfx/atk_hit.ogg index b539b58b..3b9ac2ed 100644 Binary files a/web/dungeon/audio/sfx/atk_hit.ogg and b/web/dungeon/audio/sfx/atk_hit.ogg differ diff --git a/web/dungeon/audio/sfx/echo_1.m4a b/web/dungeon/audio/sfx/echo_1.m4a index 57d7e9c1..bc19b8c0 100644 Binary files a/web/dungeon/audio/sfx/echo_1.m4a and b/web/dungeon/audio/sfx/echo_1.m4a differ diff --git a/web/dungeon/audio/sfx/echo_1.ogg b/web/dungeon/audio/sfx/echo_1.ogg index f08b32c2..7acc55cc 100644 Binary files a/web/dungeon/audio/sfx/echo_1.ogg and b/web/dungeon/audio/sfx/echo_1.ogg differ diff --git a/web/dungeon/audio/sfx/echo_2.m4a b/web/dungeon/audio/sfx/echo_2.m4a index 356307fc..8cde2331 100644 Binary files a/web/dungeon/audio/sfx/echo_2.m4a and b/web/dungeon/audio/sfx/echo_2.m4a differ diff --git a/web/dungeon/audio/sfx/echo_2.ogg b/web/dungeon/audio/sfx/echo_2.ogg index 451e9133..796181d9 100644 Binary files a/web/dungeon/audio/sfx/echo_2.ogg and b/web/dungeon/audio/sfx/echo_2.ogg differ diff --git a/web/dungeon/audio/sfx/echo_3.m4a b/web/dungeon/audio/sfx/echo_3.m4a index e8618e02..d9d6ee16 100644 Binary files a/web/dungeon/audio/sfx/echo_3.m4a and b/web/dungeon/audio/sfx/echo_3.m4a differ diff --git a/web/dungeon/audio/sfx/echo_3.ogg b/web/dungeon/audio/sfx/echo_3.ogg index 5e73fdf7..80031538 100644 Binary files a/web/dungeon/audio/sfx/echo_3.ogg and b/web/dungeon/audio/sfx/echo_3.ogg differ diff --git a/web/dungeon/audio/sfx/echo_4.m4a b/web/dungeon/audio/sfx/echo_4.m4a index 7b7eebaa..0384ee46 100644 Binary files a/web/dungeon/audio/sfx/echo_4.m4a and b/web/dungeon/audio/sfx/echo_4.m4a differ diff --git a/web/dungeon/audio/sfx/echo_4.ogg b/web/dungeon/audio/sfx/echo_4.ogg index 5c10ff1a..0344003d 100644 Binary files a/web/dungeon/audio/sfx/echo_4.ogg and b/web/dungeon/audio/sfx/echo_4.ogg differ diff --git a/web/dungeon/audio/sfx/enemy_die.m4a b/web/dungeon/audio/sfx/enemy_die.m4a index e845f9a6..0b22023c 100644 Binary files a/web/dungeon/audio/sfx/enemy_die.m4a and b/web/dungeon/audio/sfx/enemy_die.m4a differ diff --git a/web/dungeon/audio/sfx/enemy_die.ogg b/web/dungeon/audio/sfx/enemy_die.ogg index 9bd0c41b..237905fb 100644 Binary files a/web/dungeon/audio/sfx/enemy_die.ogg and b/web/dungeon/audio/sfx/enemy_die.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_bug.m4a b/web/dungeon/audio/sfx/mob_step_bug.m4a index 8f62d422..2fef2016 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug.m4a and b/web/dungeon/audio/sfx/mob_step_bug.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_bug.ogg b/web/dungeon/audio/sfx/mob_step_bug.ogg index b015e115..004503eb 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug.ogg and b/web/dungeon/audio/sfx/mob_step_bug.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_bug_c.m4a b/web/dungeon/audio/sfx/mob_step_bug_c.m4a index 5b1b4d7b..df8f2b9a 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug_c.m4a and b/web/dungeon/audio/sfx/mob_step_bug_c.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_bug_c.ogg b/web/dungeon/audio/sfx/mob_step_bug_c.ogg index 0948cfd4..be32304f 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug_c.ogg and b/web/dungeon/audio/sfx/mob_step_bug_c.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_bug_l.m4a b/web/dungeon/audio/sfx/mob_step_bug_l.m4a index 7cc36e9b..ab1f380d 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug_l.m4a and b/web/dungeon/audio/sfx/mob_step_bug_l.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_bug_l.ogg b/web/dungeon/audio/sfx/mob_step_bug_l.ogg index 27a65094..b191e9f4 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug_l.ogg and b/web/dungeon/audio/sfx/mob_step_bug_l.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_bug_r.m4a b/web/dungeon/audio/sfx/mob_step_bug_r.m4a index bc0bdd7c..615986d4 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug_r.m4a and b/web/dungeon/audio/sfx/mob_step_bug_r.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_bug_r.ogg b/web/dungeon/audio/sfx/mob_step_bug_r.ogg index de1dc63a..e1bd77e0 100644 Binary files a/web/dungeon/audio/sfx/mob_step_bug_r.ogg and b/web/dungeon/audio/sfx/mob_step_bug_r.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy.m4a b/web/dungeon/audio/sfx/mob_step_heavy.m4a index b3be4d00..cd56cb5a 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy.m4a and b/web/dungeon/audio/sfx/mob_step_heavy.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy.ogg b/web/dungeon/audio/sfx/mob_step_heavy.ogg index 92666773..984f15e2 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy.ogg and b/web/dungeon/audio/sfx/mob_step_heavy.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy_c.m4a b/web/dungeon/audio/sfx/mob_step_heavy_c.m4a index b960688b..edf6d7d1 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy_c.m4a and b/web/dungeon/audio/sfx/mob_step_heavy_c.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy_c.ogg b/web/dungeon/audio/sfx/mob_step_heavy_c.ogg index cc184115..b6c5ddde 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy_c.ogg and b/web/dungeon/audio/sfx/mob_step_heavy_c.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy_l.m4a b/web/dungeon/audio/sfx/mob_step_heavy_l.m4a index a7331692..fcd1e378 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy_l.m4a and b/web/dungeon/audio/sfx/mob_step_heavy_l.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy_l.ogg b/web/dungeon/audio/sfx/mob_step_heavy_l.ogg index 715c4463..270acc80 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy_l.ogg and b/web/dungeon/audio/sfx/mob_step_heavy_l.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy_r.m4a b/web/dungeon/audio/sfx/mob_step_heavy_r.m4a index b4be3597..b2d6af3d 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy_r.m4a and b/web/dungeon/audio/sfx/mob_step_heavy_r.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_heavy_r.ogg b/web/dungeon/audio/sfx/mob_step_heavy_r.ogg index 084e1f89..a4284f61 100644 Binary files a/web/dungeon/audio/sfx/mob_step_heavy_r.ogg and b/web/dungeon/audio/sfx/mob_step_heavy_r.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_small.m4a b/web/dungeon/audio/sfx/mob_step_small.m4a index acb1a162..cea0d410 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small.m4a and b/web/dungeon/audio/sfx/mob_step_small.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_small.ogg b/web/dungeon/audio/sfx/mob_step_small.ogg index 8e63ad0b..70fe56ea 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small.ogg and b/web/dungeon/audio/sfx/mob_step_small.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_small_c.m4a b/web/dungeon/audio/sfx/mob_step_small_c.m4a index 702609e0..c85c0050 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small_c.m4a and b/web/dungeon/audio/sfx/mob_step_small_c.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_small_c.ogg b/web/dungeon/audio/sfx/mob_step_small_c.ogg index ee27d04a..799879fb 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small_c.ogg and b/web/dungeon/audio/sfx/mob_step_small_c.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_small_l.m4a b/web/dungeon/audio/sfx/mob_step_small_l.m4a index 5777ee6e..1abaae4f 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small_l.m4a and b/web/dungeon/audio/sfx/mob_step_small_l.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_small_l.ogg b/web/dungeon/audio/sfx/mob_step_small_l.ogg index 0ad9989c..ffdd0344 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small_l.ogg and b/web/dungeon/audio/sfx/mob_step_small_l.ogg differ diff --git a/web/dungeon/audio/sfx/mob_step_small_r.m4a b/web/dungeon/audio/sfx/mob_step_small_r.m4a index 4395ef41..39d26342 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small_r.m4a and b/web/dungeon/audio/sfx/mob_step_small_r.m4a differ diff --git a/web/dungeon/audio/sfx/mob_step_small_r.ogg b/web/dungeon/audio/sfx/mob_step_small_r.ogg index 1afd1878..b574d4e9 100644 Binary files a/web/dungeon/audio/sfx/mob_step_small_r.ogg and b/web/dungeon/audio/sfx/mob_step_small_r.ogg differ diff --git a/web/dungeon/audio/sfx/stairs_down.m4a b/web/dungeon/audio/sfx/stairs_down.m4a index 94018ecf..387542c4 100644 Binary files a/web/dungeon/audio/sfx/stairs_down.m4a and b/web/dungeon/audio/sfx/stairs_down.m4a differ diff --git a/web/dungeon/audio/sfx/stairs_down.ogg b/web/dungeon/audio/sfx/stairs_down.ogg index 3f56a7eb..10a26147 100644 Binary files a/web/dungeon/audio/sfx/stairs_down.ogg and b/web/dungeon/audio/sfx/stairs_down.ogg differ diff --git a/web/dungeon/audio/sfx/tel_bug.m4a b/web/dungeon/audio/sfx/tel_bug.m4a index ee927960..5c93eb66 100644 Binary files a/web/dungeon/audio/sfx/tel_bug.m4a and b/web/dungeon/audio/sfx/tel_bug.m4a differ diff --git a/web/dungeon/audio/sfx/tel_bug.ogg b/web/dungeon/audio/sfx/tel_bug.ogg index 7c285338..628a1755 100644 Binary files a/web/dungeon/audio/sfx/tel_bug.ogg and b/web/dungeon/audio/sfx/tel_bug.ogg differ diff --git a/web/dungeon/audio/sfx/tel_bug_l.m4a b/web/dungeon/audio/sfx/tel_bug_l.m4a index 176052df..64e2a589 100644 Binary files a/web/dungeon/audio/sfx/tel_bug_l.m4a and b/web/dungeon/audio/sfx/tel_bug_l.m4a differ diff --git a/web/dungeon/audio/sfx/tel_bug_l.ogg b/web/dungeon/audio/sfx/tel_bug_l.ogg index 7d900219..902eadfb 100644 Binary files a/web/dungeon/audio/sfx/tel_bug_l.ogg and b/web/dungeon/audio/sfx/tel_bug_l.ogg differ diff --git a/web/dungeon/audio/sfx/tel_bug_r.m4a b/web/dungeon/audio/sfx/tel_bug_r.m4a index 0ee84772..bf4b7803 100644 Binary files a/web/dungeon/audio/sfx/tel_bug_r.m4a and b/web/dungeon/audio/sfx/tel_bug_r.m4a differ diff --git a/web/dungeon/audio/sfx/tel_bug_r.ogg b/web/dungeon/audio/sfx/tel_bug_r.ogg index 281e2a4d..51d032f7 100644 Binary files a/web/dungeon/audio/sfx/tel_bug_r.ogg and b/web/dungeon/audio/sfx/tel_bug_r.ogg differ diff --git a/web/dungeon/audio/sfx/tel_heavy.m4a b/web/dungeon/audio/sfx/tel_heavy.m4a index f5731d76..97fb6142 100644 Binary files a/web/dungeon/audio/sfx/tel_heavy.m4a and b/web/dungeon/audio/sfx/tel_heavy.m4a differ diff --git a/web/dungeon/audio/sfx/tel_heavy.ogg b/web/dungeon/audio/sfx/tel_heavy.ogg index 87e5703f..f9c27b61 100644 Binary files a/web/dungeon/audio/sfx/tel_heavy.ogg and b/web/dungeon/audio/sfx/tel_heavy.ogg differ diff --git a/web/dungeon/audio/sfx/tel_heavy_l.m4a b/web/dungeon/audio/sfx/tel_heavy_l.m4a index a24604e8..674695f3 100644 Binary files a/web/dungeon/audio/sfx/tel_heavy_l.m4a and b/web/dungeon/audio/sfx/tel_heavy_l.m4a differ diff --git a/web/dungeon/audio/sfx/tel_heavy_l.ogg b/web/dungeon/audio/sfx/tel_heavy_l.ogg index 4138371e..94d198f7 100644 Binary files a/web/dungeon/audio/sfx/tel_heavy_l.ogg and b/web/dungeon/audio/sfx/tel_heavy_l.ogg differ diff --git a/web/dungeon/audio/sfx/tel_heavy_r.m4a b/web/dungeon/audio/sfx/tel_heavy_r.m4a index 55324487..37a77a09 100644 Binary files a/web/dungeon/audio/sfx/tel_heavy_r.m4a and b/web/dungeon/audio/sfx/tel_heavy_r.m4a differ diff --git a/web/dungeon/audio/sfx/tel_heavy_r.ogg b/web/dungeon/audio/sfx/tel_heavy_r.ogg index 622bc4b3..2e94c218 100644 Binary files a/web/dungeon/audio/sfx/tel_heavy_r.ogg and b/web/dungeon/audio/sfx/tel_heavy_r.ogg differ diff --git a/web/dungeon/audio/sfx/tel_small.m4a b/web/dungeon/audio/sfx/tel_small.m4a index f213b4ba..b3e25b19 100644 Binary files a/web/dungeon/audio/sfx/tel_small.m4a and b/web/dungeon/audio/sfx/tel_small.m4a differ diff --git a/web/dungeon/audio/sfx/tel_small.ogg b/web/dungeon/audio/sfx/tel_small.ogg index 127eb961..c4ce60b2 100644 Binary files a/web/dungeon/audio/sfx/tel_small.ogg and b/web/dungeon/audio/sfx/tel_small.ogg differ diff --git a/web/dungeon/audio/sfx/tel_small_l.m4a b/web/dungeon/audio/sfx/tel_small_l.m4a index b9bbf298..83f24079 100644 Binary files a/web/dungeon/audio/sfx/tel_small_l.m4a and b/web/dungeon/audio/sfx/tel_small_l.m4a differ diff --git a/web/dungeon/audio/sfx/tel_small_l.ogg b/web/dungeon/audio/sfx/tel_small_l.ogg index 0d500406..120d29a7 100644 Binary files a/web/dungeon/audio/sfx/tel_small_l.ogg and b/web/dungeon/audio/sfx/tel_small_l.ogg differ diff --git a/web/dungeon/audio/sfx/tel_small_r.m4a b/web/dungeon/audio/sfx/tel_small_r.m4a index 8ad706e8..291dbf4d 100644 Binary files a/web/dungeon/audio/sfx/tel_small_r.m4a and b/web/dungeon/audio/sfx/tel_small_r.m4a differ diff --git a/web/dungeon/audio/sfx/tel_small_r.ogg b/web/dungeon/audio/sfx/tel_small_r.ogg index 0a073b0f..bcb66d0f 100644 Binary files a/web/dungeon/audio/sfx/tel_small_r.ogg and b/web/dungeon/audio/sfx/tel_small_r.ogg differ diff --git a/web/dungeon/gen/gen_sfx2.py b/web/dungeon/gen/gen_sfx2.py index e6565c6c..53d6673e 100644 --- a/web/dungeon/gen/gen_sfx2.py +++ b/web/dungeon/gen/gen_sfx2.py @@ -5,6 +5,8 @@ SA3 无 HF token 跳过（research/new8/batch1-assets-report.md §4 明示： 「dungeon 回声/脚步是核心动词，开发联调前建议先补产…或临时用 SYN 兜底」），本脚本即该兜底。 若日后补产 SA3 版，直接覆盖同名文件即可（id 不变、游戏层零改动）。 +※ 2026-09-02 更新：SA3 版已由 gen_sfx_sa3.py 补产并 A-B 后同名替换本脚本产出（28 文件； +本脚本重跑即回退 SYN 兜底版，原版备份 tools/asset_fix_backup/sa3_ab_20260902/dungeon/）。 25 个 id / 25×2 文件（mob_step_*/tel_* 的 l/r 为 mono 源经 ffmpeg pan 派生预 pan 立体声）： echo_1..4 正前回声 4 距离档：贴身亮短 → 远处暗长（stone 反射；门/水材质现有生成器 diff --git a/web/dungeon/gen/gen_sfx_sa3.py b/web/dungeon/gen/gen_sfx_sa3.py new file mode 100644 index 00000000..8fd19297 --- /dev/null +++ b/web/dungeon/gen/gen_sfx_sa3.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""《回声地窟》Stable Audio 3 small-sfx 音效后处理（SA3 补产批次，16-dungeon.md §9.2）。 + +批次 2 的 25 个 id 此前为 SYN 兜底（gen_sfx2.py，§9.2 降级条款）；本脚本产 SA3 候选版， +与在役 SYN 版做 A-B（响度/时长/包络对齐），质量明显更好才同名晋级 audio/sfx/（游戏层零改动）。 +prompt 出自 §9.2 表（echo 距离档/tel 前摇的英文稿按中文说明补全）。 + +生成（GPU，flock 串行；模型 gated 走 hf-mirror，见 env.d/sfx.sh）： + python3 gen_sfx_sa3.py --emit # 打印生成命令，原样执行 +生成 13 条 mono 源原始 wav 到 gen/staging/sfx_sa3/dungeon_.wav（顺序 = 下表 idx）。 +后处理：修整/去直流/8ms cos² 淡化/源真峰值 -4dBFS（同在役 SYN 档）→ +mob_step_*/tel_* 派生预 pan 立体声（l=能量只在 L / r=只在 R / c=双声道同相复制， +工艺同 gen_sfx2.py pan_derive/stereo_dup）→ +ogg libvorbis q4 + m4a aac 128k（highpass=f=30，前后垫 60/120ms 静音）→ +解码回读 QC：逐声道真峰值 <= -3dBFS、|直流| <= 0.001；超限降增益 + 首轮重编加 +18kHz 低通，最多 3 轮补偿。候选落 gen/staging/sfx_sa3/out/（不入库），A-B 后晋级。 + +用法：.venv-tts/bin/python web/dungeon/gen/gen_sfx_sa3.py [--emit] +""" +import os +import subprocess +import sys + +import numpy as np +import soundfile as sf + +SR = 44100 +HERE = os.path.dirname(os.path.abspath(__file__)) +STAGING = os.path.join(HERE, "staging", "sfx_sa3") +OUT = os.path.join(STAGING, "out") +DST = os.path.join(HERE, "..", "audio", "sfx") +SA3 = "/home/xsl/SoundGame/tools/stable-audio-3" +FFMPEG = "/home/xsl/SoundGame/tools/bin/ffmpeg" +SEED = 20260902 + +# (id, gen 秒, prompt) —— §9.2 标注 SA3 的 13 个 mono 源；定档全部 -4dBFS（同在役 SYN） +ITEMS = [ + ("echo_1", 2.5, "close stone wall echo slap-back, bright, short, dry"), + ("echo_2", 3.0, "mid distance cave echo, moderate tail"), + ("echo_3", 3.5, "far cave echo, darker longer"), + ("echo_4", 4.5, "very far cavern echo, long dark tail, huge space"), + ("stairs_down", 4.0, "descending old metal stairs, deep cavern reverb"), + ("atk_hit", 2.0, "blunt impact on hard insect chitin shell, dull thud hit"), + ("enemy_die", 2.5, "monster body collapsing on stone floor, heavy fall thud and debris"), + ("mob_step_small", 2.5, "small beast skittering steps, dry"), + ("mob_step_bug", 2.5, "insect chitin crawling clicks"), + ("mob_step_heavy", 2.5, "heavy armored thudding steps"), + ("tel_small", 2.5, "small beast pounce hiss, short aggressive screech"), + ("tel_bug", 2.5, "insect legs chitin rubbing scrape, creepy skitter"), + ("tel_heavy", 2.5, "huge monster stomp and low growl, heavy rumble"), +] +# pan 派生表（与在役文件集一致：mob_step 有 _c，tel 无 _c） +DERIVE = {} +for _n in ["mob_step_small", "mob_step_bug", "mob_step_heavy"]: + DERIVE[_n] = ["_l", "_c", "_r"] +for _n in ["tel_small", "tel_bug", "tel_heavy"]: + DERIVE[_n] = ["_l", "_r"] + +# 事件窗长度（秒）：raw 10.59s 定长 latent 床上按自响起点截取；None=按 -50dBFS 活动窗 +# 自然修剪。echo 档长 0.65<1.25<2.0<2.8 保持与在役 SYN 档（0.58/1.08/1.68/2.38）同序递增 +# ——回声尾长=墙距是本作核心信息，A-B 不满足递增即不晋级。 +WIN = { + "echo_1": 0.65, "echo_2": 1.25, "echo_3": 2.0, "echo_4": 2.8, + "stairs_down": None, "atk_hit": None, "enemy_die": 1.7, + "mob_step_small": 0.45, "mob_step_bug": 0.55, "mob_step_heavy": 0.65, + "tel_small": 1.0, "tel_bug": 1.1, "tel_heavy": 1.3, +} +# 响度对齐：候选按在役 SYN 版（同名 ogg）RMS 定档（-4dBFS 峰值帽不变），保游戏混音参数 +SYN_DIR = os.path.join(HERE, "..", "audio", "sfx") + + +def emit(): + durs = " ".join(str(d) for _, d, *_ in ITEMS) + ps = " ".join(f'"{p}"' for _, _, p, *_ in ITEMS) + print(f"cd {SA3} && flock /home/xsl/SoundGame/.gpu.lock .venv/bin/stable-audio " + f"--model small-sfx --steps 8 --seed {SEED} " + f"-p {ps} --duration {durs} -o {STAGING}/dungeon.wav") + + +def trim_oneshot(y, floor_db=-50.0, head_ms=50, tail_ms=140, max_s=None): + env = np.max(np.abs(y), axis=1) + act = np.where(env > 10 ** (floor_db / 20))[0] + a = max(0, act[0] - int(SR * head_ms / 1000)) + b = min(len(y), act[-1] + int(SR * tail_ms / 1000)) + if max_s: + b = min(b, a + int(SR * max_s)) + return y[a:b] + + +def prep(y, target_db=-4.0, win_len=None, syn_rms_db=None): + if y.ndim == 1: + y = y[:, None] + if win_len is not None: + # 自响起点：10ms 窗包络首个超过 (峰值-20dB) 的位置，提前 20ms 起 + w = int(SR * 0.01) + nfr = len(y) // w + eb = 20 * np.log10(np.max(np.abs(y[:nfr * w, 0]).reshape(nfr, w), axis=1) + 1e-9) + pk = 20 * np.log10(np.max(np.abs(y[:, 0])) + 1e-9) + on = np.where(eb > pk - 20)[0][0] * w + a = max(0, on - int(SR * 0.02)) + b = min(len(y), a + int(SR * win_len)) + y = y[a:b] + nf2 = max(1, int(SR * 0.050)) # 窗切在持续声中，切点 50ms 缓出 + y[-nf2:] *= (0.5 * (1 + np.cos(np.pi * np.arange(nf2) / nf2)))[:, None] + else: + y = trim_oneshot(y) + y = y - np.mean(y, axis=0, keepdims=True) # 去直流 + nf = max(1, int(SR * 0.008)) + fade = 0.5 * (1 - np.cos(np.pi * np.arange(nf) / nf)) + y[:nf] *= fade[:, None] # 8ms cos² 淡入淡出 + y[-nf:] *= fade[::-1][:, None] + y = y / (np.max(np.abs(y)) + 1e-9) * 10 ** (target_db / 20) + pad_h = np.zeros((int(SR * 0.06), y.shape[1])) # 编码器软着陆（同 gen_sfx2） + pad_t = np.zeros((int(SR * 0.12), y.shape[1])) + y = np.concatenate([pad_h, y, pad_t]) + if syn_rms_db is not None: # A-B 响度对齐（峰值帽 -4 不变）。 + rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9) # 垫后计算：与解码回读同口径 + s = min(10 ** ((syn_rms_db - rms) / 20), 1.0) + y = y * s + return y + + +def _syn_metrics(name): + """在役 SYN 版 ogg 解码回读（时长, RMS）——响度对齐基准。""" + p = os.path.join(SYN_DIR, name + ".ogg") + subprocess.run([FFMPEG, "-y", "-v", "error", "-i", p, "-c:a", "pcm_f32le", "/tmp/_syn_qc.wav"], check=True) + y, sr = sf.read("/tmp/_syn_qc.wav", always_2d=True) + return y.shape[0] / sr, 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9) + + +def derive(name, mono): + """mono -> 预 pan 立体声（能量只在 L 或 R；c = 双声道同相复制），工艺同 gen_sfx2。""" + outs = [] + for suf in DERIVE.get(name, []): + if suf == "_l": + outs.append((name + suf, np.stack([mono[:, 0], np.zeros(len(mono))], axis=1))) + elif suf == "_r": + outs.append((name + suf, np.stack([np.zeros(len(mono)), mono[:, 0]], axis=1))) + elif suf == "_c": + outs.append((name + suf, np.stack([mono[:, 0], mono[:, 0]], axis=1))) + return outs + + +def encode(name, sig): + base = sig.astype(np.float64) + vol, lpf, dc_adj = 0.0, False, 0.0 + tmp = os.path.join(OUT, "_enc.wav") + chk = os.path.join(OUT, "_qc.wav") + res = {} + for attempt in range(4): # 首发 + 最多 3 轮补偿（首轮重编加 18kHz 低通） + x = base - dc_adj + if vol < 0: + x = x * (10 ** (vol / 20)) + sf.write(tmp, x.astype(np.float32), SR) + af = "highpass=f=30,lowpass=f=18000" if lpf else "highpass=f=30" + for ext, cmd in [("ogg", ["-c:a", "libvorbis", "-q:a", "4"]), + ("m4a", ["-c:a", "aac", "-b:a", "128k"])]: + out = os.path.join(OUT, f"{name}.{ext}") + subprocess.run([FFMPEG, "-y", "-v", "error", "-i", tmp, "-af", af, + "-ar", "44100", *cmd, out], check=True) + subprocess.run([FFMPEG, "-y", "-v", "error", "-i", out, "-c:a", "pcm_f32le", chk], check=True) + y, sr = sf.read(chk, always_2d=True) + tp = 20 * np.log10(np.max(np.abs(y)) + 1e-9) + dcv = float(np.mean(y.mean(axis=1))) + rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9) + ok_tp, ok_dc = tp <= -3.0, abs(dcv) <= 0.001 + res[ext] = (tp, dcv, rms, y.shape[0] / sr) + if not (ok_tp and ok_dc): + if not ok_tp: + vol -= (tp + 3.0) + 0.25 + lpf = True # 首轮重编起加 18kHz 低通压 MDCT 过冲 + if not ok_dc: + dc_adj -= dcv + break + else: + break + else: + raise AssertionError(f"{name}: 补偿 3 轮后仍超限 {res}") + return res + + +def main(): + os.makedirs(STAGING, exist_ok=True) + os.makedirs(OUT, exist_ok=True) + if "--emit" in sys.argv: + emit() + return + print(f"| {'id':18s} | {'时长s':>6s} | {'ogg峰dBFS':>9s} | {'m4a峰dBFS':>9s} | {'RMS dBFS':>8s} | {'SYN时长':>7s} | {'SYN RMS':>8s} | 直流 | 声道 |") + for i, (name, _d, _p) in enumerate(ITEMS): + raw = os.path.join(STAGING, f"dungeon_{i}.wav") + y, sr = sf.read(raw, always_2d=True, dtype="float64") + assert sr == SR, f"{name}: 采样率 {sr}" + mono = y[:, :1] + syn_dur, syn_rms_db = _syn_metrics(name) + # 响度对齐只做一次（按 mono 基准），派生 l/r/c 直接复制基准样本—— + # 与在役 SYN 派生工艺一致：l/r 文件 RMS = 基准 -3dB（能量单声道承载），逐声道振幅不变 + base_sig = prep(mono, win_len=WIN.get(name), syn_rms_db=syn_rms_db) + for eid, sig in [(name, base_sig)] + derive(name, base_sig): + res = encode(eid, sig) + tp_o, dc, rms, dur = res["ogg"] + tp_m = res["m4a"][0] + ch = 1 if eid == name else 2 + sd = f"{syn_dur:6.2f}" if eid == name else " -" + print(f"| {eid:18s} | {dur:6.2f} | {tp_o:9.2f} | {tp_m:9.2f} | {rms:8.2f} | {sd} | {syn_rms_db:7.2f} | {dc:+.4f} | {ch:d} |") + print("dungeon SA3 候选完成 -> gen/staging/sfx_sa3/out/（A-B 复核后晋级）") + + +if __name__ == "__main__": + main() diff --git a/web/radio/audio/sfx/amb_cab.m4a b/web/radio/audio/sfx/amb_cab.m4a new file mode 100644 index 00000000..4a17be46 Binary files /dev/null and b/web/radio/audio/sfx/amb_cab.m4a differ diff --git a/web/radio/audio/sfx/amb_cab.ogg b/web/radio/audio/sfx/amb_cab.ogg new file mode 100644 index 00000000..7b735eb1 Binary files /dev/null and b/web/radio/audio/sfx/amb_cab.ogg differ diff --git a/web/radio/audio/sfx/amb_kid_room.m4a b/web/radio/audio/sfx/amb_kid_room.m4a new file mode 100644 index 00000000..80624525 Binary files /dev/null and b/web/radio/audio/sfx/amb_kid_room.m4a differ diff --git a/web/radio/audio/sfx/amb_kid_room.ogg b/web/radio/audio/sfx/amb_kid_room.ogg new file mode 100644 index 00000000..c78b0aa3 Binary files /dev/null and b/web/radio/audio/sfx/amb_kid_room.ogg differ diff --git a/web/radio/audio/sfx/amb_rain_city.m4a b/web/radio/audio/sfx/amb_rain_city.m4a new file mode 100644 index 00000000..f3f4436d Binary files /dev/null and b/web/radio/audio/sfx/amb_rain_city.m4a differ diff --git a/web/radio/audio/sfx/amb_rain_city.ogg b/web/radio/audio/sfx/amb_rain_city.ogg new file mode 100644 index 00000000..d99bebd1 Binary files /dev/null and b/web/radio/audio/sfx/amb_rain_city.ogg differ diff --git a/web/radio/audio/sfx/amb_station.m4a b/web/radio/audio/sfx/amb_station.m4a new file mode 100644 index 00000000..e8523fac Binary files /dev/null and b/web/radio/audio/sfx/amb_station.m4a differ diff --git a/web/radio/audio/sfx/amb_station.ogg b/web/radio/audio/sfx/amb_station.ogg new file mode 100644 index 00000000..aa8a7d8b Binary files /dev/null and b/web/radio/audio/sfx/amb_station.ogg differ diff --git a/web/radio/audio/sfx/amb_street_rain.m4a b/web/radio/audio/sfx/amb_street_rain.m4a new file mode 100644 index 00000000..2045ddd6 Binary files /dev/null and b/web/radio/audio/sfx/amb_street_rain.m4a differ diff --git a/web/radio/audio/sfx/amb_street_rain.ogg b/web/radio/audio/sfx/amb_street_rain.ogg new file mode 100644 index 00000000..4e6a46bc Binary files /dev/null and b/web/radio/audio/sfx/amb_street_rain.ogg differ diff --git a/web/radio/audio/sfx/amb_studio.m4a b/web/radio/audio/sfx/amb_studio.m4a new file mode 100644 index 00000000..d414010c Binary files /dev/null and b/web/radio/audio/sfx/amb_studio.m4a differ diff --git a/web/radio/audio/sfx/amb_studio.ogg b/web/radio/audio/sfx/amb_studio.ogg new file mode 100644 index 00000000..65263c60 Binary files /dev/null and b/web/radio/audio/sfx/amb_studio.ogg differ diff --git a/web/radio/audio/sfx/arg_muffled.m4a b/web/radio/audio/sfx/arg_muffled.m4a new file mode 100644 index 00000000..e971e0a1 Binary files /dev/null and b/web/radio/audio/sfx/arg_muffled.m4a differ diff --git a/web/radio/audio/sfx/arg_muffled.ogg b/web/radio/audio/sfx/arg_muffled.ogg new file mode 100644 index 00000000..28d62ae1 Binary files /dev/null and b/web/radio/audio/sfx/arg_muffled.ogg differ diff --git a/web/radio/audio/sfx/brake_screech.m4a b/web/radio/audio/sfx/brake_screech.m4a new file mode 100644 index 00000000..e7ca1022 Binary files /dev/null and b/web/radio/audio/sfx/brake_screech.m4a differ diff --git a/web/radio/audio/sfx/brake_screech.ogg b/web/radio/audio/sfx/brake_screech.ogg new file mode 100644 index 00000000..1a95ec60 Binary files /dev/null and b/web/radio/audio/sfx/brake_screech.ogg differ diff --git a/web/radio/audio/sfx/line_connect.m4a b/web/radio/audio/sfx/line_connect.m4a new file mode 100644 index 00000000..1e2aee49 Binary files /dev/null and b/web/radio/audio/sfx/line_connect.m4a differ diff --git a/web/radio/audio/sfx/line_connect.ogg b/web/radio/audio/sfx/line_connect.ogg new file mode 100644 index 00000000..3ef100cb Binary files /dev/null and b/web/radio/audio/sfx/line_connect.ogg differ diff --git a/web/radio/audio/sfx/line_drop.m4a b/web/radio/audio/sfx/line_drop.m4a new file mode 100644 index 00000000..ba2d11a2 Binary files /dev/null and b/web/radio/audio/sfx/line_drop.m4a differ diff --git a/web/radio/audio/sfx/line_drop.ogg b/web/radio/audio/sfx/line_drop.ogg new file mode 100644 index 00000000..6d23627d Binary files /dev/null and b/web/radio/audio/sfx/line_drop.ogg differ diff --git a/web/radio/audio/sfx/onair_sting.m4a b/web/radio/audio/sfx/onair_sting.m4a new file mode 100644 index 00000000..45a51556 Binary files /dev/null and b/web/radio/audio/sfx/onair_sting.m4a differ diff --git a/web/radio/audio/sfx/onair_sting.ogg b/web/radio/audio/sfx/onair_sting.ogg new file mode 100644 index 00000000..de2f9b26 Binary files /dev/null and b/web/radio/audio/sfx/onair_sting.ogg differ diff --git a/web/radio/audio/sfx/phone_ring.m4a b/web/radio/audio/sfx/phone_ring.m4a index 1ee12c8a..1673b966 100644 Binary files a/web/radio/audio/sfx/phone_ring.m4a and b/web/radio/audio/sfx/phone_ring.m4a differ diff --git a/web/radio/audio/sfx/phone_ring.ogg b/web/radio/audio/sfx/phone_ring.ogg index 219fb484..8fc3be76 100644 Binary files a/web/radio/audio/sfx/phone_ring.ogg and b/web/radio/audio/sfx/phone_ring.ogg differ diff --git a/web/radio/audio/sfx/tape_static.m4a b/web/radio/audio/sfx/tape_static.m4a new file mode 100644 index 00000000..40ea48ee Binary files /dev/null and b/web/radio/audio/sfx/tape_static.m4a differ diff --git a/web/radio/audio/sfx/tape_static.ogg b/web/radio/audio/sfx/tape_static.ogg new file mode 100644 index 00000000..4c8e482e Binary files /dev/null and b/web/radio/audio/sfx/tape_static.ogg differ diff --git a/web/radio/gen/gen_sfx.py b/web/radio/gen/gen_sfx.py index 8b753765..4e82a7f6 100644 --- a/web/radio/gen/gen_sfx.py +++ b/web/radio/gen/gen_sfx.py @@ -9,6 +9,9 @@ R2 起共 3 件（其余 11 件为 Stable Audio 3 条目，HF token 未配置， 循环周期 5.2s ≥4s 烘进文件本身；无缝 loop 规格：逐声道去直流、 首尾**不做淡化**、循环边界静音校验；游戏侧 src.loop=true 循环播放， 接听即停——SA3 真铃声到位后可整文件替换，接口不变） + ※ 2026-09-02 更新：SA3 12 件已由 gen_sfx_sa3.py 补产落地，phone_ring 已被 SA3 真铃声 + 同名替换（本脚本重跑会以 SYN 版覆盖回去——如需回退再跑本脚本即可，原 SYN 备份在 + tools/asset_fix_backup/sa3_ab_20260902/radio/）。 可复现：本脚本入库，一键出 web/radio/audio/sfx/ui_{confirm,menu,phone_ring}.{ogg,m4a}； 中间 wav 在 gen/staging/sfx/（不入库）。规格：ogg libvorbis q4 44.1kHz + m4a aac 128k； diff --git a/web/radio/gen/gen_sfx_sa3.py b/web/radio/gen/gen_sfx_sa3.py new file mode 100644 index 00000000..e549a892 --- /dev/null +++ b/web/radio/gen/gen_sfx_sa3.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""《不眠频率》Stable Audio 3 small-sfx 音效后处理（SA3 补产批次，15-radio.md §5.2）。 + +生成（GPU，flock 串行；模型 gated 走 hf-mirror，见 env.d/sfx.sh）： + python3 gen_sfx_sa3.py --emit # 打印生成命令，原样执行 +生成 12 条原始 wav 到 gen/staging/sfx_sa3/radio_.wav（顺序 = 下表 idx）。 +本脚本随后做后处理：修整/去直流/8ms cos² 淡化（loop 类只去直流不淡化）/ +响度定档（源真峰值或 RMS 定档，escape amb 床 -26dBFS RMS 先例）→ +ogg libvorbis q4 + m4a aac 128k（highpass=f=30，one-shot 前后垫 60/120ms 静音）→ +解码回读 QC：逐声道真峰值 <= -3dBFS、|直流| <= 0.001；超限降增益 + 首轮重编 +加 18kHz 低通，最多 3 轮补偿。loop 类（phone_ring）另验解码首/末样本 <= -60dBFS。 +候选 ogg/m4a 落 gen/staging/sfx_sa3/out/（不入库），人工/A-B 复核后晋级 audio/sfx/。 + +prompt 均出自 15-radio.md §5.2（英文原稿）；定档理由： + onair_sting/brake_screech/phone_ring -4（台标/危机拍/铃声，主层次） + line_connect/line_drop -7（line_down 细节咔哒，subtle/soft） + arg_muffled -8（C3 信息层线索，垫在人声下但可闻） + tape_static -5（p9 直面道具声） + amb_* RMS -26 / 峰值帽 -8（通话侧环境床，escape amb_cabinet/-26dBFS 先例） + +用法：.venv-tts/bin/python web/radio/gen/gen_sfx_sa3.py [--emit] +""" +import os +import subprocess +import sys + +import numpy as np +import soundfile as sf + +SR = 44100 +HERE = os.path.dirname(os.path.abspath(__file__)) +STAGING = os.path.join(HERE, "staging", "sfx_sa3") +OUT = os.path.join(STAGING, "out") +DST = os.path.join(HERE, "..", "audio", "sfx") +SA3 = "/home/xsl/SoundGame/tools/stable-audio-3" +FFMPEG = "/home/xsl/SoundGame/tools/bin/ffmpeg" +SEED = 20260902 + +# (id, gen 秒, prompt, 模式 oneshot/loop/amb, 定档 dBFS) +ITEMS = [ + ("onair_sting", 3.0, "Short broadcast on-air light sting, warm chime with vintage radio click, 1 second", "oneshot", -4.0), + ("phone_ring", 6.0, "Old studio telephone bell, two-strike bright metallic ring, clean", "loop", -4.0), + ("line_connect", 3.0, "Telephone line connect click and short dial tone fade, subtle", "oneshot", -7.0), + ("line_drop", 3.0, "Telephone hang-up click with short busy tone, soft", "oneshot", -7.0), + ("amb_rain_city", 10.0, "Heavy rain on city street at night, distant traffic hiss, seamless loop", "amb", -26.0), + ("amb_station", 10.0, "Empty train station hall at night, distant echo and rolling suitcase, seamless loop", "amb", -26.0), + ("amb_kid_room", 10.0, "Quiet child bedroom at night, soft lamp buzz, muffled house silence, seamless loop", "amb", -26.0), + ("arg_muffled", 5.0, "Muffled argument behind a door, dull thud of door slam, indistinct voices, 4 seconds", "oneshot", -8.0), + ("amb_street_rain", 10.0, "Rain at an empty bus stop at night, distant wet road traffic, seamless loop", "amb", -26.0), + ("amb_cab", 10.0, "Inside a moving truck cab at night, steady engine hum, rhythmic windshield wipers, rain outside, seamless loop", "amb", -26.0), + ("brake_screech", 3.0, "Sudden brake screech on wet asphalt, brief tire friction, 2 seconds", "oneshot", -4.0), + ("tape_static", 4.0, "Old cassette tape static and tape hiss with a click stop, 3 seconds", "oneshot", -5.0), + # amb_studio：文档 §5.2 有、game.js 未接线（直播间底噪由 bgm/onair 承担）——产出备用 + ("amb_studio", 10.0, "Quiet radio studio room tone, console hum and soft air conditioning, seamless loop", "amb", -26.0), +] +RING_PERIOD = 5.2 # RING 循环周期 >=4s（设计 §5.2），烘进文件 + +# 事件窗（raw 16.16s 定长 latent 床上截取目标事件；None=按 -50dBFS 活动窗自然修剪）。 +# 窗中心取实测事件：arg_muffled 摔门峰 @5.00s、brake_screech 最强吱咬 @3.91s、 +# tape_static 停带咔哒 @3.72s（seed 20260902 实测）。 +WIN = { + "arg_muffled": (1.0, 5.4), # 争吵声垫 -> 摔门 -> 余渍，≈设计 4s + "brake_screech": (3.6, 5.8), # 最强急刹吱咬 2.2s（设计 2s 危机拍） + "tape_static": (0.8, 4.15), # 电流嘶声 -> 停带咔哒，≈设计 3s +} + + +def emit(): + durs = " ".join(str(d) for _, d, *_ in ITEMS) + ps = " ".join(f'"{p}"' for _, _, p, *_ in ITEMS) + print(f"cd {SA3} && flock /home/xsl/SoundGame/.gpu.lock .venv/bin/stable-audio " + f"--model small-sfx --steps 8 --seed {SEED} " + f"-p {ps} --duration {durs} -o {STAGING}/radio.wav") + + +def trim_oneshot(y, floor_db=-50.0, head_ms=50, tail_ms=140, max_s=None): + env = np.max(np.abs(y), axis=1) + act = np.where(env > 10 ** (floor_db / 20))[0] + a = max(0, act[0] - int(SR * head_ms / 1000)) + b = min(len(y), act[-1] + int(SR * tail_ms / 1000)) + if max_s: + b = min(b, a + int(SR * max_s)) + return y[a:b] + + +def prep_oneshot(y, target_db, win=None): + if win: + a, b = int(SR * win[0]), int(SR * win[1]) + y = y[a:b] + nf2 = max(1, int(SR * 0.050)) # 窗切在持续声中，切点 50ms 缓出 + y[-nf2:] *= (0.5 * (1 + np.cos(np.pi * np.arange(nf2) / nf2)))[:, None] + else: + y = trim_oneshot(y) + y = y - np.mean(y, axis=0, keepdims=True) # 去直流 + nf = max(1, int(SR * 0.008)) + fade = 0.5 * (1 - np.cos(np.pi * np.arange(nf) / nf)) + y[:nf] *= fade[:, None] # 8ms cos² 淡入淡出 + y[-nf:] *= fade[::-1][:, None] + y = y / (np.max(np.abs(y)) + 1e-9) * 10 ** (target_db / 20) + pad_h = np.zeros((int(SR * 0.06), y.shape[1])) # 编码器软着陆 + pad_t = np.zeros((int(SR * 0.12), y.shape[1])) + return np.concatenate([pad_h, y, pad_t]) + + +def prep_loop(y, target_db, period_s): + """铃声循环垫：修剪衰减尾 -> 循环边界 10ms 压零（不做整段淡化）-> 补静音到周期。""" + y = trim_oneshot(y, floor_db=-50.0, head_ms=30, tail_ms=30, max_s=2.6) + y = y - np.mean(y, axis=0, keepdims=True) + nf = max(1, int(SR * 0.010)) + y[:nf] *= np.linspace(0, 1, nf)[:, None] # 仅边界压零，防 loop 爆点 + y[-nf:] *= np.linspace(1, 0, nf)[:, None] + y = y / (np.max(np.abs(y)) + 1e-9) * 10 ** (target_db / 20) + need = int(SR * period_s) + if len(y) < need: + y = np.concatenate([y, np.zeros((need - len(y), y.shape[1]))]) + return y + + +def prep_amb(y, rms_db, peak_db): + """环境床：只去直流，不淡化（游戏层 0.4s 重触发重叠掩接缝）；RMS 定档 + 峰值帽。""" + y = y - np.mean(y, axis=0, keepdims=True) + rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9) + peak = 20 * np.log10(np.max(np.abs(y)) + 1e-9) + s = min(10 ** ((rms_db - rms) / 20), 10 ** ((peak_db - peak) / 20)) + return y * s + + +def encode(name, sig, loop=False): + base = sig.astype(np.float64) + vol, lpf, dc_adj = 0.0, False, 0.0 + edge_n = 0 # loop 边界超限时逐轮加宽压零区 + tmp = os.path.join(OUT, "_enc.wav") + chk = os.path.join(OUT, "_qc.wav") + res = {} + for attempt in range(4): # 首发 + 最多 3 轮补偿（首轮重编加 18kHz 低通） + x = base - dc_adj + if vol < 0: + x = x * (10 ** (vol / 20)) + if edge_n: + ne = int(SR * 0.005 * edge_n) + x[:ne] = 0.0 + x[-ne:] = 0.0 + sf.write(tmp, x.astype(np.float32), SR) + af = "highpass=f=30,lowpass=f=18000" if lpf else "highpass=f=30" + for ext, cmd in [("ogg", ["-c:a", "libvorbis", "-q:a", "4"]), + ("m4a", ["-c:a", "aac", "-b:a", "128k"])]: + out = os.path.join(OUT, f"{name}.{ext}") + subprocess.run([FFMPEG, "-y", "-v", "error", "-i", tmp, "-af", af, + "-ar", "44100", *cmd, out], check=True) + subprocess.run([FFMPEG, "-y", "-v", "error", "-i", out, "-c:a", "pcm_f32le", chk], check=True) + y, sr = sf.read(chk, always_2d=True) + tp = 20 * np.log10(np.max(np.abs(y)) + 1e-9) + dcv = float(np.mean(y.mean(axis=1))) + rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9) + ok_tp, ok_dc = tp <= -3.0, abs(dcv) <= 0.001 + ok_loop = True + if loop: + edge = 10 ** (-60.0 / 20) + ok_loop = (float(np.max(np.abs(y[0, :]))) <= edge + and float(np.max(np.abs(y[-1, :]))) <= edge) + res[ext] = (tp, dcv, rms, y.shape[0] / sr) + if not (ok_tp and ok_dc and ok_loop): + if not ok_tp: + vol -= (tp + 3.0) + 0.25 + lpf = True # 首轮重编起加 18kHz 低通压 MDCT 过冲 + if not ok_dc: + dc_adj -= dcv + if not ok_loop: + edge_n += 1 # 循环边界再压 5ms 零 + break + else: + break + else: + raise AssertionError(f"{name}: 补偿 3 轮后仍超限 {res}") + return res + + +def main(): + os.makedirs(STAGING, exist_ok=True) + os.makedirs(OUT, exist_ok=True) + if "--emit" in sys.argv: + emit() + return + print(f"| {'id':16s} | {'时长s':>6s} | {'ogg峰dBFS':>9s} | {'m4a峰dBFS':>9s} | {'RMS dBFS':>8s} | 直流 |") + for i, (name, _d, _p, mode, tgt) in enumerate(ITEMS): + raw = os.path.join(STAGING, f"radio_{i}.wav") + y, sr = sf.read(raw, always_2d=True, dtype="float64") + assert sr == SR, f"{name}: 采样率 {sr}" + if y.shape[1] == 1: + y = np.stack([y[:, 0], y[:, 0]], axis=1) + if mode == "oneshot": + sig = prep_oneshot(y, tgt, win=WIN.get(name)) + elif mode == "loop": + sig = prep_loop(y, tgt, RING_PERIOD) + else: + sig = prep_amb(y, tgt, -8.0) + res = encode(name, sig, loop=(mode == "loop")) + tp_o, dc, rms, dur = res["ogg"] + tp_m = res["m4a"][0] + print(f"| {name:16s} | {dur:6.2f} | {tp_o:9.2f} | {tp_m:9.2f} | {rms:8.2f} | {dc:+.4f} |") + print("radio SA3 候选完成 -> gen/staging/sfx_sa3/out/（晋级前先复核）") + + +if __name__ == "__main__": + main()