Clean up

app.py

@@ -55,7 +55,7 @@ def show_warning(warning_text: str) -> gr.Blocks:
 
 
 pipe = InferencePipeline(HF_TOKEN)
-trainer = Trainer(HF_TOKEN)
+trainer = Trainer()
 
 with gr.Blocks(css='style.css') as demo:
     if IS_SHARED_UI:
@@ -75,7 +75,7 @@ with gr.Blocks(css='style.css') as demo:
             gr.Markdown('''
             - You can use this tab to upload models later if you choose not to upload models in training time or if upload in training time failed.
             ''')
-            create_upload_demo(HF_TOKEN)
+            create_upload_demo()
 
     if not HF_TOKEN:
         show_warning(HF_TOKEN_NOT_SPECIFIED_WARNING)

app_training.py

@@ -20,7 +20,6 @@ def create_training_demo(trainer: Trainer,
             lines = f.readlines()
         return ''.join(lines[-10:])
 
-    hf_token = os.getenv('HF_TOKEN')
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
@@ -48,9 +47,9 @@ def create_training_demo(trainer: Trainer,
                                                  label='Resolution',
                                                  visible=False)
 
-                    input_hf_token = gr.Text(label='Hugging Face Write Token',
-                                             placeholder='',
-                                             visible=hf_token is None)
+                    hf_token = gr.Text(label='Hugging Face Write Token',
+                                       placeholder='',
+                                       visible=os.getenv('HF_TOKEN') is None)
                     with gr.Accordion('Advanced settings', open=False):
                         num_training_steps = gr.Number(
                             label='Number of Training Steps',
@@ -150,13 +149,12 @@ def create_training_demo(trainer: Trainer,
                              delete_existing_repo,
                              upload_to,
                              remove_gpu_after_training,
-                             input_hf_token,
+                             hf_token,
                          ])
     return demo
 
 
 if __name__ == '__main__':
-    hf_token = os.getenv('HF_TOKEN')
-    trainer = Trainer(hf_token)
+    trainer = Trainer()
     demo = create_training_demo(trainer)
     demo.queue(api_open=False, max_size=1).launch()

app_upload.py

@@ -2,56 +2,21 @@
 
 from __future__ import annotations
 
-import pathlib
+import os
 
 import gradio as gr
-import slugify
 
 from constants import MODEL_LIBRARY_ORG_NAME, UploadTarget
-from uploader import Uploader
+from uploader import upload
 from utils import find_exp_dirs
 
 
-class ModelUploader(Uploader):
-    def upload_model(
-        self,
-        folder_path: str,
-        repo_name: str,
-        upload_to: str,
-        private: bool,
-        delete_existing_repo: bool,
-        input_hf_token: str | None = None,
-        return_html_link: bool = True,
-    ) -> str:
-        if not folder_path:
-            raise ValueError
-        if not repo_name:
-            repo_name = pathlib.Path(folder_path).name
-        repo_name = slugify.slugify(repo_name)
-
-        if upload_to == UploadTarget.PERSONAL_PROFILE.value:
-            organization = ''
-        elif upload_to == UploadTarget.MODEL_LIBRARY.value:
-            organization = MODEL_LIBRARY_ORG_NAME
-        else:
-            raise ValueError
-
-        return self.upload(folder_path,
-                           repo_name,
-                           organization=organization,
-                           private=private,
-                           delete_existing_repo=delete_existing_repo,
-                           input_hf_token=input_hf_token,
-                           return_html_link=return_html_link)
-
-
 def load_local_model_list() -> dict:
     choices = find_exp_dirs()
     return gr.update(choices=choices, value=choices[0] if choices else None)
 
 
-def create_upload_demo(hf_token: str | None) -> gr.Blocks:
-    uploader = ModelUploader(hf_token)
+def create_upload_demo() -> gr.Blocks:
     model_dirs = find_exp_dirs()
 
     with gr.Blocks() as demo:
@@ -72,9 +37,8 @@ def create_upload_demo(hf_token: str | None) -> gr.Blocks:
                                  choices=[_.value for _ in UploadTarget],
                                  value=UploadTarget.MODEL_LIBRARY.value)
             model_name = gr.Textbox(label='Model Name')
-            input_hf_token = gr.Text(label='Hugging Face Write Token',
-                                     placeholder='',
-                                     visible=False if hf_token else True)
+            hf_token = gr.Text(label='Hugging Face Write Token',
+                               visible=os.getenv('HF_TOKEN') is None)
         upload_button = gr.Button('Upload')
         gr.Markdown(f'''
             - You can upload your trained model to your personal profile (i.e. https://huggingface.co/{{your_username}}/{{model_name}}) or to the public [Tune-A-Video Library](https://huggingface.co/{MODEL_LIBRARY_ORG_NAME}) (i.e. https://huggingface.co/{MODEL_LIBRARY_ORG_NAME}/{{model_name}}).
@@ -86,23 +50,19 @@ def create_upload_demo(hf_token: str | None) -> gr.Blocks:
         reload_button.click(fn=load_local_model_list,
                             inputs=None,
                             outputs=model_dir)
-        upload_button.click(fn=uploader.upload_model,
+        upload_button.click(fn=upload,
                             inputs=[
                                 model_dir,
                                 model_name,
                                 upload_to,
                                 use_private_repo,
                                 delete_existing_repo,
-                                input_hf_token,
+                                hf_token,
                             ],
                             outputs=output_message)
-
     return demo
 
 
 if __name__ == '__main__':
-    import os
-
-    hf_token = os.getenv('HF_TOKEN')
-    demo = create_upload_demo(hf_token)
+    demo = create_upload_demo()
     demo.queue(api_open=False, max_size=1).launch()

constants.py

@@ -8,3 +8,4 @@ class UploadTarget(enum.Enum):
 
 MODEL_LIBRARY_ORG_NAME = 'Tune-A-Video-library'
 SAMPLE_MODEL_REPO = 'Tune-A-Video-library/a-man-is-surfing'
+URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/Tune-A-Video-library/share/YjTcaNJmKyeHFpMBioHhzBcTzCYddVErEk'

trainer.py

@@ -14,19 +14,14 @@ import torch
 from huggingface_hub import HfApi
 from omegaconf import OmegaConf
 
-from app_upload import ModelUploader
+from uploader import upload
 from utils import save_model_card
 
 sys.path.append('Tune-A-Video')
 
-URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/Tune-A-Video-library/share/YjTcaNJmKyeHFpMBioHhzBcTzCYddVErEk'
-
 
 class Trainer:
-    def __init__(self, hf_token: str | None = None):
-        self.hf_token = hf_token
-        self.model_uploader = ModelUploader(hf_token)
-
+    def __init__(self):
         self.checkpoint_dir = pathlib.Path('checkpoints')
         self.checkpoint_dir.mkdir(exist_ok=True)
 
@@ -44,12 +39,6 @@ class Trainer:
                            cwd=org_dir)
         return model_dir.as_posix()
 
-    def join_model_library_org(self, token: str) -> None:
-        subprocess.run(
-            shlex.split(
-                f'curl -X POST -H "Authorization: Bearer {token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}'
-            ))
-
     def run(
         self,
         training_video: str,
@@ -72,7 +61,7 @@ class Trainer:
         delete_existing_repo: bool,
         upload_to: str,
         remove_gpu_after_training: bool,
-        input_hf_token: str,
+        hf_token: str,
     ) -> None:
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
@@ -96,10 +85,6 @@ class Trainer:
             shutil.rmtree(output_dir, ignore_errors=True)
         output_dir.mkdir(parents=True)
 
-        if upload_to_hub:
-            self.join_model_library_org(
-                self.hf_token if self.hf_token else input_hf_token)
-
         config = OmegaConf.load('Tune-A-Video/configs/man-surfing.yaml')
         config.pretrained_model_path = self.download_base_model(base_model)
         config.output_dir = output_dir.as_posix()
@@ -146,20 +131,18 @@ class Trainer:
             f.write('Training completed!\n')
 
         if upload_to_hub:
-            upload_message = self.model_uploader.upload_model(
-                folder_path=output_dir.as_posix(),
-                repo_name=output_model_name,
-                upload_to=upload_to,
-                private=use_private_repo,
-                delete_existing_repo=delete_existing_repo,
-                input_hf_token=input_hf_token,
-                return_html_link=False)
+            upload_message = upload(local_folder_path=output_dir.as_posix(),
+                                    target_repo_name=output_model_name,
+                                    upload_to=upload_to,
+                                    private=use_private_repo,
+                                    delete_existing_repo=delete_existing_repo,
+                                    hf_token=hf_token)
             with open(self.log_file, 'a') as f:
                 f.write(upload_message)
 
         if remove_gpu_after_training:
             space_id = os.getenv('SPACE_ID')
             if space_id:
-                api = HfApi(token=self.hf_token or input_hf_token)
+                api = HfApi(token=os.getenv('HF_TOKEN') or hf_token)
                 api.request_space_hardware(repo_id=space_id,
                                            hardware='cpu-basic')

uploader.py

@@ -1,47 +1,63 @@
 from __future__ import annotations
 
+import os
+import pathlib
+import shlex
+import subprocess
+
+import slugify
 from huggingface_hub import HfApi
 
+from constants import (MODEL_LIBRARY_ORG_NAME, URL_TO_JOIN_MODEL_LIBRARY_ORG,
+                       UploadTarget)
+
+
+def join_model_library_org(hf_token: str) -> None:
+    subprocess.run(
+        shlex.split(
+            f'curl -X POST -H "Authorization: Bearer {hf_token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}'
+        ))
+
+
+def upload(local_folder_path: str,
+           target_repo_name: str,
+           upload_to: str,
+           private: bool = True,
+           delete_existing_repo: bool = False,
+           hf_token: str = '') -> str:
+    hf_token = os.getenv('HF_TOKEN') or hf_token
+    if not hf_token:
+        raise ValueError
+    api = HfApi(token=hf_token)
+
+    if not local_folder_path:
+        raise ValueError
+    if not target_repo_name:
+        target_repo_name = pathlib.Path(local_folder_path).name
+    target_repo_name = slugify.slugify(target_repo_name)
+
+    if upload_to == UploadTarget.PERSONAL_PROFILE.value:
+        organization = api.whoami()['name']
+    elif upload_to == UploadTarget.MODEL_LIBRARY.value:
+        organization = MODEL_LIBRARY_ORG_NAME
+        join_model_library_org(hf_token)
+    else:
+        raise ValueError
 
-class Uploader:
-    def __init__(self, hf_token: str | None):
-        self.hf_token = hf_token
-
-    def upload(self,
-               folder_path: str,
-               repo_name: str,
-               organization: str = '',
-               repo_type: str = 'model',
-               private: bool = True,
-               delete_existing_repo: bool = False,
-               input_hf_token: str | None = None,
-               return_html_link: bool = True) -> str:
-
-        api = HfApi(token=self.hf_token or input_hf_token)
-
-        if not folder_path:
-            raise ValueError
-        if not repo_name:
-            raise ValueError
-        if not organization:
-            organization = api.whoami()['name']
-
-        repo_id = f'{organization}/{repo_name}'
-        if delete_existing_repo:
-            try:
-                api.delete_repo(repo_id, repo_type=repo_type)
-            except Exception:
-                pass
+    repo_id = f'{organization}/{target_repo_name}'
+    if delete_existing_repo:
         try:
-            api.create_repo(repo_id, repo_type=repo_type, private=private)
-            api.upload_folder(repo_id=repo_id,
-                              folder_path=folder_path,
-                              path_in_repo='.',
-                              repo_type=repo_type)
-            url = f'https://huggingface.co/{repo_id}'
-            if return_html_link:
-                url = f'<a href="{url}" target="_blank">{url}</a>'
-            message = f'Your model was successfully uploaded to {url}.'
-        except Exception as e:
-            message = str(e)
-        return message
+            api.delete_repo(repo_id, repo_type='model')
+        except Exception:
+            pass
+    try:
+        api.create_repo(repo_id, repo_type='model', private=private)
+        api.upload_folder(repo_id=repo_id,
+                          folder_path=local_folder_path,
+                          path_in_repo='.',
+                          repo_type='model')
+        url = f'https://huggingface.co/{repo_id}'
+        message = f'Your model was successfully uploaded to {url}.'
+    except Exception as e:
+        message = str(e)
+    return message