mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix surrogate loss return in waveform_gen2

Browse Source
This commit is contained in:
James Betker 2022-04-28 10:10:11 -06:00
parent e208d9fb80
commit 8aa6651fc7
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
codes/models/audio/music/unet_diffusion_waveform_gen2.py
View File

@ -376,7 +376,7 @@ class DiffusionTts(nn.Module):
aligned_conditioning = F.pad(aligned_conditioning, (0,int(pc*aligned_conditioning.shape[-1])))
return x, aligned_conditioning
def forward(self, x, timesteps, conditioning, conditioning_free=False):
def forward(self, x, timesteps, conditioning, return_surrogate=True, conditioning_free=False):
"""
Apply the model to an input batch.
@ -433,7 +433,10 @@ class DiffusionTts(nn.Module):
extraneous_addition = extraneous_addition + p.mean()
out = out + extraneous_addition * 0
return out[:, :, :orig_x_shape], surrogate[:, :, :orig_x_shape]
if return_surrogate:
return out[:, :, :orig_x_shape], surrogate[:, :, :orig_x_shape]
else:
return out[:, :, :orig_x_shape]
@register_model