Fix dtype mismatch in vocoder convolutions

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_ltx2.py

@@ -94,6 +94,7 @@ def __call__(self, x: Array) -> Array:
       x = jnp.pad(x, ((0, 0), (self.pad_left, self.pad_right), (0, 0)), mode='edge')
 
     filter_expanded = jnp.repeat(self.filter, num_channels, axis=2)
+    filter_expanded = filter_expanded.astype(x.dtype)
 
     x_filtered = jax.lax.conv_general_dilated(
         x,
@@ -149,6 +150,7 @@ def __call__(self, x: Array) -> Array:
     x = jnp.pad(x, ((0, 0), (self.pad, self.pad), (0, 0)), mode='edge')
 
     filter_expanded = jnp.repeat(self.filter, num_channels, axis=2)
+    filter_expanded = filter_expanded.astype(x.dtype)
 
     x_upsampled = jax.lax.conv_general_dilated(
         x,
@@ -486,6 +488,7 @@ def __call__(self, waveform: Array) -> tuple[Array, Array]:
 
     left_pad = max(0, self.window_length - self.hop_length)
     waveform = jnp.pad(waveform, ((0, 0), (left_pad, 0), (0, 0)))
+    waveform = waveform.astype(self.forward_basis.value.dtype)
 
     spec = jax.lax.conv_general_dilated(
         waveform,