Update app.py

app.py

@@ -130,4 +130,86 @@ def train_step(state, batch, rng):
             pixel_values,
             method=pipeline.vae.encode
         ).latent_dist.sample(rng)
-        latents = latents *
+        latents = latents * jnp.float32(0.18215)
+        print(f"latents shape: {latents.shape}")
+
+        noise = jax.random.normal(rng, latents.shape, dtype=jnp.float32)
+        
+        timesteps = jax.random.randint(
+            rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
+        )
+        
+        noisy_latents = pipeline.scheduler.add_noise(
+            pipeline.scheduler.create_state(),
+            original_samples=latents,
+            noise=noise,
+            timesteps=timesteps
+        )
+        
+        encoder_hidden_states = jax.random.normal(
+            rng, 
+            (latents.shape[0], pipeline.text_encoder.config.hidden_size),
+            dtype=jnp.float32
+        )
+        
+        print(f"noisy_latents shape: {noisy_latents.shape}")
+        print(f"timesteps shape: {timesteps.shape}")
+        print(f"encoder_hidden_states shape: {encoder_hidden_states.shape}")
+        
+        # Use the correct method to call the UNet
+        model_output = unet.apply(
+            {'params': params["unet"]},
+            noisy_latents,
+            jnp.array(timesteps, dtype=jnp.int32),
+            encoder_hidden_states,
+            train=True,
+        ).sample
+        
+        return jnp.mean((model_output - noise) ** 2)
+
+    grad_fn = jax.grad(compute_loss, argnums=0, allow_int=True)
+    rng, step_rng = jax.random.split(rng)
+    
+    grads = grad_fn(state.params, batch["pixel_values"], step_rng)
+    loss = compute_loss(state.params, batch["pixel_values"], step_rng)
+    state = state.apply_gradients(grads=grads)
+    return state, loss
+
+# Initialize training state
+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)
+state = train_state.TrainState.create(
+    apply_fn=unet.apply,
+    params=float32_params,
+    tx=optimizer,
+)
+
+# Training loop
+num_epochs = 3
+batch_size = 1
+rng = jax.random.PRNGKey(0)
+
+for epoch in range(num_epochs):
+    epoch_loss = 0
+    num_batches = 0
+    for batch in tqdm(processed_dataset.batch(batch_size)):
+        batch['pixel_values'] = jnp.array(batch['pixel_values'][0], dtype=jnp.float32)
+        rng, step_rng = jax.random.split(rng)
+        state, loss = train_step(state, batch, step_rng)
+        epoch_loss += loss
+        num_batches += 1
+        
+        if num_batches % 10 == 0:
+            jax.clear_caches()
+    
+    avg_loss = epoch_loss / num_batches
+    print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
+    jax.clear_caches()
+
+# 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"])
+
+print(f"Model saved to {output_dir}")