Add upload functionality

app.py

@@ -16,6 +16,7 @@ import torch
 
 from inference import InferencePipeline
 from trainer import Trainer
+from uploader import upload
 
 TITLE = '# LoRA + StableDiffusion Training UI'
 DESCRIPTION = 'This is an unofficial demo for [https://github.com/cloneofsimo/lora](https://github.com/cloneofsimo/lora).'
@@ -205,6 +206,27 @@ def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
     return demo
 
 
+def create_upload_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        model_name = gr.Textbox(label='Model Name')
+        hf_token = gr.Textbox(
+            label='Hugging Face Token (with write permission)')
+        upload_button = gr.Button('Upload')
+        with gr.Box():
+            gr.Markdown('Message')
+            result = gr.Markdown()
+        gr.Markdown('''
+            - You can upload your trained model to your private Model repo (i.e. https://huggingface.co/{your_username}/{model_name}).
+            - You can find your Hugging Face token [here](https://huggingface.co/settings/tokens).
+            ''')
+
+    upload_button.click(fn=upload,
+                        inputs=[model_name, hf_token],
+                        outputs=result)
+
+    return demo
+
+
 pipe = InferencePipeline()
 trainer = Trainer()
 
@@ -222,6 +244,8 @@ with gr.Blocks(css='style.css') as demo:
             create_training_demo(trainer, pipe)
         with gr.TabItem('Test'):
             create_inference_demo(pipe)
+        with gr.TabItem('Upload'):
+            create_upload_demo()
 
     with gr.Accordion('Usage', open=False):
         gr.Markdown(USAGE_INFO)

uploader.py

@@ -0,0 +1,21 @@
+import gradio as gr
+from huggingface_hub import HfApi
+
+
+def upload(model_name: str, hf_token: str) -> None:
+    api = HfApi(token=hf_token)
+    user_name = api.whoami()['name']
+    model_id = f'{user_name}/{model_name}'
+    try:
+        api.create_repo(model_id, repo_type='model', private=True)
+        api.upload_folder(repo_id=model_id,
+                          folder_path='results',
+                          path_in_repo='results',
+                          repo_type='model')
+        url = f'https://huggingface.co/{model_id}'
+        message = f'Your model was successfully uploaded to [{url}]({url}).'
+    except Exception as e:
+        message = str(e)
+        #raise gr.Error(str(e))
+
+    return gr.update(value=message, visible=True)