|
1 | 1 | """ |
2 | 2 | Conditional weight_norm import for RVC fork compatibility. |
3 | 3 |
|
4 | | -By default, uses the old-style ``torch.nn.utils.weight_norm`` which stores |
5 | | -weights as ``.weight_g`` / ``.weight_v`` — the format used by virtually every |
6 | | -RVC fork. This ensures trained models are interchangeable across projects. |
| 4 | +Uses ``torch.nn.utils.parametrizations.weight_norm`` (PyTorch 2.0+) by default, |
| 5 | +which stores weights as ``.parametrizations.weight.original0`` / ``original1``. |
| 6 | +This matches the approach used by Applio and Vietnamese-RVC. |
7 | 7 |
|
8 | | -Set the config option ``new_pytorch_weight_norm: true`` to switch to |
9 | | -``torch.nn.utils.parametrizations.weight_norm`` (PyTorch 2.0+), which stores |
10 | | -weights as ``.parametrizations.weight.original0`` / ``original1``. |
| 8 | +All checkpoints saved to disk use the old ``.weight_g`` / ``.weight_v`` format |
| 9 | +for maximum backward compatibility across RVC forks. Key conversion happens |
| 10 | +automatically at save/load boundaries. |
11 | 11 | """ |
12 | 12 |
|
13 | 13 | import logging |
14 | 14 |
|
15 | 15 | logger = logging.getLogger(__name__) |
16 | 16 |
|
17 | 17 | # ── Global flag ────────────────────────────────────────────────────────── |
18 | | -# False = old-style weight_norm (default, compatible with most RVC forks) |
19 | | -# True = new PyTorch 2.0+ parametrizations.weight_norm |
20 | | -_new_pytorch_mode = False |
| 18 | +# True = new PyTorch 2.0+ parametrizations.weight_norm (default, matches Applio/VRVC) |
| 19 | +# False = old-style weight_norm (legacy compatibility) |
| 20 | +_new_pytorch_mode = True |
21 | 21 |
|
22 | 22 |
|
23 | | -def configure_weight_norm(new_pytorch: bool = False) -> None: |
| 23 | +def configure_weight_norm(new_pytorch: bool = True) -> None: |
24 | 24 | """Set the weight_norm mode. Call once at startup before model creation.""" |
25 | 25 | global _new_pytorch_mode |
26 | 26 | _new_pytorch_mode = bool(new_pytorch) |
27 | | - mode = "new PyTorch 2.0+ parametrizations" if _new_pytorch_mode else "old-style (RVC fork compatible)" |
| 27 | + mode = "new PyTorch 2.0+ parametrizations" if _new_pytorch_mode else "old-style (legacy)" |
28 | 28 | logger.info(f"weight_norm mode: {mode}") |
29 | 29 |
|
30 | 30 |
|
@@ -94,60 +94,45 @@ def weight_norm_g(module): |
94 | 94 |
|
95 | 95 |
|
96 | 96 | # ── Key conversion helpers ─────────────────────────────────────────────── |
97 | | -# These are only needed when new_pytorch mode is active. |
| 97 | +# Always active regardless of mode. This ensures seamless interoperability |
| 98 | +# between old-format checkpoints and new-format models (and vice versa). |
98 | 99 |
|
99 | | -def needs_key_conversion() -> bool: |
100 | | - """Return True if checkpoint key conversion between old/new format is needed. |
101 | | -
|
102 | | - Conversion is needed only when the model uses new PyTorch parametrizations |
103 | | - but checkpoints are stored in old ``.weight_g``/``.weight_v`` format. |
104 | | - """ |
105 | | - return _new_pytorch_mode |
| 100 | +def _replace_keys_in_dict(d, old_part, new_part): |
| 101 | + """Recursively replace key substrings in a (possibly nested) dict.""" |
| 102 | + from collections import OrderedDict |
| 103 | + updated = OrderedDict() if isinstance(d, OrderedDict) else {} |
| 104 | + for key, value in d.items(): |
| 105 | + updated[( |
| 106 | + key.replace(old_part, new_part) if isinstance(key, str) else key |
| 107 | + )] = ( |
| 108 | + _replace_keys_in_dict(value, old_part, new_part) if isinstance(value, dict) else value |
| 109 | + ) |
| 110 | + return updated |
106 | 111 |
|
107 | 112 |
|
108 | 113 | def convert_old_to_new(state_dict): |
109 | 114 | """Convert old-style weight_norm keys to new parametrizations keys. |
110 | 115 |
|
111 | | - Only applies when ``new_pytorch_weight_norm`` is enabled. |
112 | | - No-op otherwise (model and checkpoint both use old format). |
113 | | - """ |
114 | | - if not _new_pytorch_mode: |
115 | | - return state_dict |
116 | | - |
117 | | - from collections import OrderedDict |
118 | | - |
119 | | - def _replace(d, old_part, new_part): |
120 | | - updated = OrderedDict() if isinstance(d, OrderedDict) else {} |
121 | | - for key, value in d.items(): |
122 | | - updated[key.replace(old_part, new_part) if isinstance(key, str) else key] = ( |
123 | | - _replace(value, old_part, new_part) if isinstance(value, dict) else value |
124 | | - ) |
125 | | - return updated |
| 116 | + Converts ``.weight_v`` → ``.parametrizations.weight.original1`` |
| 117 | + and ``.weight_g`` → ``.parametrizations.weight.original0`` |
126 | 118 |
|
127 | | - result = _replace(state_dict, ".weight_v", ".parametrizations.weight.original1") |
128 | | - result = _replace(result, ".weight_g", ".parametrizations.weight.original0") |
| 119 | + Always active — ensures old-format checkpoints load correctly into |
| 120 | + models that use the new PyTorch parametrization format. |
| 121 | + """ |
| 122 | + result = _replace_keys_in_dict(state_dict, ".weight_v", ".parametrizations.weight.original1") |
| 123 | + result = _replace_keys_in_dict(result, ".weight_g", ".parametrizations.weight.original0") |
129 | 124 | return result |
130 | 125 |
|
131 | 126 |
|
132 | 127 | def convert_new_to_old(state_dict): |
133 | 128 | """Convert new parametrizations keys to old-style weight_norm keys. |
134 | 129 |
|
135 | | - Only applies when ``new_pytorch_weight_norm`` is enabled. |
136 | | - No-op otherwise (model and checkpoint both use old format). |
137 | | - """ |
138 | | - if not _new_pytorch_mode: |
139 | | - return state_dict |
| 130 | + Converts ``.parametrizations.weight.original1`` → ``.weight_v`` |
| 131 | + and ``.parametrizations.weight.original0`` → ``.weight_g`` |
140 | 132 |
|
141 | | - from collections import OrderedDict |
142 | | - |
143 | | - def _replace(d, old_part, new_part): |
144 | | - updated = OrderedDict() if isinstance(d, OrderedDict) else {} |
145 | | - for key, value in d.items(): |
146 | | - updated[key.replace(old_part, new_part) if isinstance(key, str) else key] = ( |
147 | | - _replace(value, old_part, new_part) if isinstance(value, dict) else value |
148 | | - ) |
149 | | - return updated |
150 | | - |
151 | | - result = _replace(state_dict, ".parametrizations.weight.original1", ".weight_v") |
152 | | - result = _replace(result, ".parametrizations.weight.original0", ".weight_g") |
| 133 | + Always active — ensures saved checkpoints use the old format for |
| 134 | + maximum backward compatibility with other RVC forks. |
| 135 | + """ |
| 136 | + result = _replace_keys_in_dict(state_dict, ".parametrizations.weight.original1", ".weight_v") |
| 137 | + result = _replace_keys_in_dict(result, ".parametrizations.weight.original0", ".weight_g") |
153 | 138 | return result |
0 commit comments