Update app.py

app.py

@@ -161,7 +161,7 @@ def train_step(state, batch, rng):
         print(f"encoder_hidden_states shape: {encoder_hidden_states.shape}")
         
         # Use the correct method to call the UNet
-        model_output = unet.apply(
+        model_output = state.apply_fn(
             {'params': params["unet"]},
             noisy_latents,
             jnp.array(timesteps, dtype=jnp.int32),
@@ -183,9 +183,15 @@ def train_step(state, batch, rng):
 learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
 float32_params = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype != jnp.int32 else x, params)
+
+# Create a new UNet with the adjusted parameters
+adjusted_unet = FlaxUNet2DConditionModel(**unet.config)
+adjusted_params = adjusted_unet.init(jax.random.PRNGKey(0), jnp.ones((1, 4, 64, 64)), jnp.ones((1,)), jnp.ones((1, 77, 768)))
+adjusted_params['params'] = float32_params['unet']
+
 state = train_state.TrainState.create(
-    apply_fn=unet.apply,
-    params=float32_params,
+    apply_fn=adjusted_unet.apply,
+    params=adjusted_params,
     tx=optimizer,
 )
 
@@ -214,6 +220,6 @@ for epoch in range(num_epochs):
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"
 os.makedirs(output_dir, exist_ok=True)
-unet.save_pretrained(output_dir, params=state.params["unet"])
+adjusted_unet.save_pretrained(output_dir, params=state.params["params"])
 
 print(f"Model saved to {output_dir}")