Update app.py

app.py

@@ -29,12 +29,11 @@ def get_model(model_id, revision):
             return pickle.load(f)
     else:
         print("Downloading model...")
-        pipeline = FlaxStableDiffusionPipeline.from_pretrained(
+        pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
             model_id, 
             revision=revision,
             dtype=jnp.float32,
         )
-        params = pipeline.params
         with open(model_cache_file, 'wb') as f:
             pickle.dump((pipeline, params), f)
         return pipeline, params
@@ -102,5 +101,73 @@ except Exception as e:
     else:
         raise ValueError(f"Local path {local_path} does not exist.")
 
-# Rest of your code (training loop, etc.) remains the same
-...
+# Training function
+def train_step(state, batch, rng):
+    def compute_loss(params):
+        # Convert batch to JAX array
+        pixel_values = jnp.array(batch["pixel_values"])
+        batch_size = pixel_values.shape[0]
+        
+        # Generate random noise
+        noise_rng, timestep_rng = jax.random.split(rng)
+        noise = jax.random.normal(noise_rng, pixel_values.shape)
+        
+        # Sample random timesteps
+        timesteps = jax.random.randint(
+            timestep_rng, (batch_size,), 0, pipeline.scheduler.config.num_train_timesteps
+        )
+        
+        # Add noise to images using the scheduler
+        noisy_images = pipeline.scheduler.add_noise(
+            original_samples=pixel_values,
+            noise=noise,
+            timesteps=timesteps
+        )
+        
+        # Predict noise
+        model_output = state.apply_fn.apply(
+            {'params': params},
+            jnp.array(noisy_images),
+            jnp.array(timesteps),
+            train=True,
+        )
+        
+        # Compute loss
+        loss = jnp.mean((model_output - noise) ** 2)
+        return loss
+
+    loss, grads = jax.value_and_grad(compute_loss)(state.params)
+    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,
+    params=params["unet"],  # Use only UNet params
+    tx=optimizer,
+)
+
+# Training loop
+num_epochs = 10
+batch_size = 4
+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)):
+        rng, step_rng = jax.random.split(rng)
+        state, loss = train_step(state, batch, step_rng)
+        epoch_loss += loss
+        num_batches += 1
+    avg_loss = epoch_loss / num_batches
+    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)
+unet.save_pretrained(output_dir, params=state.params)
+
+print(f"Model saved to {output_dir}")