Update app.py

app.py

@@ -136,7 +136,7 @@ def train_step(state, batch, rng):
         )
         
         # Predict noise
-        model_output = state.apply_fn(
+        model_output = unet.apply(
             {'params': params["unet"]},
             noisy_latents,
             timesteps,
@@ -160,7 +160,7 @@ learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
 float32_params = jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype != jnp.int32 else x, params)
 state = train_state.TrainState.create(
-    apply_fn=unet.__call__,
+    apply_fn=unet.apply,
     params=float32_params,
     tx=optimizer,
 )