simplified-ui

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 exclude: patch
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.4.0
+  rev: v4.2.0
   hooks:
   - id: check-executables-have-shebangs
   - id: check-json
@@ -9,43 +9,29 @@ repos:
   - id: check-shebang-scripts-are-executable
   - id: check-toml
   - id: check-yaml
+  - id: double-quote-string-fixer
   - id: end-of-file-fixer
   - id: mixed-line-ending
-    args: ["--fix=lf"]
+    args: ['--fix=lf']
   - id: requirements-txt-fixer
   - id: trailing-whitespace
 - repo: https://github.com/myint/docformatter
-  rev: v1.7.5
+  rev: v1.4
   hooks:
   - id: docformatter
-    args: ["--in-place"]
+    args: ['--in-place']
 - repo: https://github.com/pycqa/isort
   rev: 5.12.0
   hooks:
     - id: isort
-      args: ["--profile", "black"]
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.5.1
+  rev: v0.991
   hooks:
     - id: mypy
-      args: ["--ignore-missing-imports"]
-      additional_dependencies: ["types-python-slugify", "types-requests", "types-PyYAML"]
-- repo: https://github.com/psf/black
-  rev: 23.9.0
+      args: ['--ignore-missing-imports']
+      additional_dependencies: ['types-python-slugify']
+- repo: https://github.com/google/yapf
+  rev: v0.32.0
   hooks:
-    - id: black
-      language_version: python3.10
-      args: ["--line-length", "119"]
-- repo: https://github.com/kynan/nbstripout
-  rev: 0.6.1
-  hooks:
-    - id: nbstripout
-      args: ["--extra-keys", "metadata.interpreter metadata.kernelspec cell.metadata.pycharm"]
-- repo: https://github.com/nbQA-dev/nbQA
-  rev: 1.7.0
-  hooks:
-    - id: nbqa-black
-    - id: nbqa-pyupgrade
-      args: ["--py37-plus"]
-    - id: nbqa-isort
-      args: ["--float-to-top"]
+  - id: yapf
+    args: ['--parallel', '--in-place']

.style.yapf

@@ -0,0 +1,5 @@
+[style]
+based_on_style = pep8
+blank_line_before_nested_class_or_def = false
+spaces_before_comment = 2
+split_before_logical_operator = true

.vscode/settings.json

@@ -1,21 +0,0 @@
-{
-    "[python]": {
-        "editor.defaultFormatter": "ms-python.black-formatter",
-        "editor.formatOnType": true,
-        "editor.codeActionsOnSave": {
-            "source.organizeImports": true
-        }
-    },
-    "black-formatter.args": [
-        "--line-length=119"
-    ],
-    "isort.args": ["--profile", "black"],
-    "flake8.args": [
-        "--max-line-length=119"
-    ],
-    "ruff.args": [
-        "--line-length=119"
-    ],
-    "editor.formatOnSave": true,
-    "files.insertFinalNewline": true
-}

Dockerfile

@@ -35,7 +35,7 @@ WORKDIR ${HOME}/app
 
 RUN curl https://pyenv.run | bash
 ENV PATH=${HOME}/.pyenv/shims:${HOME}/.pyenv/bin:${PATH}
-ARG PYTHON_VERSION=3.10.11
+ENV PYTHON_VERSION=3.10.9
 RUN pyenv install ${PYTHON_VERSION} && \
     pyenv global ${PYTHON_VERSION} && \
     pyenv rehash && \
@@ -44,8 +44,6 @@ RUN pyenv install ${PYTHON_VERSION} && \
 RUN pip install --no-cache-dir -U torch==1.13.1 torchvision==0.14.1
 COPY --chown=1000 requirements.txt /tmp/requirements.txt
 RUN pip install --no-cache-dir -U -r /tmp/requirements.txt
-COPY --chown=1000 requirements-monitor.txt /tmp/requirements-monitor.txt
-RUN pip install --no-cache-dir -U -r /tmp/requirements-monitor.txt
 
 COPY --chown=1000 . ${HOME}/app
 RUN cd Tune-A-Video && patch -p1 < ../patch

app.py

@@ -9,43 +9,43 @@ import gradio as gr
 import torch
 
 from app_inference import create_inference_demo
-from app_system_monitor import create_monitor_demo
 from app_training import create_training_demo
 from app_upload import create_upload_demo
 from inference import InferencePipeline
 from trainer import Trainer
 
-TITLE = "# [Tune-A-Video](https://tuneavideo.github.io/)"
+TITLE = '# [Tune-A-Video](https://tuneavideo.github.io/) UI'
 
-ORIGINAL_SPACE_ID = "Tune-A-Video-library/Tune-A-Video-Training-UI"
-SPACE_ID = os.getenv("SPACE_ID")
-GPU_DATA = getoutput("nvidia-smi")
-SHARED_UI_WARNING = f"""## Attention - Training doesn't work in this shared UI. You can duplicate and use it with a paid private T4 GPU.
+ORIGINAL_SPACE_ID = 'Tune-A-Video-library/Tune-A-Video-Training-UI'
+SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
+GPU_DATA = getoutput('nvidia-smi')
+SHARED_UI_WARNING = f'''## Attention - Training doesn't work in this shared UI. You can duplicate and use it with a paid private T4 GPU.
 
 <center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></center>
-"""
+'''
 
-IS_SHARED_UI = SPACE_ID == ORIGINAL_SPACE_ID
-if os.getenv("SYSTEM") == "spaces" and SPACE_ID != ORIGINAL_SPACE_ID:
+if os.getenv('SYSTEM') == 'spaces' and SPACE_ID != ORIGINAL_SPACE_ID:
     SETTINGS = f'<a href="https://huggingface.co/spaces/{SPACE_ID}/settings">Settings</a>'
 else:
-    SETTINGS = "Settings"
+    SETTINGS = 'Settings'
 
-INVALID_GPU_WARNING = f"""## Attention - the specified GPU is invalid. Training may not work. Make sure you have selected a `T4 GPU` for this task."""
+INVALID_GPU_WARNING = f'''## Attention - the specified GPU is invalid. Training may not work. Make sure you have selected a `T4 GPU` for this task.'''
 
-CUDA_NOT_AVAILABLE_WARNING = f"""## Attention - Running on CPU.
+CUDA_NOT_AVAILABLE_WARNING = f'''## Attention - Running on CPU.
 <center>
 You can assign a GPU in the {SETTINGS} tab if you are running this on HF Spaces.
 You can use "T4 small/medium" to run this demo.
 </center>
-"""
+'''
 
-HF_TOKEN_NOT_SPECIFIED_WARNING = f"""The environment variable `HF_TOKEN` is not specified. Feel free to specify your Hugging Face token with write permission if you don't want to manually provide it for every run.
-
-You can check and create your Hugging Face tokens <a href="https://huggingface.co/settings/tokens" target="_blank">here</a>. You can specify environment variables in the "Repository secrets" section of the {SETTINGS} tab.
-"""
+HF_TOKEN_NOT_SPECIFIED_WARNING = f'''The environment variable `HF_TOKEN` is not specified. Feel free to specify your Hugging Face token with write permission if you don't want to manually provide it for every run. 
+<center>
+You can check and create your Hugging Face tokens <a href="https://huggingface.co/settings/tokens" target="_blank">here</a>.
+You can specify environment variables in the "Repository secrets" section of the {SETTINGS} tab.
+</center>
+'''
 
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv('HF_TOKEN')
 
 
 def show_warning(warning_text: str) -> gr.Blocks:
@@ -56,36 +56,30 @@ def show_warning(warning_text: str) -> gr.Blocks:
 
 
 pipe = InferencePipeline(HF_TOKEN)
-trainer = Trainer()
+trainer = Trainer(HF_TOKEN)
 
-with gr.Blocks(css="style.css") as demo:
-    if IS_SHARED_UI:
+with gr.Blocks(css='style.css') as demo:
+    if SPACE_ID == ORIGINAL_SPACE_ID:
         show_warning(SHARED_UI_WARNING)
     elif not torch.cuda.is_available():
         show_warning(CUDA_NOT_AVAILABLE_WARNING)
-    elif "T4" not in GPU_DATA:
+    elif(not "T4" in GPU_DATA):
         show_warning(INVALID_GPU_WARNING)
+    
 
     gr.Markdown(TITLE)
     with gr.Tabs():
-        with gr.TabItem("Train"):
-            create_training_demo(trainer, pipe, disable_run_button=IS_SHARED_UI)
-        with gr.TabItem("Run"):
-            create_inference_demo(pipe, HF_TOKEN, disable_run_button=IS_SHARED_UI)
-        with gr.TabItem("Upload"):
-            gr.Markdown(
-                """
+        with gr.TabItem('Train'):
+            create_training_demo(trainer, pipe)
+        with gr.TabItem('Run'):
+            create_inference_demo(pipe, HF_TOKEN)
+        with gr.TabItem('Upload'):
+            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(disable_run_button=IS_SHARED_UI)
-
-    with gr.Row():
-        if not IS_SHARED_UI and not os.getenv("DISABLE_SYSTEM_MONITOR"):
-            with gr.Accordion(label="System info", open=False):
-                create_monitor_demo()
-
+            ''')
+            create_upload_demo(HF_TOKEN)
+    
     if not HF_TOKEN:
         show_warning(HF_TOKEN_NOT_SPECIFIED_WARNING)
 
-demo.queue(api_open=False, max_size=1).launch()
+demo.queue(max_size=1).launch(share=False)

app_inference.py

@@ -14,7 +14,7 @@ from utils import find_exp_dirs
 
 class ModelSource(enum.Enum):
     HUB_LIB = UploadTarget.MODEL_LIBRARY.value
-    LOCAL = "Local"
+    LOCAL = 'Local'
 
 
 class InferenceUtil:
@@ -23,13 +23,18 @@ class InferenceUtil:
 
     def load_hub_model_list(self) -> dict:
         api = HfApi(token=self.hf_token)
-        choices = [info.modelId for info in api.list_models(author=MODEL_LIBRARY_ORG_NAME)]
-        return gr.update(choices=choices, value=choices[0] if choices else None)
+        choices = [
+            info.modelId
+            for info in api.list_models(author=MODEL_LIBRARY_ORG_NAME)
+        ]
+        return gr.update(choices=choices,
+                         value=choices[0] if choices else None)
 
     @staticmethod
     def load_local_model_list() -> dict:
         choices = find_exp_dirs()
-        return gr.update(choices=choices, value=choices[0] if choices else None)
+        return gr.update(choices=choices,
+                         value=choices[0] if choices else None)
 
     def reload_model_list(self, model_source: str) -> dict:
         if model_source == ModelSource.HUB_LIB.value:
@@ -43,21 +48,21 @@ class InferenceUtil:
         try:
             card = InferencePipeline.get_model_card(model_id, self.hf_token)
         except Exception:
-            return "", ""
-        base_model = getattr(card.data, "base_model", "")
-        training_prompt = getattr(card.data, "training_prompt", "")
+            return '', ''
+        base_model = getattr(card.data, 'base_model', '')
+        training_prompt = getattr(card.data, 'training_prompt', '')
         return base_model, training_prompt
 
-    def reload_model_list_and_update_model_info(self, model_source: str) -> tuple[dict, str, str]:
+    def reload_model_list_and_update_model_info(
+            self, model_source: str) -> tuple[dict, str, str]:
         model_list_update = self.reload_model_list(model_source)
-        model_list = model_list_update["choices"]
-        model_info = self.load_model_info(model_list[0] if model_list else "")
+        model_list = model_list_update['choices']
+        model_info = self.load_model_info(model_list[0] if model_list else '')
         return model_list_update, *model_info
 
 
-def create_inference_demo(
-    pipe: InferencePipeline, hf_token: str | None = None, disable_run_button: bool = False
-) -> gr.Blocks:
+def create_inference_demo(pipe: InferencePipeline,
+                          hf_token: str | None = None) -> gr.Blocks:
     app = InferenceUtil(hf_token)
 
     with gr.Blocks() as demo:
@@ -65,60 +70,83 @@ def create_inference_demo(
             with gr.Column():
                 with gr.Box():
                     model_source = gr.Radio(
-                        label="Model Source", choices=[_.value for _ in ModelSource], value=ModelSource.HUB_LIB.value
-                    )
-                    reload_button = gr.Button("Reload Model List")
-                    model_id = gr.Dropdown(label="Model ID", choices=None, value=None)
-                    with gr.Accordion(label="Model info (Base model and prompt used for training)", open=False):
+                        label='Model Source',
+                        choices=[_.value for _ in ModelSource],
+                        value=ModelSource.HUB_LIB.value)
+                    reload_button = gr.Button('Reload Model List')
+                    model_id = gr.Dropdown(label='Model ID',
+                                           choices=None,
+                                           value=None)
+                    with gr.Accordion(
+                            label=
+                            'Model info (Base model and prompt used for training)',
+                            open=False):
                         with gr.Row():
-                            base_model_used_for_training = gr.Text(label="Base model", interactive=False)
-                            prompt_used_for_training = gr.Text(label="Training prompt", interactive=False)
-                prompt = gr.Textbox(label="Prompt", max_lines=1, placeholder='Example: "A panda is surfing"')
-                video_length = gr.Slider(label="Video length", minimum=4, maximum=12, step=1, value=8)
-                fps = gr.Slider(label="FPS", minimum=1, maximum=12, step=1, value=1)
-                seed = gr.Slider(label="Seed", minimum=0, maximum=100000, step=1, value=0)
-                with gr.Accordion("Advanced options", open=False):
-                    num_steps = gr.Slider(label="Number of Steps", minimum=0, maximum=100, step=1, value=50)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0, maximum=50, step=0.1, value=7.5)
-
-                run_button = gr.Button("Generate", interactive=not disable_run_button)
-
-                gr.Markdown(
-                    """
+                            base_model_used_for_training = gr.Text(
+                                label='Base model', interactive=False)
+                            prompt_used_for_training = gr.Text(
+                                label='Training prompt', interactive=False)
+                prompt = gr.Textbox(
+                    label='Prompt',
+                    max_lines=1,
+                    placeholder='Example: "A panda is surfing"')
+                video_length = gr.Slider(label='Video length',
+                                         minimum=4,
+                                         maximum=12,
+                                         step=1,
+                                         value=8)
+                fps = gr.Slider(label='FPS',
+                                minimum=1,
+                                maximum=12,
+                                step=1,
+                                value=1)
+                seed = gr.Slider(label='Seed',
+                                 minimum=0,
+                                 maximum=100000,
+                                 step=1,
+                                 value=0)
+                with gr.Accordion('Other Parameters', open=False):
+                    num_steps = gr.Slider(label='Number of Steps',
+                                          minimum=0,
+                                          maximum=100,
+                                          step=1,
+                                          value=50)
+                    guidance_scale = gr.Slider(label='CFG Scale',
+                                               minimum=0,
+                                               maximum=50,
+                                               step=0.1,
+                                               value=7.5)
+
+                run_button = gr.Button('Generate')
+
+                gr.Markdown('''
                 - After training, you can press "Reload Model List" button to load your trained model names.
                 - It takes a few minutes to download model first.
                 - Expected time to generate an 8-frame video: 70 seconds with T4, 24 seconds with A10G, (10 seconds with A100)
-                """
-                )
+                ''')
             with gr.Column():
-                result = gr.Video(label="Result")
-
-        model_source.change(
-            fn=app.reload_model_list_and_update_model_info,
-            inputs=model_source,
-            outputs=[
-                model_id,
-                base_model_used_for_training,
-                prompt_used_for_training,
-            ],
-        )
-        reload_button.click(
-            fn=app.reload_model_list_and_update_model_info,
-            inputs=model_source,
-            outputs=[
-                model_id,
-                base_model_used_for_training,
-                prompt_used_for_training,
-            ],
-        )
-        model_id.change(
-            fn=app.load_model_info,
-            inputs=model_id,
-            outputs=[
-                base_model_used_for_training,
-                prompt_used_for_training,
-            ],
-        )
+                result = gr.Video(label='Result')
+
+        model_source.change(fn=app.reload_model_list_and_update_model_info,
+                            inputs=model_source,
+                            outputs=[
+                                model_id,
+                                base_model_used_for_training,
+                                prompt_used_for_training,
+                            ])
+        reload_button.click(fn=app.reload_model_list_and_update_model_info,
+                            inputs=model_source,
+                            outputs=[
+                                model_id,
+                                base_model_used_for_training,
+                                prompt_used_for_training,
+                            ])
+        model_id.change(fn=app.load_model_info,
+                        inputs=model_id,
+                        outputs=[
+                            base_model_used_for_training,
+                            prompt_used_for_training,
+                        ])
         inputs = [
             model_id,
             prompt,
@@ -133,10 +161,10 @@ def create_inference_demo(
     return demo
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     import os
 
-    hf_token = os.getenv("HF_TOKEN")
+    hf_token = os.getenv('HF_TOKEN')
     pipe = InferencePipeline(hf_token)
     demo = create_inference_demo(pipe, hf_token)
-    demo.queue(api_open=False, max_size=10).launch()
+    demo.queue(max_size=10).launch(share=False)

app_system_monitor.py

@@ -1,86 +0,0 @@
-#!/usr/bin/env python
-
-from __future__ import annotations
-
-import collections
-
-import gradio as gr
-import nvitop
-import pandas as pd
-import plotly.express as px
-import psutil
-
-
-class SystemMonitor:
-    MAX_SIZE = 61
-
-    def __init__(self):
-        self.devices = nvitop.Device.all()
-        self.cpu_memory_usage = collections.deque([0 for _ in range(self.MAX_SIZE)], maxlen=self.MAX_SIZE)
-        self.cpu_memory_usage_str = ""
-        self.gpu_memory_usage = collections.deque([0 for _ in range(self.MAX_SIZE)], maxlen=self.MAX_SIZE)
-        self.gpu_util = collections.deque([0 for _ in range(self.MAX_SIZE)], maxlen=self.MAX_SIZE)
-        self.gpu_memory_usage_str = ""
-        self.gpu_util_str = ""
-
-    def update(self) -> None:
-        self.update_cpu()
-        self.update_gpu()
-
-    def update_cpu(self) -> None:
-        memory = psutil.virtual_memory()
-        self.cpu_memory_usage.append(memory.percent)
-        self.cpu_memory_usage_str = (
-            f"{memory.used / 1024**3:0.2f}GiB / {memory.total / 1024**3:0.2f}GiB ({memory.percent}%)"
-        )
-
-    def update_gpu(self) -> None:
-        if not self.devices:
-            return
-        device = self.devices[0]
-        self.gpu_memory_usage.append(device.memory_percent())
-        self.gpu_util.append(device.gpu_utilization())
-        self.gpu_memory_usage_str = f"{device.memory_usage()} ({device.memory_percent()}%)"
-        self.gpu_util_str = f"{device.gpu_utilization()}%"
-
-    def get_json(self) -> dict[str, str]:
-        return {
-            "CPU memory usage": self.cpu_memory_usage_str,
-            "GPU memory usage": self.gpu_memory_usage_str,
-            "GPU Util": self.gpu_util_str,
-        }
-
-    def get_graph_data(self) -> dict[str, list[int | float]]:
-        return {
-            "index": list(range(-self.MAX_SIZE + 1, 1)),
-            "CPU memory usage": self.cpu_memory_usage,
-            "GPU memory usage": self.gpu_memory_usage,
-            "GPU Util": self.gpu_util,
-        }
-
-    def get_graph(self):
-        df = pd.DataFrame(self.get_graph_data())
-        return px.line(
-            df,
-            x="index",
-            y=[
-                "CPU memory usage",
-                "GPU memory usage",
-                "GPU Util",
-            ],
-            range_y=[-5, 105],
-        ).update_layout(xaxis_title="Time", yaxis_title="Percentage")
-
-
-def create_monitor_demo() -> gr.Blocks:
-    monitor = SystemMonitor()
-    with gr.Blocks() as demo:
-        gr.JSON(value=monitor.update, every=1, visible=False)
-        gr.JSON(value=monitor.get_json, show_label=False, every=1)
-        gr.Plot(value=monitor.get_graph, show_label=False, every=1)
-    return demo
-
-
-if __name__ == "__main__":
-    demo = create_monitor_demo()
-    demo.queue(api_open=False).launch()

app_training.py

@@ -6,125 +6,135 @@ import os
 
 import gradio as gr
 
-from constants import UploadTarget
+from constants import MODEL_LIBRARY_ORG_NAME, SAMPLE_MODEL_REPO, UploadTarget
 from inference import InferencePipeline
 from trainer import Trainer
 
 
-def create_training_demo(
-    trainer: Trainer, pipe: InferencePipeline | None = None, disable_run_button: bool = False
-) -> gr.Blocks:
-    def read_log() -> str:
-        with open(trainer.log_file) as f:
-            lines = f.readlines()
-        return "".join(lines[-10:])
-
+def create_training_demo(trainer: Trainer,
+                         pipe: InferencePipeline | None = None) -> gr.Blocks:
+    hf_token = os.getenv('HF_TOKEN')
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 with gr.Box():
-                    gr.Markdown("Training Data")
-                    training_video = gr.File(label="Training video")
-                    training_prompt = gr.Textbox(label="Training prompt", max_lines=1, placeholder="A man is surfing")
-                    gr.Markdown(
-                        """
+                    gr.Markdown('Training Data')
+                    training_video = gr.File(label='Training video')
+                    training_prompt = gr.Textbox(
+                        label='Training prompt',
+                        max_lines=1,
+                        placeholder='A man is surfing')
+                    gr.Markdown('''
                         - Upload a video and write a `Training Prompt` that describes the video.
-                        """
-                    )
+                        ''')
 
             with gr.Column():
                 with gr.Box():
-                    gr.Markdown("Training Parameters")
+                    gr.Markdown('Training Parameters')
                     with gr.Row():
-                        base_model = gr.Text(label="Base Model", value="CompVis/stable-diffusion-v1-4", max_lines=1)
-                        resolution = gr.Dropdown(
-                            choices=["512", "768"], value="512", label="Resolution", visible=False
-                        )
-
-                    hf_token = gr.Text(
-                        label="Hugging Face Write Token", type="password", visible=os.getenv("HF_TOKEN") is None
-                    )
-                    with gr.Accordion(label="Advanced options", open=False):
-                        num_training_steps = gr.Number(label="Number of Training Steps", value=300, precision=0)
-                        learning_rate = gr.Number(label="Learning Rate", value=0.000035)
+                        base_model = gr.Text(label='Base Model',
+                                             value='CompVis/stable-diffusion-v1-4',
+                                             max_lines=1)
+                        resolution = gr.Dropdown(choices=['512', '768'],
+                                                 value='512',
+                                                 label='Resolution',
+                                                 visible=False)
+                    
+                    input_token = gr.Text(label="Hugging Face Write Token", placeholder="", visible=False if hf_token else True)
+                    with gr.Accordion("Advanced settings", open=False):
+                        num_training_steps = gr.Number(
+                            label='Number of Training Steps', value=300, precision=0)
+                        learning_rate = gr.Number(label='Learning Rate',
+                                                  value=0.000035)
                         gradient_accumulation = gr.Number(
-                            label="Number of Gradient Accumulation", value=1, precision=0
-                        )
-                        seed = gr.Slider(label="Seed", minimum=0, maximum=100000, step=1, randomize=True, value=0)
-                        fp16 = gr.Checkbox(label="FP16", value=True)
-                        use_8bit_adam = gr.Checkbox(label="Use 8bit Adam", value=False)
-                        checkpointing_steps = gr.Number(label="Checkpointing Steps", value=1000, precision=0)
-                        validation_epochs = gr.Number(label="Validation Epochs", value=100, precision=0)
-                    gr.Markdown(
-                        """
+                            label='Number of Gradient Accumulation',
+                            value=1,
+                            precision=0)
+                        seed = gr.Slider(label='Seed',
+                                         minimum=0,
+                                         maximum=100000,
+                                         step=1,
+                                         randomize=True,
+                                         value=0)
+                        fp16 = gr.Checkbox(label='FP16', value=True)
+                        use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=False)
+                        checkpointing_steps = gr.Number(label='Checkpointing Steps',
+                                                        value=1000,
+                                                        precision=0)
+                        validation_epochs = gr.Number(label='Validation Epochs',
+                                                      value=100,
+                                                      precision=0)
+                    gr.Markdown('''
                         - The base model must be a Stable Diffusion model compatible with [diffusers](https://github.com/huggingface/diffusers) library.
                         - Expected time to train a model for 300 steps: ~20 minutes with T4
                         - You can check the training status by pressing the "Open logs" button if you are running this on your Space.
-                        """
-                    )
-
+                        ''')
+                    
         with gr.Row():
             with gr.Column():
-                gr.Markdown("Output Model")
-                output_model_name = gr.Text(label="Name of your model", placeholder="The surfer man", max_lines=1)
-                validation_prompt = gr.Text(
-                    label="Validation Prompt", placeholder="prompt to test the model, e.g: a dog is surfing"
-                )
+                gr.Markdown('Output Model')
+                output_model_name = gr.Text(label='Name of your model',
+                                            placeholder='The surfer man',
+                                            max_lines=1)
+                validation_prompt = gr.Text(label='Validation Prompt', placeholder='prompt to test the model, e.g: a dog is surfing')
             with gr.Column():
-                gr.Markdown("Upload Settings")
+                gr.Markdown('Upload Settings')
                 with gr.Row():
-                    upload_to_hub = gr.Checkbox(label="Upload model to Hub", value=True)
-                    use_private_repo = gr.Checkbox(label="Private", value=True)
-                    delete_existing_repo = gr.Checkbox(label="Delete existing repo of the same name", value=False)
+                    upload_to_hub = gr.Checkbox(
+                        label='Upload model to Hub', value=True)
+                    use_private_repo = gr.Checkbox(label='Private',
+                                                   value=True)
+                    delete_existing_repo = gr.Checkbox(
+                        label='Delete existing repo of the same name',
+                        value=False)
                     upload_to = gr.Radio(
-                        label="Upload to",
+                        label='Upload to',
                         choices=[_.value for _ in UploadTarget],
-                        value=UploadTarget.MODEL_LIBRARY.value,
-                    )
-
-        pause_space_after_training = gr.Checkbox(
-            label="Pause this Space after training",
+                        value=UploadTarget.MODEL_LIBRARY.value)
+                    
+        remove_gpu_after_training = gr.Checkbox(
+            label='Remove GPU after training',
             value=False,
-            interactive=bool(os.getenv("SPACE_ID")),
-            visible=False,
-        )
-        run_button = gr.Button("Start Training", interactive=not disable_run_button)
-
+            interactive=bool(os.getenv('SPACE_ID')),
+            visible=False)
+        run_button = gr.Button('Start Training')
+    
         with gr.Box():
-            gr.Text(label="Log", value=read_log, lines=10, max_lines=10, every=1)
+            gr.Markdown('Output message')
+            output_message = gr.Markdown()
 
         if pipe is not None:
             run_button.click(fn=pipe.clear)
-        run_button.click(
-            fn=trainer.run,
-            inputs=[
-                training_video,
-                training_prompt,
-                output_model_name,
-                delete_existing_repo,
-                validation_prompt,
-                base_model,
-                resolution,
-                num_training_steps,
-                learning_rate,
-                gradient_accumulation,
-                seed,
-                fp16,
-                use_8bit_adam,
-                checkpointing_steps,
-                validation_epochs,
-                upload_to_hub,
-                use_private_repo,
-                delete_existing_repo,
-                upload_to,
-                pause_space_after_training,
-                hf_token,
-            ],
-        )
+        run_button.click(fn=trainer.run,
+                         inputs=[
+                             training_video,
+                             training_prompt,
+                             output_model_name,
+                             delete_existing_repo,
+                             validation_prompt,
+                             base_model,
+                             resolution,
+                             num_training_steps,
+                             learning_rate,
+                             gradient_accumulation,
+                             seed,
+                             fp16,
+                             use_8bit_adam,
+                             checkpointing_steps,
+                             validation_epochs,
+                             upload_to_hub,
+                             use_private_repo,
+                             delete_existing_repo,
+                             upload_to,
+                             remove_gpu_after_training,
+                             input_token
+                         ],
+                         outputs=output_message)
     return demo
 
 
-if __name__ == "__main__":
-    trainer = Trainer()
+if __name__ == '__main__':
+    hf_token = os.getenv('HF_TOKEN')
+    trainer = Trainer(hf_token)
     demo = create_training_demo(trainer)
-    demo.queue(api_open=False, max_size=1).launch()
+    demo.queue(max_size=1).launch(share=False)

app_upload.py

@@ -2,68 +2,103 @@
 
 from __future__ import annotations
 
-import os
+import pathlib
 
 import gradio as gr
+import slugify
 
 from constants import MODEL_LIBRARY_ORG_NAME, UploadTarget
-from uploader import upload
+from uploader import Uploader
 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_token: str | None = None,
+    ) -> 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_token=input_token)
+
+
 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(disable_run_button: bool = False) -> gr.Blocks:
+def create_upload_demo(hf_token: str | None) -> gr.Blocks:
+    uploader = ModelUploader(hf_token)
     model_dirs = find_exp_dirs()
 
     with gr.Blocks() as demo:
         with gr.Box():
-            gr.Markdown("Local Models")
-            reload_button = gr.Button("Reload Model List")
+            gr.Markdown('Local Models')
+            reload_button = gr.Button('Reload Model List')
             model_dir = gr.Dropdown(
-                label="Model names", choices=model_dirs, value=model_dirs[0] if model_dirs else None
-            )
+                label='Model names',
+                choices=model_dirs,
+                value=model_dirs[0] if model_dirs else None)
         with gr.Box():
-            gr.Markdown("Upload Settings")
+            gr.Markdown('Upload Settings')
             with gr.Row():
-                use_private_repo = gr.Checkbox(label="Private", value=True)
-                delete_existing_repo = gr.Checkbox(label="Delete existing repo of the same name", value=False)
-            upload_to = gr.Radio(
-                label="Upload to", choices=[_.value for _ in UploadTarget], value=UploadTarget.MODEL_LIBRARY.value
-            )
-            model_name = gr.Textbox(label="Model Name")
-            hf_token = gr.Text(
-                label="Hugging Face Write Token", type="password", visible=os.getenv("HF_TOKEN") is None
-            )
-        upload_button = gr.Button("Upload", interactive=not disable_run_button)
-        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}}`).
-            """
-        )
+                use_private_repo = gr.Checkbox(label='Private', value=True)
+                delete_existing_repo = gr.Checkbox(
+                    label='Delete existing repo of the same name', value=False)
+            upload_to = gr.Radio(label='Upload to',
+                                 choices=[_.value for _ in UploadTarget],
+                                 value=UploadTarget.MODEL_LIBRARY.value)
+            model_name = gr.Textbox(label='Model Name')
+            input_token = gr.Text(label="Hugging Face Write Token", placeholder="", visible=False if hf_token else True)
+        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}}).
+            ''')
         with gr.Box():
-            gr.Markdown("Output message")
+            gr.Markdown('Output message')
             output_message = gr.Markdown()
 
-        reload_button.click(fn=load_local_model_list, inputs=None, outputs=model_dir)
-        upload_button.click(
-            fn=upload,
-            inputs=[
-                model_dir,
-                model_name,
-                upload_to,
-                use_private_repo,
-                delete_existing_repo,
-                hf_token,
-            ],
-            outputs=output_message,
-        )
+        reload_button.click(fn=load_local_model_list,
+                            inputs=None,
+                            outputs=model_dir)
+        upload_button.click(fn=uploader.upload_model,
+                            inputs=[
+                                model_dir,
+                                model_name,
+                                upload_to,
+                                use_private_repo,
+                                delete_existing_repo,
+                                input_token,
+                            ],
+                            outputs=output_message)
+
     return demo
 
 
-if __name__ == "__main__":
-    demo = create_upload_demo()
-    demo.queue(api_open=False, max_size=1).launch()
+if __name__ == '__main__':
+    import os
+
+    hf_token = os.getenv('HF_TOKEN')
+    demo = create_upload_demo(hf_token)
+    demo.queue(max_size=1).launch(share=False)

constants.py

@@ -2,12 +2,9 @@ import enum
 
 
 class UploadTarget(enum.Enum):
-    PERSONAL_PROFILE = "Personal Profile"
-    MODEL_LIBRARY = "Tune-A-Video Library"
+    PERSONAL_PROFILE = 'Personal Profile'
+    MODEL_LIBRARY = 'Tune-A-Video Library'
 
 
-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"
-)
+MODEL_LIBRARY_ORG_NAME = 'Tune-A-Video-library'
+SAMPLE_MODEL_REPO = 'Tune-A-Video-library/a-man-is-surfing'

inference.py

@@ -13,7 +13,7 @@ from diffusers.utils.import_utils import is_xformers_available
 from einops import rearrange
 from huggingface_hub import ModelCard
 
-sys.path.append("Tune-A-Video")
+sys.path.append('Tune-A-Video')
 
 from tuneavideo.models.unet import UNet3DConditionModel
 from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
@@ -23,7 +23,8 @@ class InferencePipeline:
     def __init__(self, hf_token: str | None = None):
         self.hf_token = hf_token
         self.pipe = None
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(
+            'cuda:0' if torch.cuda.is_available() else 'cpu')
         self.model_id = None
 
     def clear(self) -> None:
@@ -38,9 +39,10 @@ class InferencePipeline:
         return pathlib.Path(model_id).exists()
 
     @staticmethod
-    def get_model_card(model_id: str, hf_token: str | None = None) -> ModelCard:
+    def get_model_card(model_id: str,
+                       hf_token: str | None = None) -> ModelCard:
         if InferencePipeline.check_if_model_is_local(model_id):
-            card_path = (pathlib.Path(model_id) / "README.md").as_posix()
+            card_path = (pathlib.Path(model_id) / 'README.md').as_posix()
         else:
             card_path = model_id
         return ModelCard.load(card_path, token=hf_token)
@@ -55,11 +57,14 @@ class InferencePipeline:
             return
         base_model_id = self.get_base_model_info(model_id, self.hf_token)
         unet = UNet3DConditionModel.from_pretrained(
-            model_id, subfolder="unet", torch_dtype=torch.float16, use_auth_token=self.hf_token
-        )
-        pipe = TuneAVideoPipeline.from_pretrained(
-            base_model_id, unet=unet, torch_dtype=torch.float16, use_auth_token=self.hf_token
-        )
+            model_id,
+            subfolder='unet',
+            torch_dtype=torch.float16,
+            use_auth_token=self.hf_token)
+        pipe = TuneAVideoPipeline.from_pretrained(base_model_id,
+                                                  unet=unet,
+                                                  torch_dtype=torch.float16,
+                                                  use_auth_token=self.hf_token)
         pipe = pipe.to(self.device)
         if is_xformers_available():
             pipe.unet.enable_xformers_memory_efficient_attention()
@@ -77,7 +82,7 @@ class InferencePipeline:
         guidance_scale: float,
     ) -> PIL.Image.Image:
         if not torch.cuda.is_available():
-            raise gr.Error("CUDA is not available.")
+            raise gr.Error('CUDA is not available.')
 
         self.load_pipe(model_id)
 
@@ -92,10 +97,10 @@ class InferencePipeline:
             generator=generator,
         )  # type: ignore
 
-        frames = rearrange(out.videos[0], "c t h w -> t h w c")
+        frames = rearrange(out.videos[0], 'c t h w -> t h w c')
         frames = (frames * 255).to(torch.uint8).numpy()
 
-        out_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+        out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
         writer = imageio.get_writer(out_file.name, fps=fps)
         for frame in frames:
             writer.append_data(frame)

requirements-monitor.txt

@@ -1,4 +0,0 @@
-nvitop==1.1.1
-pandas==2.0.0
-plotly==5.14.1
-psutil==5.9.4

requirements.txt

@@ -1,19 +1,19 @@
-accelerate==0.18.0
-bitsandbytes==0.37.2
+accelerate==0.15.0
+bitsandbytes==0.35.4
 decord==0.6.0
 diffusers[torch]==0.11.1
 einops==0.6.0
 ftfy==6.1.1
-gradio==3.43.2
-huggingface-hub==0.17.1
-imageio==2.27.0
+gradio==3.16.2
+huggingface-hub==0.12.0
+imageio==2.25.0
 imageio-ffmpeg==0.4.8
 omegaconf==2.3.0
-Pillow==9.5.0
-python-slugify==8.0.1
+Pillow==9.4.0
+python-slugify==7.0.0
 tensorboard==2.11.2
 torch==1.13.1
 torchvision==0.14.1
 transformers==4.26.0
-triton==2.0.0
+triton==2.0.0.dev20221202
 xformers==0.0.16

trainer.py

@@ -8,34 +8,46 @@ import shutil
 import subprocess
 import sys
 
+import gradio as gr
 import slugify
 import torch
 from huggingface_hub import HfApi
 from omegaconf import OmegaConf
 
-from uploader import upload
+from app_upload import ModelUploader
 from utils import save_model_card
 
-sys.path.append("Tune-A-Video")
+sys.path.append('Tune-A-Video')
 
+URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/Tune-A-Video-library/share/YjTcaNJmKyeHFpMBioHhzBcTzCYddVErEk'
+ORIGINAL_SPACE_ID = 'Tune-A-Video-library/Tune-A-Video-Training-UI'
+SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
 
 class Trainer:
-    def __init__(self):
-        self.checkpoint_dir = pathlib.Path("checkpoints")
-        self.checkpoint_dir.mkdir(exist_ok=True)
+    def __init__(self, hf_token: str | None = None):
+        self.hf_token = hf_token
+        self.model_uploader = ModelUploader(hf_token)
 
-        self.log_file = pathlib.Path("log.txt")
-        self.log_file.touch(exist_ok=True)
+        self.checkpoint_dir = pathlib.Path('checkpoints')
+        self.checkpoint_dir.mkdir(exist_ok=True)
 
     def download_base_model(self, base_model_id: str) -> str:
         model_dir = self.checkpoint_dir / base_model_id
         if not model_dir.exists():
-            org_name = base_model_id.split("/")[0]
+            org_name = base_model_id.split('/')[0]
             org_dir = self.checkpoint_dir / org_name
             org_dir.mkdir(exist_ok=True)
-            subprocess.run(shlex.split(f"git clone https://huggingface.co/{base_model_id}"), cwd=org_dir)
+            subprocess.run(shlex.split(
+                f'git clone https://huggingface.co/{base_model_id}'),
+                           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,
@@ -57,32 +69,37 @@ class Trainer:
         use_private_repo: bool,
         delete_existing_repo: bool,
         upload_to: str,
-        pause_space_after_training: bool,
-        hf_token: str,
-    ) -> None:
+        remove_gpu_after_training: bool,
+        input_token: str,
+    ) -> str:
+        if SPACE_ID == ORIGINAL_SPACE_ID:
+            raise gr.Error('This Space does not work on this Shared UI. Duplicate the Space and attribute a GPU')
         if not torch.cuda.is_available():
-            raise RuntimeError("CUDA is not available.")
+            raise gr.Error('CUDA is not available.')
         if training_video is None:
-            raise ValueError("You need to upload a video.")
+            raise gr.Error('You need to upload a video.')
         if not training_prompt:
-            raise ValueError("The training prompt is missing.")
+            raise gr.Error('The training prompt is missing.')
         if not validation_prompt:
-            raise ValueError("The validation prompt is missing.")
+            raise gr.Error('The validation prompt is missing.')
 
         resolution = int(resolution_s)
 
         if not output_model_name:
-            timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-            output_model_name = f"tune-a-video-{timestamp}"
+            timestamp = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
+            output_model_name = f'tune-a-video-{timestamp}'
         output_model_name = slugify.slugify(output_model_name)
 
         repo_dir = pathlib.Path(__file__).parent
-        output_dir = repo_dir / "experiments" / output_model_name
+        output_dir = repo_dir / 'experiments' / output_model_name
         if overwrite_existing_model or upload_to_hub:
             shutil.rmtree(output_dir, ignore_errors=True)
         output_dir.mkdir(parents=True)
 
-        config = OmegaConf.load("Tune-A-Video/configs/man-surfing.yaml")
+        if upload_to_hub:
+            self.join_model_library_org(self.hf_token if self.hf_token else input_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()
         config.train_data.video_path = training_video.name  # type: ignore
@@ -105,40 +122,40 @@ class Trainer:
         config.checkpointing_steps = checkpointing_steps
         config.validation_steps = validation_epochs
         config.seed = seed
-        config.mixed_precision = "fp16" if fp16 else ""
+        config.mixed_precision = 'fp16' if fp16 else ''
         config.use_8bit_adam = use_8bit_adam
 
-        config_path = output_dir / "config.yaml"
-        with open(config_path, "w") as f:
+        config_path = output_dir / 'config.yaml'
+        with open(config_path, 'w') as f:
             OmegaConf.save(config, f)
 
-        command = f"accelerate launch Tune-A-Video/train_tuneavideo.py --config {config_path}"
-        with open(self.log_file, "w") as f:
-            subprocess.run(shlex.split(command), stdout=f, stderr=subprocess.STDOUT, text=True)
-        save_model_card(
-            save_dir=output_dir,
-            base_model=base_model,
-            training_prompt=training_prompt,
-            test_prompt=validation_prompt,
-            test_image_dir="samples",
-        )
+        command = f'accelerate launch Tune-A-Video/train_tuneavideo.py --config {config_path}'
+        subprocess.run(shlex.split(command))
+        save_model_card(save_dir=output_dir,
+                        base_model=base_model,
+                        training_prompt=training_prompt,
+                        test_prompt=validation_prompt,
+                        test_image_dir='samples')
 
-        with open(self.log_file, "a") as f:
-            f.write("Training completed!\n")
+        message = 'Training completed!'
+        print(message)
 
         if upload_to_hub:
-            upload_message = upload(
-                local_folder_path=output_dir.as_posix(),
-                target_repo_name=output_model_name,
+            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,
-                hf_token=hf_token,
-            )
-            with open(self.log_file, "a") as f:
-                f.write(upload_message)
-
-        if pause_space_after_training:
-            if space_id := os.getenv("SPACE_ID"):
-                api = HfApi(token=os.getenv("HF_TOKEN") or hf_token)
-                api.pause_space(repo_id=space_id)
+                input_token=input_token)
+            print(upload_message)
+            message = message + '\n' + upload_message
+
+        if remove_gpu_after_training:
+            space_id = os.getenv('SPACE_ID')
+            if space_id:
+                api = HfApi(token=self.hf_token if self.hf_token else input_token)
+                api.request_space_hardware(repo_id=space_id,
+                                           hardware='cpu-basic')
+
+        return message

uploader.py

@@ -1,66 +1,44 @@
 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
 
-    repo_id = f"{organization}/{target_repo_name}"
-    if delete_existing_repo:
-        try:
-            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
+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_token: str | None = None) -> str:
+        
+        api = HfApi(token=self.hf_token if self.hf_token else input_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:
+                self.api.delete_repo(repo_id, repo_type=repo_type)
+            except Exception:
+                pass
+        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}'
+            message = f'Your model was successfully uploaded to <a href="{url}" target="_blank">{url}</a>.'
+        except Exception as e:
+            message = str(e)
+        return message

utils.py

@@ -5,11 +5,14 @@ import pathlib
 
 def find_exp_dirs() -> list[str]:
     repo_dir = pathlib.Path(__file__).parent
-    exp_root_dir = repo_dir / "experiments"
+    exp_root_dir = repo_dir / 'experiments'
     if not exp_root_dir.exists():
         return []
-    exp_dirs = sorted(exp_root_dir.glob("*"))
-    exp_dirs = [exp_dir for exp_dir in exp_dirs if (exp_dir / "model_index.json").exists()]
+    exp_dirs = sorted(exp_root_dir.glob('*'))
+    exp_dirs = [
+        exp_dir for exp_dir in exp_dirs
+        if (exp_dir / 'model_index.json').exists()
+    ]
     return [path.relative_to(repo_dir).as_posix() for path in exp_dirs]
 
 
@@ -17,21 +20,21 @@ def save_model_card(
     save_dir: pathlib.Path,
     base_model: str,
     training_prompt: str,
-    test_prompt: str = "",
-    test_image_dir: str = "",
+    test_prompt: str = '',
+    test_image_dir: str = '',
 ) -> None:
-    image_str = ""
+    image_str = ''
     if test_prompt and test_image_dir:
-        image_paths = sorted((save_dir / test_image_dir).glob("*.gif"))
+        image_paths = sorted((save_dir / test_image_dir).glob('*.gif'))
         if image_paths:
             image_path = image_paths[-1]
             rel_path = image_path.relative_to(save_dir)
-            image_str = f"""## Samples
+            image_str = f'''## Samples
 Test prompt: {test_prompt}
 
-![{image_path.stem}]({rel_path})"""
+![{image_path.stem}]({rel_path})'''
 
-    model_card = f"""---
+    model_card = f'''---
 license: creativeml-openrail-m
 base_model: {base_model}
 training_prompt: {training_prompt}
@@ -56,7 +59,7 @@ inference: false
 ## Related papers:
 - [Tune-A-Video](https://arxiv.org/abs/2212.11565): One-Shot Tuning of Image Diffusion Models for Text-to-Video Generation
 - [Stable-Diffusion](https://arxiv.org/abs/2112.10752): High-Resolution Image Synthesis with Latent Diffusion Models
-"""
+'''
 
-    with open(save_dir / "README.md", "w") as f:
+    with open(save_dir / 'README.md', 'w') as f:
         f.write(model_card)