Update app.py

app.py

@@ -57,8 +57,26 @@ pipeline.scheduler = custom_scheduler
 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(unet_config)
+unet = FlaxUNet2DConditionModel(modified_unet_config)
 
 # Initialize the new UNet with random weights
 rng = jax.random.PRNGKey(0)