Update app.py

app.py

@@ -56,12 +56,18 @@ pipeline.scheduler = custom_scheduler
 # Extract UNet from pipeline
 unet = pipeline.unet
 
+# Print UNet configuration
+print("UNet configuration:")
+print(unet.config)
+
 # Adjust the input layer of the UNet
 def adjust_unet_input_layer(params):
     conv_in_weight = params['unet']['conv_in']['kernel']
+    print(f"Original conv_in weight shape: {conv_in_weight.shape}")
     new_conv_in_weight = jnp.zeros((3, 3, 4, 320), dtype=jnp.float32)
     new_conv_in_weight = new_conv_in_weight.at[:, :, :3, :].set(conv_in_weight[:, :, :3, :])
     params['unet']['conv_in']['kernel'] = new_conv_in_weight
+    print(f"New conv_in weight shape: {params['unet']['conv_in']['kernel'].shape}")
     return params
 
 params = adjust_unet_input_layer(params)
@@ -103,6 +109,10 @@ else:
 
 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}")
+
 # Training function
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
@@ -134,6 +144,10 @@ def train_step(state, batch, rng):
             dtype=jnp.float32
         )
         
+        print(f"noisy_latents shape: {noisy_latents.shape}")
+        print(f"timesteps shape: {timesteps.shape}")
+        print(f"encoder_hidden_states shape: {encoder_hidden_states.shape}")
+        
         # Use the correct method to call the UNet
         model_output = unet.apply(
             {'params': params["unet"]},