Add shape matching workaround for BWE residual and skip

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_ltx2.py

@@ -648,6 +648,12 @@ def __call__(self, mel_spec: Array) -> Array:
 
     skip = self.resampler(x)
     residual = jnp.transpose(residual, (0, 2, 1))
+    
+    if residual.shape[1] < skip.shape[1]:
+      residual = jnp.pad(residual, ((0, 0), (0, skip.shape[1] - residual.shape[1]), (0, 0)), mode='edge')
+    elif residual.shape[1] > skip.shape[1]:
+      residual = residual[:, :skip.shape[1], :]
+      
     waveform = jnp.clip(residual + skip, -1, 1)
 
     output_samples = num_samples * self.output_sampling_rate // self.input_sampling_rate