Update app.py

app.py

@@ -145,6 +145,15 @@ def train_step(state, batch, rng):
     state = state.apply_gradients(grads=grads)
     return state, loss
 
+# Initialize training state
+learning_rate = 1e-5
+optimizer = optax.adam(learning_rate)
+state = train_state.TrainState.create(
+    apply_fn=unet.__call__,
+    params={"unet": params["unet"], "vae": params["vae"]},
+    tx=optimizer,
+)
+
 # Training loop
 num_epochs = 3
 batch_size = 1  # Reduced batch size due to memory constraints
@@ -167,6 +176,8 @@ for epoch in range(num_epochs):
     print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
     clear_jit_cache()
 
+
+    
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"
 os.makedirs(output_dir, exist_ok=True)