Push to hub.

Suggest to open a discussion when an error occurs.

app.py

@@ -1,5 +1,7 @@
 import gradio as gr
 import json
+import subprocess
+import urllib.parse
 from pathlib import Path
 
 from huggingface_hub import hf_hub_download, HfApi
@@ -38,14 +40,53 @@ tolerance_mapping = {
 }
 tolerance_labels = list(tolerance_mapping.keys())
 
-def error_str(error, title="Error"):
-    return f"""#### {title}
-            {error}"""  if error else ""
+push_mapping = {
+    "Submit a PR to the original repo": "pr",
+    "Create a new repo": "new",
+}
+push_labels = list(push_mapping.keys())
+
+def error_str(error, title="Error", model=None, task=None, framework=None, compute_units=None, precision=None, tolerance=None, destination=None):
+    if not error: return ""
+
+    issue_title = urllib.parse.quote(f"Error converting {model}")
+    issue_description = urllib.parse.quote(f"""Conversion Settings:
+
+        Model: {model}
+        Task: {task}
+        Framework: {framework}
+        Compute Units: {compute_units}
+        Precision: {precision}
+        Tolerance: {tolerance}
+        Push to: {destination}
+
+        Error: {error}
+        """)
+    issue_url = f"https://huggingface.co/spaces/pcuenq/transformers-to-coreml/discussions/new?title={issue_title}&description={issue_description}"
+    return f"""
+        #### {title}
+        {error}
+
+        It could be that the model is not yet compatible with the Core ML exporter. Please, open a discussion on the [Hugging Face Hub]({issue_url}) to report this issue.
+        """
 
 def url_to_model_id(model_id_str):
     if not model_id_str.startswith("https://huggingface.co/"): return model_id_str
     return model_id_str.split("/")[-2] + "/" + model_id_str.split("/")[-1]
 
+def get_pr_url(api, repo_id, title):
+    try:
+        discussions = api.get_repo_discussions(repo_id=repo_id)
+    except Exception:
+        return None
+    for discussion in discussions:
+        if (
+            discussion.status == "open"
+            and discussion.is_pull_request
+            and discussion.title == title
+        ):
+            return f"https://huggingface.co/{repo_id}/discussions/{discussion.num}"
+        
 def supported_frameworks(model_id):
     """
     Return a list of supported frameworks (`PyTorch` or `TensorFlow`) for a given model_id.
@@ -82,7 +123,7 @@ def on_model_change(model):
             gr.update(visible=bool(model_type)),                                                    # Settings column
             gr.update(choices=tasks, value=tasks[0] if tasks else None),                            # Tasks
             gr.update(visible=len(frameworks)>1, choices=frameworks, value=selected_framework),     # Frameworks
-            gr.update(value=error_str(error)),                                                      # Error
+            gr.update(value=error_str(error, model=model)),                                         # Error
         )
     except Exception as e:
         error = e
@@ -121,23 +162,49 @@ def convert_model(preprocessor, model, model_coreml_config,
     progress(progress_end, desc=f"Done converting {model_label}")
 
 
-def convert(model, task, compute_units, precision, tolerance, framework, progress=gr.Progress()):
-    model = url_to_model_id(model)
+def push_to_hub(destination, directory, task, precision, token=None):
+    api = HfApi(token=token)
+    api.create_repo(destination, token=token, exist_ok=True)
+    commit_message="Add Core ML conversion"
+    api.upload_folder(
+        folder_path=directory,
+        repo_id=destination,
+        token=token,
+        create_pr=True,
+        commit_message=commit_message,
+        commit_description=f"Core ML conversion, task={task}, precision={precision}",
+    )
+
+    subprocess.run(["rm", "-rf", directory])
+    return f"""Successfully converted! We opened a PR to add the Core ML weights to the model repo.
+            Please, view and merge the PR [here]({get_pr_url(HfApi(token=token), destination, commit_message)})."""
+
+
+def convert(model_id, task,
+            compute_units, precision, tolerance, framework,
+            push_destination, destination_model, token,
+            progress=gr.Progress()):
+    model_id = url_to_model_id(model_id)
     compute_units = compute_units_mapping[compute_units]
     precision = precision_mapping[precision]
     tolerance = tolerance_mapping[tolerance]
     framework = framework_mapping[framework]
+    push_destination = push_mapping[push_destination]
+    if push_destination == "pr":
+        destination_model = model_id
+        token = None
 
     # TODO: support legacy format
-    output = Path("exported")/model/"coreml"/task
+    base = Path("exported")/model_id
+    output = base/"coreml"/task
     output.mkdir(parents=True, exist_ok=True)
     output = output/f"{precision}_model.mlpackage"
 
     try:
         progress(0, desc="Downloading model")
 
-        preprocessor = get_preprocessor(model)
-        model = FeaturesManager.get_model_from_feature(task, model, framework=framework)
+        preprocessor = get_preprocessor(model_id)
+        model = FeaturesManager.get_model_from_feature(task, model_id, framework=framework)
         _, model_coreml_config = FeaturesManager.check_supported_model_or_raise(model, feature=task)
 
         if task in ["seq2seq-lm", "speech-seq2seq"]:
@@ -152,9 +219,9 @@ def convert(model, task, compute_units, precision, tolerance, framework, progres
                 seq2seq="encoder",
                 progress=progress,
                 progress_start=0.1,
-                progress_end=0.45,
+                progress_end=0.4,
             )
-            progress(0.6, desc="Converting decoder")
+            progress(0.4, desc="Converting decoder")
             convert_model(
                 preprocessor,
                 model,
@@ -165,8 +232,8 @@ def convert(model, task, compute_units, precision, tolerance, framework, progres
                 output,
                 seq2seq="decoder",
                 progress=progress,
-                progress_start=0.45,
-                progress_end=0.8,
+                progress_start=0.4,
+                progress_end=0.7,
             )
         else:
             convert_model(
@@ -178,14 +245,15 @@ def convert(model, task, compute_units, precision, tolerance, framework, progres
                 tolerance,
                 output,
                 progress=progress,
-                progress_end=0.8,
+                progress_end=0.7,
             )
 
-        # TODO: push to hub, whatever
+        progress(0.7, "Uploading model to Hub")
+        result = push_to_hub(destination_model, base, task, precision, token=token)
         progress(1, "Done")
-        return "Done"
+        return result
     except Exception as e:
-        return error_str(e)
+        return error_str(e, model=model_id, task=task, framework=framework, compute_units=compute_units, precision=precision, tolerance=tolerance)
 
 DESCRIPTION = """
 ## Convert a transformers model to Core ML
@@ -235,6 +303,17 @@ with gr.Blocks() as demo:
                     choices=tolerance_labels,
                     value=tolerance_labels[0],
                 )
+
+                radio_push = gr.Radio(
+                    label="Destination Model",
+                    choices=push_labels,
+                    value=push_labels[0],
+                )
+                with gr.Row(visible=False) as row_destination:
+                    # TODO: public/private
+                    text_destination = gr.Textbox(label="Destination model name", value="")
+                    text_token = gr.Textbox(label="Token (write permissions)", value="")
+
                 btn_convert = gr.Button("Convert")
                 gr.Markdown("Conversion will take a few minutes.")
 
@@ -251,18 +330,25 @@ with gr.Blocks() as demo:
         
     btn_convert.click(
         fn=convert,
-        inputs=[input_model, radio_tasks, radio_compute, radio_precision, radio_tolerance, radio_framework],
+        inputs=[input_model, radio_tasks, radio_compute, radio_precision, radio_tolerance, radio_framework, radio_push, text_destination, text_token],
         outputs=error_output,
         scroll_to_output=True
     )
 
-    # gr.HTML("""
-    # <div style="border-top: 1px solid #303030;">
-    #   <br>
-    #   <p>Footer</p><br>
-    #   <p><img src="https://visitor-badge.glitch.me/badge?page_id=pcuenq.transformers-to-coreml" alt="visitors"></p>
-    # </div>
-    # """)
+    radio_push.change(
+        lambda x: gr.update(visible=x == "Create a new repo"),
+        inputs=radio_push,
+        outputs=row_destination,
+        queue=False,
+        scroll_to_output=True
+    )
+
+    gr.HTML("""
+    <div style="border-top: 0.5px solid #303030;">
+      <br>
+      <p style="color:gray;font-size:smaller;font-style:italic">Adapted from https://huggingface.co/spaces/diffusers/sd-to-diffusers/tree/main</p><br>
+    </div>
+    """)
     
 demo.queue(concurrency_count=1, max_size=10)
 demo.launch(debug=True, share=False)