Update app.py

app.py

@@ -3,7 +3,7 @@ import jax.numpy as jnp
 from flax.jax_utils import replicate
 from flax.training import train_state
 import optax
-from diffusers import FlaxStableDiffusionPipeline
+from diffusers import FlaxStableDiffusionPipeline, FlaxUNet2DConditionModel
 from diffusers.schedulers import FlaxPNDMScheduler
 from datasets import load_dataset
 from tqdm.auto import tqdm
@@ -53,8 +53,22 @@ pipeline, params = get_model(model_id, "flax")
 custom_scheduler = CustomFlaxPNDMScheduler.from_config(pipeline.scheduler.config)
 pipeline.scheduler = custom_scheduler
 
-# Extract UNet from pipeline
-unet = pipeline.unet
+# Modify UNet configuration
+unet_config = pipeline.unet.config
+unet_config.in_channels = 4  # Set to match the latent space dimensions
+
+# Create a new UNet with the modified configuration
+unet = FlaxUNet2DConditionModel(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"]
+
+# Replace the UNet params in the pipeline
+params["unet"] = new_unet_params
 
 # Load and preprocess your dataset
 def preprocess_images(examples):
@@ -129,11 +143,6 @@ def train_step(state, batch, rng):
             dtype=jnp.float32
         )
         
-        # Ensure noisy_latents has the correct number of channels
-        if noisy_latents.shape[-1] != pipeline.unet.config.in_channels:
-            pad_width = [(0, 0)] * (noisy_latents.ndim - 1) + [(0, pipeline.unet.config.in_channels - noisy_latents.shape[-1])]
-            noisy_latents = jnp.pad(noisy_latents, pad_width, mode='constant')
-        
         # Use the correct method to call the UNet
         model_output = unet.apply(
             {'params': params["unet"]},