Update app.py

app.py

@@ -15,7 +15,6 @@ import numpy as np
 # Custom Scheduler
 class CustomFlaxPNDMScheduler(FlaxPNDMScheduler):
     def add_noise(self, state, original_samples, noise, timesteps):
-        # Explicitly cast timesteps to int32
         timesteps = timesteps.astype(jnp.int32)
         return super().add_noise(state, original_samples, noise, timesteps)
 
@@ -97,10 +96,8 @@ def train_step(state, batch, rng):
         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,
@@ -108,10 +105,8 @@ def train_step(state, batch, rng):
         ).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
         )
@@ -119,8 +114,8 @@ def train_step(state, batch, rng):
         print("timesteps dtype:", timesteps.dtype)
         print("latents dtype:", latents.dtype)
         print("noise dtype:", noise.dtype)
+        print("latents shape:", latents.shape)
         
-        # Add noise to latents
         noisy_latents = pipeline.scheduler.add_noise(
             pipeline.scheduler.create_state(),
             original_samples=latents,
@@ -128,14 +123,12 @@ def train_step(state, batch, rng):
             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 = unet.apply(
             {'params': params["unet"]},
             noisy_latents,
@@ -144,7 +137,6 @@ def train_step(state, batch, rng):
             train=True,
         )
         
-        # Compute loss
         return jnp.mean((model_output - noise) ** 2)
 
     grad_fn = jax.grad(compute_loss, argnums=0, allow_int=True)