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,24 +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
 
-# Modify the UNet's input layer
-def modify_unet_input_layer(params):
-    conv_in_weight = params['unet']['conv_in']['kernel']
-    conv_in_bias = params['unet']['conv_in']['bias']
-    
-    # Adjust the weight tensor
-    new_weight = jnp.zeros((3, 3, 4, 320), dtype=jnp.float32)
-    new_weight = new_weight.at[:, :, :4, :].set(conv_in_weight[:, :, :4, :])
-    
-    # Update the parameters
-    params['unet']['conv_in']['kernel'] = new_weight
-    
-    return params
+# 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"]
 
-params = modify_unet_input_layer(params)
+# Replace the UNet params in the pipeline
+params["unet"] = new_unet_params
 
 # Load and preprocess your dataset
 def preprocess_images(examples):