Remove cached and exported models after conversion.

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
 import json
+import shutil
 import subprocess
 import urllib.parse
 from pathlib import Path
 
-from huggingface_hub import hf_hub_download, HfApi
+from huggingface_hub import hf_hub_download, HfApi, scan_cache_dir
 from coremltools import ComputeUnit
 from coremltools.models.utils import _is_macos, _macos_version
 
@@ -201,7 +202,6 @@ def on_model_change(model):
     )
 
 
-
 def convert_model(preprocessor, model, model_coreml_config,
                   compute_units, precision, tolerance, output,
                   use_past=False, seq2seq=None,
@@ -252,6 +252,26 @@ def push_to_hub(destination, directory, task, precision, token=None):
     return get_pr_url(HfApi(token=token), destination, commit_message)
 
 
+def cleanup(model_id, exported):
+    if exported:
+        shutil.rmtree(exported)
+
+    # We remove the model from the huggingface cache, so it will have to be downloaded again
+    # if the user wants to convert it for a different task or precision.
+    # Alternatively, we could remove models older than 1 day or so.
+    cache_info = scan_cache_dir()
+    try:
+        repo = next(repo for repo in cache_info.repos if repo.repo_id==model_id)
+    except StopIteration:
+        # The model was not in the cache!
+        return
+
+    if repo is not None:
+        for revision in repo.revisions:
+            delete_strategy = cache_info.delete_revisions(revision.commit_hash)
+            delete_strategy.execute()
+
+
 def convert(model_id, task,
             compute_units, precision, tolerance, framework,
             push_destination, destination_model, token,
@@ -270,8 +290,8 @@ def convert(model_id, task,
         return error_str("Please provide a token to push to the Hub.", open_discussion=False)
 
     # TODO: support legacy format
-    base = Path("exported")/model_id
-    output = base/"coreml"/task
+    exported_base = Path("exported")/model_id
+    output = exported_base/"coreml"/task
     output.mkdir(parents=True, exist_ok=True)
     output = output/f"{precision}_model.mlpackage"
 
@@ -324,16 +344,17 @@ def convert(model_id, task,
             )
 
         progress(0.7, "Uploading model to Hub")
-        pr_url = push_to_hub(destination_model, base, task, precision, token=token)
+        pr_url = push_to_hub(destination_model, exported_base, task, precision, token=token)
         progress(1, "Done")
 
+        cleanup(model_id, exported_base)
+
         did_validate = _is_macos() and _macos_version() >= (12, 0)
         result = f"""### Successfully converted!
         We opened a PR to add the Core ML weights to the model repo. Please, view and merge the PR [here]({pr_url}).
 
         {f"**Note**: model could not be automatically validated as this Space is not running on macOS." if not did_validate else ""}
         """
-
         return result
     except Exception as e:
         return error_str(e, model=model_id, task=task, framework=framework, compute_units=compute_units, precision=precision, tolerance=tolerance)