Update app.py

app.py

@@ -53,40 +53,18 @@ pipeline, params = get_model(model_id, "flax")
 custom_scheduler = CustomFlaxPNDMScheduler.from_config(pipeline.scheduler.config)
 pipeline.scheduler = custom_scheduler
 
-# Modify UNet configuration
-unet_config = pipeline.unet.config
-unet_config.in_channels = 4  # Set to match the latent space dimensions
-
-# Modify the UNet architecture
-def modify_unet_config(config):
-    config.down_block_types = [
-        "CrossAttnDownBlock2D",
-        "CrossAttnDownBlock2D",
-        "CrossAttnDownBlock2D",
-        "DownBlock2D"
-    ]
-    config.up_block_types = [
-        "UpBlock2D",
-        "CrossAttnUpBlock2D",
-        "CrossAttnUpBlock2D",
-        "CrossAttnUpBlock2D"
-    ]
-    return config
-
-modified_unet_config = modify_unet_config(unet_config)
-
-# Create a new UNet with the modified configuration
-unet = FlaxUNet2DConditionModel(modified_unet_config)
-
-# Initialize the new UNet with random weights
-rng = jax.random.PRNGKey(0)
-sample_input = jnp.ones((1, 64, 64, 4))
-sample_t = jnp.ones((1,))
-sample_encoder_hidden_states = jnp.ones((1, 77, 768))
-new_unet_params = unet.init(rng, sample_input, sample_t, sample_encoder_hidden_states)["params"]
+# Extract UNet from pipeline
+unet = pipeline.unet
+
+# Adjust the input layer of the UNet
+def adjust_unet_input_layer(params):
+    conv_in_weight = params['unet']['conv_in']['kernel']
+    new_conv_in_weight = jnp.zeros((3, 3, 4, 320), dtype=jnp.float32)
+    new_conv_in_weight = new_conv_in_weight.at[:, :, :4, :].set(conv_in_weight[:, :, :4, :])
+    params['unet']['conv_in']['kernel'] = new_conv_in_weight
+    return params
 
-# Replace the UNet params in the pipeline
-params["unet"] = new_unet_params
+params = adjust_unet_input_layer(params)
 
 # Load and preprocess your dataset
 def preprocess_images(examples):
@@ -124,10 +102,6 @@ print(f"Processed dataset size: {len(processed_dataset)}")
 # 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)
-        
         pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
         
         latents = pipeline.vae.apply(
@@ -143,11 +117,6 @@ def train_step(state, batch, rng):
             rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
         )
         
-        print("timesteps dtype:", timesteps.dtype)
-        print("latents dtype:", latents.dtype)
-        print("noise dtype:", noise.dtype)
-        print("latents shape:", latents.shape)
-        
         noisy_latents = pipeline.scheduler.add_noise(
             pipeline.scheduler.create_state(),
             original_samples=latents,