Update app.py

app.py

@@ -11,6 +11,16 @@ from PIL import Image
 import numpy as np
 import gc
 
+
+from diffusers.schedulers import PNDMScheduler
+
+class CustomPNDMScheduler(PNDMScheduler):
+    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)
+
+        
 # Force JAX to use CPU
 jax.config.update('jax_platform_name', 'cpu')
 
@@ -50,6 +60,13 @@ pipeline, params = get_model(model_id, "flax")
 # Extract UNet from pipeline
 unet = pipeline.unet
 
+
+
+# After loading the pipeline
+custom_scheduler = CustomPNDMScheduler.from_config(pipeline.scheduler.config)
+pipeline.scheduler = custom_scheduler
+
+
 # Load and preprocess your dataset
 def preprocess_images(examples):
     def process_image(image):
@@ -97,6 +114,10 @@ def clear_jit_cache():
 # Training function
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
+        print("pixel_values dtype:", pixel_values.dtype)
+        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)
         
@@ -111,13 +132,14 @@ def train_step(state, batch, rng):
         # Generate random noise
         noise = jax.random.normal(rng, latents.shape, dtype=jnp.float32)
         
-        # Sample random timesteps (keep as integers)
+        # Sample random timesteps
         timesteps = jax.random.randint(
             rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
         )
         
-        # Explicitly cast timesteps to int32
-        timesteps = timesteps.astype(jnp.int32)
+        print("timesteps dtype:", timesteps.dtype)
+        print("latents dtype:", latents.dtype)
+        print("noise dtype:", noise.dtype)
         
         # Add noise to latents
         noisy_latents = pipeline.scheduler.add_noise(
@@ -157,9 +179,10 @@ def train_step(state, batch, rng):
 # Initialize training state
 learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
+float32_params = jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype != jnp.int32 else x, params)
 state = train_state.TrainState.create(
-    apply_fn=unet.__call__,  # Use __call__ directly
-    params=params,  # Pass all params
+    apply_fn=unet.__call__,
+    params=float32_params,
     tx=optimizer,
 )