Update app.py

app.py

@@ -94,4 +94,102 @@ print(f"Processed dataset size: {len(processed_dataset)}")
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
         print("pixel_values dtype:", pixel_values.dtype)
-        print("params dtypes:", jax.tree_map
+        print("params dtypes:", jax.tree_map(lambda x: x.dtype, params))
+        print("rng dtype:", rng.dtype)
+        
+        # Ensure pixel_values are float32
+        pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
+        
+        # Encode images to latent space
+        latents = pipeline.vae.apply(
+            {"params": params["vae"]},
+            pixel_values,
+            method=pipeline.vae.encode
+        ).latent_dist.sample(rng)
+        latents = latents * jnp.float32(0.18215)
+
+        # Generate random noise
+        noise = jax.random.normal(rng, latents.shape, dtype=jnp.float32)
+        
+        # Sample random timesteps
+        timesteps = jax.random.randint(
+            rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
+        )
+        
+        print("timesteps dtype:", timesteps.dtype)
+        print("latents dtype:", latents.dtype)
+        print("noise dtype:", noise.dtype)
+        
+        # Add noise to latents
+        noisy_latents = pipeline.scheduler.add_noise(
+            pipeline.scheduler.create_state(),
+            original_samples=latents,
+            noise=noise,
+            timesteps=timesteps
+        )
+        
+        # Generate random encoder hidden states (simulating text embeddings)
+        encoder_hidden_states = jax.random.normal(
+            rng, 
+            (latents.shape[0], pipeline.text_encoder.config.hidden_size),
+            dtype=jnp.float32
+        )
+        
+        # Predict noise
+        model_output = state.apply_fn(
+            {'params': params["unet"]},
+            noisy_latents,
+            timesteps,
+            encoder_hidden_states=encoder_hidden_states,
+            train=True,
+        )
+        
+        # Compute loss
+        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_map(lambda x: x.astype(jnp.float32) if x.dtype != jnp.int32 else x, params)
+state = train_state.TrainState.create(
+    apply_fn=unet.__call__,
+    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'], 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}")