Update app.py

app.py

@@ -74,13 +74,28 @@ print(unet.config)
 
 # Adjust the input layer of the UNet
 def adjust_unet_input_layer(params):
-    conv_in_weight = params['conv_in']['kernel']
+    if 'unet' in params:
+        unet_params = params['unet']
+    else:
+        unet_params = params
+
+    if 'conv_in' not in unet_params:
+        print("Warning: 'conv_in' not found in UNet params. Skipping input layer adjustment.")
+        return params
+
+    conv_in_weight = unet_params['conv_in']['kernel']
     print(f"Original conv_in weight shape: {conv_in_weight.shape}")
     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['conv_in']['kernel'] = new_conv_in_weight
-    print(f"New conv_in weight shape: {params['conv_in']['kernel'].shape}")
+        unet_params['conv_in']['kernel'] = new_conv_in_weight
+    print(f"New conv_in weight shape: {unet_params['conv_in']['kernel'].shape}")
+
+    if 'unet' in params:
+        params['unet'] = unet_params
+    else:
+        params = unet_params
+
     return params
 
 params = adjust_unet_input_layer(params)