Repository: NVIDIA/NeMo

Issue:
Spectrogram Enhancer doesn't generalize to spectrogram lengths unseen during training
**Describe the bug**

If one trains TTS spectrogram enhancer (#5565) on short spectrograms, at inference time it doesn't generalize to longer ones. A patch in the beginning gets enhanced but further frames do not.

Example (before, after):
![before](https://user-images.githubusercontent.com/8864149/218520151-230501b1-4e9d-4307-b51e-526440b915c8.png)
![after](https://user-images.githubusercontent.com/8864149/218520199-22187a4b-9cf8-4ac9-b3ce-07c320d3d3b6.png)

**Steps/Code to reproduce bug**

1. Train a spectrogram enhancer
2. Apply it to a spectrogram that's longer than anything from the training set
3. Only a patch in the beginning gets enhanced

**Expected behavior**

The whole spectrogram should have got additional details, not just the first patch


Discussion:
A temporary fix: given a trained model, clone first patch of the initial tensor length-wise:
```
max_init_length = enhancer.generator.initial_block.shape[-1]
m = 6
for i in range(1, max_init_length // m):
    enhancer.generator.initial_block.data[:,:,:,i*m:(i+1)*m] = enhancer.generator.initial_block.data[:,:,:,0:m]
```

Visually fix works, not sure what about downstream tasks
![after_temporary_fix](https://user-images.githubusercontent.com/8864149/218524157-9283395d-f196-41aa-81fe-f8868c2acb67.png)

As-is most likely breaks #5659 (setup with enhancer)

Resolving patch:
diff --git a/nemo/collections/tts/modules/spectrogram_enhancer.py b/nemo/collections/tts/modules/spectrogram_enhancer.py
--- a/nemo/collections/tts/modules/spectrogram_enhancer.py
+++ b/nemo/collections/tts/modules/spectrogram_enhancer.py
@@ -230,14 +230,7 @@ def forward(self, x):
 
 class Generator(torch.nn.Module):
     def __init__(
-        self,
-        n_bands,
-        latent_dim,
-        style_depth,
-        network_capacity=16,
-        channels=1,
-        fmap_max=512,
-        max_spectrogram_length=2000,
+        self, n_bands, latent_dim, style_depth, network_capacity=16, channels=1, fmap_max=512,
     ):
         super().__init__()
         self.image_size = n_bands
@@ -277,7 +270,7 @@ def __init__(
             torch.nn.init.zeros_(block.to_noise2.weight)
             torch.nn.init.zeros_(block.to_noise2.bias)
 
-        initial_block_size = n_bands // self.upsample_factor, math.ceil(max_spectrogram_length / self.upsample_factor)
+        initial_block_size = n_bands // self.upsample_factor, 1
         self.initial_block = torch.nn.Parameter(
             torch.randn((1, init_channels, *initial_block_size)), requires_grad=False
         )
@@ -304,8 +297,7 @@ def upsample_factor(self):
     def forward(self, condition: torch.Tensor, lengths: torch.Tensor, ws: List[torch.Tensor], noise: torch.Tensor):
         batch_size, _, _, max_length = condition.shape
 
-        x = self.initial_block.expand(batch_size, -1, -1, -1)
-        x = x[:, :, :, : max_length // self.upsample_factor]
+        x = self.initial_block.expand(batch_size, -1, -1, max_length // self.upsample_factor)
 
         rgb = None
         x = self.initial_conv(x)


Test patch:
diff --git a/tests/collections/tts/test_spectrogram_enhancer.py b/tests/collections/tts/test_spectrogram_enhancer.py
--- a/tests/collections/tts/test_spectrogram_enhancer.py
+++ b/tests/collections/tts/test_spectrogram_enhancer.py
@@ -37,7 +37,6 @@ def enhancer_config():
             "network_capacity": network_capacity,
             "mixed_prob": 0.9,
             "fmap_max": fmap_max,
-            "max_spectrogram_length": 2000,
             "generator": {
                 "_target_": "nemo.collections.tts.modules.spectrogram_enhancer.Generator",
                 "n_bands": n_bands,
@@ -45,7 +44,6 @@ def enhancer_config():
                 "network_capacity": network_capacity,
                 "style_depth": style_depth,
                 "fmap_max": fmap_max,
-                "max_spectrogram_length": 2000,
             },
             "discriminator": {
                 "_target_": "nemo.collections.tts.modules.spectrogram_enhancer.Discriminator",