Update app.py

app.py

@@ -129,10 +129,15 @@ def train_step(state, batch, rng):
             dtype=jnp.float32
         )
         
+        # Ensure noisy_latents has the correct number of channels
+        if noisy_latents.shape[-1] != pipeline.unet.config.in_channels:
+            pad_width = [(0, 0)] * (noisy_latents.ndim - 1) + [(0, pipeline.unet.config.in_channels - noisy_latents.shape[-1])]
+            noisy_latents = jnp.pad(noisy_latents, pad_width, mode='constant')
+        
         # Use the correct method to call the UNet
         model_output = unet.apply(
             {'params': params["unet"]},
-            jnp.array(noisy_latents),
+            noisy_latents,
             jnp.array(timesteps, dtype=jnp.int32),
             encoder_hidden_states,
             train=True,