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, FlaxUNet2DConditionModel
+from diffusers import FlaxStableDiffusionPipeline
 from diffusers.schedulers import FlaxPNDMScheduler
 from datasets import load_dataset
 from tqdm.auto import tqdm
@@ -53,22 +53,8 @@ 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
-
-# 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
+# Extract UNet from pipeline
+unet = pipeline.unet
 
 # Load and preprocess your dataset
 def preprocess_images(examples):