Add shape matching workaround in ResBlock

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_ltx2.py

@@ -314,6 +314,12 @@ def __call__(self, x: Array) -> Array:
       xt = conv1(xt)
       xt = act2(xt)
       xt = conv2(xt)
+      
+      if xt.shape[1] < x.shape[1]:
+        xt = jnp.pad(xt, ((0, 0), (0, x.shape[1] - xt.shape[1]), (0, 0)), mode='edge')
+      elif xt.shape[1] > x.shape[1]:
+        xt = xt[:, :x.shape[1], :]
+        
       x = x + xt
     return x