Update app.py

app.py

@@ -64,8 +64,11 @@ print(unet.config)
 def adjust_unet_input_layer(params):
     conv_in_weight = params['unet']['conv_in']['kernel']
     print(f"Original conv_in weight shape: {conv_in_weight.shape}")
-    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, :])
+    if conv_in_weight.shape[2] == 64:
+        new_conv_in_weight = conv_in_weight[:, :, :4, :]
+    else:
+        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}")
     return params