This commit is contained in:
tuna2134
2026-07-26 22:26:11 +09:00
parent db37175b87
commit aa70540aef

View File

@@ -134,21 +134,47 @@ def unpack_batch(
hps: HyperParameters,
device: torch.device,
) -> tuple[dict[str, Any], torch.Tensor]:
(
x,
x_lengths,
spec,
spec_lengths,
waveform,
_,
speakers,
tone,
language,
bert,
ja_bert,
en_bert,
style_vec,
) = (item.to(device, non_blocking=True) for item in batch)
batch = tuple(item.to(device, non_blocking=True) for item in batch)
if hps.data.use_jp_extra:
(
x,
x_lengths,
spec,
spec_lengths,
waveform,
_,
speakers,
tone,
language,
bert,
style_vec,
) = batch
encoder_args = (tone, language, bert, style_vec)
else:
(
x,
x_lengths,
spec,
spec_lengths,
waveform,
_,
speakers,
tone,
language,
bert,
ja_bert,
en_bert,
style_vec,
) = batch
encoder_args = (
tone,
language,
bert,
ja_bert,
en_bert,
style_vec,
speakers,
)
mel = (
spec
if hps.model.use_mel_posterior_encoder
@@ -161,11 +187,6 @@ def unpack_batch(
hps.data.mel_fmax,
)
)
encoder_args = (
(tone, language, bert, style_vec)
if hps.data.use_jp_extra
else (tone, language, bert, ja_bert, en_bert, style_vec, speakers)
)
return {
"x": x,
"x_lengths": x_lengths,
@@ -260,7 +281,7 @@ def run() -> None:
shuffle=True,
num_workers=args.num_workers,
pin_memory=device.type == "cuda",
collate_fn=TextAudioSpeakerCollate(),
collate_fn=TextAudioSpeakerCollate(use_jp_extra=hps.data.use_jp_extra),
drop_last=True,
)
output_dir = Path(args.output_dir or Path(args.config).parent / "models")