Update app.py

app.py

@@ -56,6 +56,22 @@ pipeline.scheduler = custom_scheduler
 # Extract UNet from pipeline
 unet = pipeline.unet
 
+# 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
+
+params = modify_unet_input_layer(params)
+
 # Load and preprocess your dataset
 def preprocess_images(examples):
     def process_image(image):