Update app.py

app.py

@@ -10,6 +10,9 @@ import os
 import pickle
 from PIL import Image
 import numpy as np
+ 
+ 
+
 
 # Set up cache directories
 cache_dir = "/tmp/huggingface_cache"
@@ -120,9 +123,13 @@ def train_step(state, batch, rng):
         batch_size = pixel_values.shape[0]
         
         # Encode images to latent space
-        latents = pipeline.vae.encode(pixel_values).latent_dist.sample(rng)
-        latents = latents * pipeline.vae.config.scaling_factor
-        
+        latents = pipeline.vae.apply(
+            {"params": params["vae"]},
+            pixel_values,
+            method=pipeline.vae.encode
+        ).latent_dist.sample(rng)
+        latents = latents * 0.18215  # scaling factor
+
         # Generate random noise
         noise_rng, timestep_rng, latents_rng = jax.random.split(rng, 3)
         noise = jax.random.normal(noise_rng, latents.shape)
@@ -148,7 +155,7 @@ def train_step(state, batch, rng):
         
         # Predict noise
         model_output = state.apply_fn.apply(
-            {'params': params},
+            {'params': params["unet"]},
             jnp.array(noisy_latents),
             jnp.array(timesteps),
             encoder_hidden_states=encoder_hidden_states,
@@ -168,7 +175,7 @@ learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
 state = train_state.TrainState.create(
     apply_fn=unet,
-    params=params["unet"],  # Use only UNet params
+    params={"unet": params["unet"], "vae": params["vae"]},  # Include both UNet and VAE params
     tx=optimizer,
 )
 
@@ -191,6 +198,7 @@ for epoch in range(num_epochs):
     print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
 
 
+
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"
 os.makedirs(output_dir, exist_ok=True)