Update app.py

app.py

@@ -111,12 +111,17 @@ print(f"Processed dataset size: {len(processed_dataset)}")
 
 # Print sample input shape
 sample_batch = next(iter(processed_dataset.batch(1)))
-print(f"Sample input shape: {sample_batch['pixel_values'].shape}")
+print(f"Sample batch keys: {sample_batch.keys()}")
+print(f"Sample pixel_values type: {type(sample_batch['pixel_values'])}")
+print(f"Sample pixel_values length: {len(sample_batch['pixel_values'])}")
+if len(sample_batch['pixel_values']) > 0:
+    print(f"Sample pixel_values[0] shape: {np.array(sample_batch['pixel_values'][0]).shape}")
 
 # Training function
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
         pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
+        print(f"pixel_values shape in compute_loss: {pixel_values.shape}")
         
         latents = pipeline.vae.apply(
             {"params": params["vae"]},
@@ -124,6 +129,7 @@ def train_step(state, batch, rng):
             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)
         
@@ -186,7 +192,7 @@ 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)
+        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