Update app.py

app.py

@@ -27,6 +27,17 @@ os.makedirs(model_cache_dir, exist_ok=True)
 print(f"Cache directory: {cache_dir}")
 print(f"Model cache directory: {model_cache_dir}")
 
+
+
+def filter_dict(dict_to_filter, target_callable):
+    """Filter a dictionary to only include keys that are valid parameters for the target callable."""
+    valid_params = signature(target_callable).parameters.keys()
+    return {k: v for k, v in dict_to_filter.items() if k in valid_params}
+
+
+
+
+
 # Function to load or download the model
 def get_model(model_id, revision):
     model_cache_file = os.path.join(model_cache_dir, f"{model_id.replace('/', '_')}_{revision}.pkl")
@@ -227,6 +238,6 @@ for epoch in range(num_epochs):
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"
 os.makedirs(output_dir, exist_ok=True)
-adjusted_unet.save_pretrained(output_dir, params=state.params["params"])
+adjusted_unet.save_pretrained(output_dir, params=state.params)
 
 print(f"Model saved to {output_dir}")