-
Notifications
You must be signed in to change notification settings - Fork 129
Expand file tree
/
Copy pathgenerate_weights_a2.py
More file actions
354 lines (285 loc) · 13.6 KB
/
generate_weights_a2.py
File metadata and controls
354 lines (285 loc) · 13.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
#!/usr/bin/env python3
"""
Generate weights for wavenet_a2_max.nam file.
This script handles the full A2 architecture including:
- FiLM (Feature-wise Linear Modulation) modules
- head1x1 modules
- condition_dsp (nested WaveNet)
- Advanced gating modes (GATED, BLENDED, NONE)
- Complex activation configurations
"""
import json
import random
from pathlib import Path
from typing import Dict, Any, List
def count_conv1d_weights(in_channels: int, out_channels: int, kernel_size: int,
has_bias: bool, groups: int = 1) -> int:
"""Count weights for a Conv1D layer."""
weight_count = kernel_size * (out_channels * in_channels // groups)
if has_bias:
weight_count += out_channels
return weight_count
def count_conv1x1_weights(in_channels: int, out_channels: int,
has_bias: bool, groups: int = 1) -> int:
"""Count weights for a Conv1x1 layer (kernel_size=1)."""
weight_count = (out_channels * in_channels // groups)
if has_bias:
weight_count += out_channels
return weight_count
def count_film_weights(condition_dim: int, input_dim: int, has_shift: bool, groups: int = 1) -> int:
"""
Count weights for a FiLM (Feature-wise Linear Modulation) module.
FiLM uses a Conv1x1: condition_dim -> (2*input_dim if shift else input_dim), with bias.
Args:
condition_dim: Size of the conditioning input
input_dim: Size of the input to be modulated
has_shift: Whether to apply both scale and shift (true) or only scale (false)
groups: Number of groups for grouped convolution (default: 1)
"""
out_channels = (2 * input_dim) if has_shift else input_dim
return count_conv1x1_weights(condition_dim, out_channels, has_bias=True, groups=groups)
def parse_gating_mode(layer_config: Dict[str, Any], layer_index: int = 0) -> str:
"""
Parse gating mode from layer config (handles both old and new formats).
Args:
layer_config: Layer configuration dictionary
layer_index: Index of the layer (for array-valued gating_mode)
"""
if "gating_mode" in layer_config:
gating_mode_value = layer_config["gating_mode"]
if isinstance(gating_mode_value, list):
# Array of gating modes - use the one at layer_index
gating_mode_str = gating_mode_value[layer_index]
else:
# Single gating mode - use for all layers
gating_mode_str = gating_mode_value
if isinstance(gating_mode_str, str):
if gating_mode_str in ["GATED", "BLENDED", "NONE"]:
return gating_mode_str
# Handle lowercase versions
return gating_mode_str.upper()
return "NONE"
elif "gated" in layer_config:
# Backward compatibility
return "GATED" if layer_config["gated"] else "NONE"
else:
return "NONE"
def count_layer_weights(layer_config: Dict[str, Any], condition_size: int, layer_index: int = 0) -> int:
"""
Count weights for a single layer (one dilation).
A layer consists of:
1. Conv1D: (channels, bottleneck*(2 if gated/blended else 1), kernel_size, bias=True, groups_input)
2. Input mixin Conv1x1: (condition_size, bottleneck*(2 if gated/blended else 1), bias=False, groups_input_mixin)
3. Optional layer1x1 Conv1x1: (bottleneck, channels, bias=True, layer1x1_groups)
4. Optional head1x1 Conv1x1: (bottleneck, head1x1_out_channels, bias=True, head1x1_groups)
5. FiLM modules (optional, various configurations)
Args:
layer_config: Layer configuration dictionary
condition_size: Size of the conditioning input
layer_index: Index of the layer within the layer array (for array-valued configs)
"""
channels = layer_config["channels"]
bottleneck = layer_config.get("bottleneck", channels)
kernel_size = layer_config["kernel_size"]
groups_input = layer_config.get("groups_input", 1)
groups_input_mixin = layer_config.get("groups_input_mixin", 1)
# Parse layer1x1 parameters
layer1x1_active = True # default to active if not present
layer1x1_groups = 1
if "layer1x1" in layer_config:
layer1x1_config = layer_config["layer1x1"]
layer1x1_active = layer1x1_config.get("active", True) # default to active
layer1x1_groups = layer1x1_config.get("groups", 1)
gating_mode = parse_gating_mode(layer_config, layer_index)
# Output channels are doubled for GATED and BLENDED modes
conv_out_channels = 2 * bottleneck if gating_mode in ["GATED", "BLENDED"] else bottleneck
weight_count = 0
# 1. Conv1D weights
weight_count += count_conv1d_weights(
channels, conv_out_channels, kernel_size,
has_bias=True, groups=groups_input
)
# 2. Input mixin Conv1x1 weights
weight_count += count_conv1x1_weights(
condition_size, conv_out_channels,
has_bias=False, groups=groups_input_mixin
)
# 3. layer1x1 Conv1x1 weights (only if active)
if layer1x1_active:
weight_count += count_conv1x1_weights(
bottleneck, channels,
has_bias=True, groups=layer1x1_groups
)
# 4. Optional head1x1 weights
head1x1_config = layer_config.get("head1x1")
if head1x1_config and head1x1_config.get("active", False):
head1x1_out_channels = head1x1_config.get("out_channels", channels)
head1x1_groups = head1x1_config.get("groups", 1)
weight_count += count_conv1x1_weights(
bottleneck, head1x1_out_channels,
has_bias=True, groups=head1x1_groups
)
# 5. FiLM module weights
# Parse all possible FiLM configurations
film_configs = [
("conv_pre_film", channels),
("conv_post_film", conv_out_channels),
("input_mixin_pre_film", condition_size),
("input_mixin_post_film", conv_out_channels),
("activation_pre_film", conv_out_channels),
("activation_post_film", bottleneck),
("layer1x1_post_film", channels if layer1x1_active else 0), # Only count if layer1x1 is active
("head1x1_post_film", head1x1_config.get("out_channels", channels) if head1x1_config and head1x1_config.get("active") else 0)
]
for film_key, input_dim in film_configs:
if film_key in layer_config and layer_config[film_key]:
film_params = layer_config[film_key]
if isinstance(film_params, dict) and film_params.get("active", True):
has_shift = film_params.get("shift", True)
groups = film_params.get("groups", 1)
if input_dim > 0: # Only count if input_dim is valid
weight_count += count_film_weights(condition_size, input_dim, has_shift, groups)
return weight_count
def count_layer_array_weights(layer_config: Dict[str, Any]) -> int:
"""
Count the total number of weights for a layer array.
Each layer array consists of:
1. Rechannel Conv1x1: (input_size, channels, bias=False)
2. Layers (one per dilation)
3. Head rechannel Conv1x1: (head_output_size, head_size, bias=head_bias)
where head_output_size = head1x1.out_channels if head1x1 active, else bottleneck
"""
input_size = layer_config["input_size"]
condition_size = layer_config["condition_size"]
head_size = layer_config["head_size"]
channels = layer_config["channels"]
bottleneck = layer_config.get("bottleneck", channels)
dilations = layer_config["dilations"]
head_bias = layer_config.get("head_bias", False)
# Determine head output size: head1x1.out_channels if active, else bottleneck
head1x1_config = layer_config.get("head1x1")
if head1x1_config and head1x1_config.get("active", False):
head_output_size = head1x1_config.get("out_channels", channels)
else:
head_output_size = bottleneck
num_layers = len(dilations)
# Validate array-valued configs match number of layers
if "activation" in layer_config and isinstance(layer_config["activation"], list):
if len(layer_config["activation"]) != num_layers:
raise ValueError(f"activation array size ({len(layer_config['activation'])}) must match dilations size ({num_layers})")
if "gating_mode" in layer_config and isinstance(layer_config["gating_mode"], list):
if len(layer_config["gating_mode"]) != num_layers:
raise ValueError(f"gating_mode array size ({len(layer_config['gating_mode'])}) must match dilations size ({num_layers})")
if "secondary_activation" in layer_config and isinstance(layer_config["secondary_activation"], list):
if len(layer_config["secondary_activation"]) != num_layers:
raise ValueError(f"secondary_activation array size ({len(layer_config['secondary_activation'])}) must match dilations size ({num_layers})")
weight_count = 0
# 1. Rechannel weights
weight_count += count_conv1x1_weights(input_size, channels, has_bias=False, groups=1)
# 2. For each layer in the array
for layer_idx in range(num_layers):
weight_count += count_layer_weights(layer_config, condition_size, layer_idx)
# 3. Head rechannel weights (input is head_output_size, not bottleneck)
weight_count += count_conv1x1_weights(
head_output_size, head_size,
has_bias=head_bias, groups=1
)
return weight_count
def count_wavenet_weights(config: Dict[str, Any]) -> int:
"""
Count total weights for a WaveNet model (including optional condition_dsp).
"""
weight_count = 0
# Count weights for each layer array
for layer_config in config["layers"]:
weight_count += count_layer_array_weights(layer_config)
# Add head_scale (1 float)
weight_count += 1
return weight_count
def generate_weights(weight_count: int, seed: int = None,
weight_range: tuple = (-1.0, 1.0)) -> List[float]:
"""Generate random weights in the specified range."""
if seed is not None:
random.seed(seed)
return [random.uniform(*weight_range) for _ in range(weight_count)]
def add_weights_to_model(model_data: Dict[str, Any], seed: int = None) -> None:
print(f"Architecture: {model_data.get('architecture', 'Unknown')}")
# Process condition_dsp if present
if "config" in model_data and "condition_dsp" in model_data["config"]:
condition_dsp = model_data["config"]["condition_dsp"]
if condition_dsp and "config" in condition_dsp:
print("\nCounting weights for condition_dsp...")
condition_weights = count_wavenet_weights(condition_dsp["config"])
print(f" Condition DSP weights: {condition_weights}")
# Generate weights for condition_dsp
condition_dsp["weights"] = generate_weights(condition_weights, seed)
print(f" Generated {len(condition_dsp['weights'])} weights for condition_dsp")
# Count main model weights
print("\nCounting weights for main model...")
main_weights = count_wavenet_weights(model_data["config"])
print(f" Main model weights: {main_weights}")
# Generate weights for main model
model_data["weights"] = generate_weights(main_weights, seed)
print(f" Generated {len(model_data['weights'])} weights for main model")
# Print detailed breakdown
print("\nWeight breakdown:")
total_weights = 0
# Condition DSP breakdown
if "config" in model_data and "condition_dsp" in model_data["config"]:
condition_dsp = model_data["config"]["condition_dsp"]
if condition_dsp and "config" in condition_dsp:
print(" Condition DSP:")
for i, layer in enumerate(condition_dsp["config"]["layers"]):
layer_weights = count_layer_array_weights(layer)
print(f" Layer array {i+1}: {layer_weights} weights")
total_weights += layer_weights
total_weights += 1 # head_scale
# Main model breakdown
print(" Main model:")
for i, layer in enumerate(model_data["config"]["layers"]):
layer_weights = count_layer_array_weights(layer)
print(f" Layer array {i+1}: {layer_weights} weights")
total_weights += layer_weights
total_weights += 1 # head_scale
print(f"\nTotal weights generated: {total_weights}")
def process_model(input_path: Path, output_path: Path, seed: int = None) -> None:
"""
Load a .nam file with empty weights and generate random weights for it.
"""
# Load the input file
with open(input_path, 'r') as f:
model_data = json.load(f)
print(f"Processing: {input_path}")
add_weights_to_model(model_data, seed)
# Write output file
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, 'w') as f:
json.dump(model_data, f, indent=4)
print(f"\nOutput written to: {output_path}")
def main():
import argparse
parser = argparse.ArgumentParser(
description="Generate weights for A2 WaveNet models with empty weight arrays"
)
parser.add_argument(
"--input",
type=Path,
default=Path("example_models/wavenet_a2_max.nam"),
help="Input .nam file with empty weights (default: example_models/wavenet_a2_max.nam)"
)
parser.add_argument(
"--output",
type=Path,
default=Path("example_models/wavenet_a2_max_withweights.nam"),
help="Output .nam file (default: example_models/wavenet_a2_max_withweights.nam)"
)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Random seed for weight generation (default: 42)"
)
args = parser.parse_args()
process_model(args.input, args.output, args.seed)
if __name__ == "__main__":
main()