Update app.py

app.py

@@ -122,6 +122,7 @@ if len(sample_batch['pixel_values']) > 0:
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
         pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
+        pixel_values = jnp.expand_dims(pixel_values, axis=0)  # Add batch dimension
         print(f"pixel_values shape in compute_loss: {pixel_values.shape}")
         
         latents = pipeline.vae.apply(
@@ -129,86 +130,4 @@ def train_step(state, batch, rng):
             pixel_values,
             method=pipeline.vae.encode
         ).latent_dist.sample(rng)
-        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}")
+        latents = latents *