Update app.py

app.py

@@ -74,15 +74,13 @@ print(unet.config)
 
 # Adjust the input layer of the UNet
 def adjust_unet_input_layer(params):
-    conv_in_weight = params['unet']['conv_in']['kernel']
+    conv_in_weight = params['conv_in']['kernel']
     print(f"Original conv_in weight shape: {conv_in_weight.shape}")
-    if conv_in_weight.shape[2] == 64:
-        new_conv_in_weight = conv_in_weight[:, :, :4, :]
-    else:
+    if conv_in_weight.shape[2] != 4:
         new_conv_in_weight = jnp.zeros((3, 3, 4, 320), dtype=jnp.float32)
         new_conv_in_weight = new_conv_in_weight.at[:, :, :3, :].set(conv_in_weight[:, :, :3, :])
-    params['unet']['conv_in']['kernel'] = new_conv_in_weight
-    print(f"New conv_in weight shape: {params['unet']['conv_in']['kernel'].shape}")
+        params['conv_in']['kernel'] = new_conv_in_weight
+    print(f"New conv_in weight shape: {params['conv_in']['kernel'].shape}")
     return params
 
 params = adjust_unet_input_layer(params)
@@ -205,7 +203,7 @@ print("Filtered UNet config keys:", filtered_unet_config.keys())
 
 adjusted_unet = FlaxUNet2DConditionModel(**filtered_unet_config)
 adjusted_params = adjusted_unet.init(jax.random.PRNGKey(0), jnp.ones((1, 4, 64, 64)), jnp.ones((1,)), jnp.ones((1, 77, 768)))
-adjusted_params = float32_params['unet']  # Use only UNet params
+adjusted_params = adjust_unet_input_layer(adjusted_params)  # Adjust the input layer
 
 state = train_state.TrainState.create(
     apply_fn=adjusted_unet.apply,