* Unify training and ONNX-export environment, lift PyTorch 1.13 pin
The historical PyTorch 1.13 pin existed because two bugs surfaced in newer
torch when exporting to ONNX. Both are now fixed; training and export can
share a single environment (PyTorch >= 2.0).
Bug 1 - WaveNet diffusion ONNX export crashes on PyTorch >= 2.0
modules/backbones/wavenet.py:83 used spec.squeeze(1), which the ONNX
tracer lowers to an onnx::If whose two branches have different ranks
(block0 Squeeze -> rank-3, block1 Identity -> rank-4). Shape inference
for the downstream Conv then fails with SymbolicValueError. Replaced
with spec[:, 0] - an unconditional rank-reducing gather, semantically
identical (eager max-diff = 0) and producing a clean ONNX graph.
Bug 2 - non-RoPE encoder ONNX inference fails at dynamic lengths
torch.nn.MultiheadAttention's multi_head_attention_forward gained an
SDPA-branched implementation in torch 2.0. The branching caused the
tracer to specialize tgt_len as a Python int constant and bake it into
the output Reshape, so a model traced at T=40 errored with
'requested shape:{40,2,32}' at any other length. The historical
comment blaming espnet_positional_embedding.py was incorrect: the
failure reproduces with a bare nn.MultiheadAttention and zero PE code,
and survives even with the sinusoidal PE path which never touches
the espnet module.
Routed both non-RoPE paths through the in-house manual attention
(MultiheadSelfAttentionWithRoPE with rotary_embed=None) that was
already used on the RoPE path. It is fully dynamic-safe and produces
identical eager output (max-diff 7e-7) at T=40/80/160.
Checkpoint compatibility
Manual attention uses state_dict key 'in_proj.weight' whereas
nn.MultiheadAttention used 'in_proj_weight'. Same shape and same
Q/K/V-stacked-along-dim-0 semantics; utils.load_ckpt now renames the
old key on load, so legacy ckpts continue to work with strict=True.
Diffusion graph simplification
Each diffusion sub-graph was simplified twice: once before
graph_extract_conditioner_projections and once after. The pre-surgery
pass is removed. The conditioner-projection extraction rewrites the
graph in a way that can collide with the first simplifier's node
ordering and make the merged model fail onnx topological-sort
validation downstream (a latent merge bug). The post-surgery
simplifier subsumes the dropped pass, so the final graph is unchanged
on the models that already worked, and the merge bug is avoided.
Applies to acoustic (main diffusion), variance (pitch and multi-
variance diffusions).
Dependency cleanup
- Removed requirements-onnx.txt entirely (training and export share
requirements.txt with PyTorch >= 2.0).
- Replaced onnxsim with onnxslim (>=0.1.93) via a thin
utils.onnx_helper.simplify_onnx wrapper. onnxslim is easier to
install across environments and has no native build chain.
- All torch.onnx.export calls stay on the TorchScript exporter that
utils.onnx_helper's graph surgery was written against. The dynamo
backend's availability differs across PyTorch versions: it first
shipped as a separate torch.onnx.dynamo_export API in 2.1, and
torch.onnx.export gained a 'dynamo' kwarg in 2.4 (default False,
flipped to True in 2.9). Versions 2.0-2.3 have no such kwarg. To
stay correct on all of them we probe
inspect.signature(torch.onnx.export) once at import time and only
pass dynamo=False when the kwarg exists - exposed as
utils.onnx_helper.TORCHSCRIPT_EXPORT_KWARGS, splatted into every
export call. Verified across torch 2.1 (no kwarg -> empty dict) and
2.8 (kwarg present -> dynamo=False forwarded).
- opset 15 -> 17. Verified with onnx.checker on all three model
families.
* Unify ONNX env and bump PyTorch to >=2.4
Remove instructions to use a separate environment and requirements-onnx.txt for ONNX export; docs now recommend using the same environment for training and ONNX export and installing dependencies via the Installation section. Update requirements.txt comment to require PyTorch >= 2.4.
* Clarify PyTorch and environment recommendations
Update documentation and requirements comments to clarify environment setup: add 'uv' to the recommended virtual environment options, explicitly recommend using the latest stable PyTorch release (>= 2.4.0) in GettingStarted.md, and remove a redundant paragraph about a unified training/ONNX environment. Also adjust the top comment in requirements.txt to state that PyTorch >= 2.4 is recommended rather than required.
* Bump ONNX requirement to >=1.21.0
Update requirements.txt to change the onnx constraint from ~=1.16.0 to >=1.21.0, allowing newer ONNX releases for compatibility with updated dependencies/features.
* Unpin MonkeyType in requirements
Remove the exact version constraint for MonkeyType in requirements.txt (changed from MonkeyType==23.3.0 to MonkeyType) to allow installation of newer/compatible releases and relax strict dependency pinning.
Update requirements.txt
Kakaru
2a8a013162