Update for inference

app.py

@@ -1,84 +0,0 @@
-#!/usr/bin/env python
-
-from __future__ import annotations
-
-import os
-from subprocess import getoutput
-
-import gradio as gr
-import torch
-
-from app_inference import create_inference_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/) UI'
-
-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>
-'''
-
-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'
-
-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.
-<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.
-<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')
-
-
-def show_warning(warning_text: str) -> gr.Blocks:
-    with gr.Blocks() as demo:
-        with gr.Box():
-            gr.Markdown(warning_text)
-    return demo
-
-
-pipe = InferencePipeline(HF_TOKEN)
-trainer = Trainer(HF_TOKEN)
-
-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 (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)
-        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(HF_TOKEN)
-
-    if not HF_TOKEN:
-        show_warning(HF_TOKEN_NOT_SPECIFIED_WARNING)
-
-demo.queue(max_size=1).launch(share=False)

app_inference.py

@@ -2,19 +2,13 @@
 
 from __future__ import annotations
 
-import enum
+import os
 
 import gradio as gr
 from huggingface_hub import HfApi
 
-from constants import MODEL_LIBRARY_ORG_NAME, UploadTarget
+from constants import MODEL_LIBRARY_ORG_NAME
 from inference import InferencePipeline
-from utils import find_exp_dirs
-
-
-class ModelSource(enum.Enum):
-    HUB_LIB = UploadTarget.MODEL_LIBRARY.value
-    LOCAL = 'Local'
 
 
 class InferenceUtil:
@@ -30,20 +24,6 @@ class InferenceUtil:
         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)
-
-    def reload_model_list(self, model_source: str) -> dict:
-        if model_source == ModelSource.HUB_LIB.value:
-            return self.load_hub_model_list()
-        elif model_source == ModelSource.LOCAL.value:
-            return self.load_local_model_list()
-        else:
-            raise ValueError
-
     def load_model_info(self, model_id: str) -> tuple[str, str]:
         try:
             card = InferencePipeline.get_model_card(model_id, self.hf_token)
@@ -53,118 +33,99 @@ class InferenceUtil:
         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]:
-        model_list_update = self.reload_model_list(model_source)
+    def reload_model_list_and_update_model_info(self) -> tuple[dict, str, str]:
+        model_list_update = self.load_hub_model_list()
         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) -> gr.Blocks:
-    app = InferenceUtil(hf_token)
-
-    with gr.Blocks() as demo:
-        with gr.Row():
-            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):
-                        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('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,
+TITLE = '# [Tune-A-Video](https://tuneavideo.github.io/)'
+HF_TOKEN = os.getenv('HF_TOKEN')
+pipe = InferencePipeline(HF_TOKEN)
+app = InferenceUtil(HF_TOKEN)
+
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(TITLE)
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Box():
+                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('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('''
+            - 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')
+
+    reload_button.click(fn=app.reload_model_list_and_update_model_info,
+                        inputs=None,
                         outputs=[
+                            model_id,
                             base_model_used_for_training,
                             prompt_used_for_training,
                         ])
-        inputs = [
-            model_id,
-            prompt,
-            video_length,
-            fps,
-            seed,
-            num_steps,
-            guidance_scale,
-        ]
-        prompt.submit(fn=pipe.run, inputs=inputs, outputs=result)
-        run_button.click(fn=pipe.run, inputs=inputs, outputs=result)
-    return demo
-
-
-if __name__ == '__main__':
-    import os
-
-    hf_token = os.getenv('HF_TOKEN')
-    pipe = InferencePipeline(hf_token)
-    demo = create_inference_demo(pipe, hf_token)
-    demo.queue(max_size=10).launch(share=False)
+    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,
+        video_length,
+        fps,
+        seed,
+        num_steps,
+        guidance_scale,
+    ]
+    prompt.submit(fn=pipe.run, inputs=inputs, outputs=result)
+    run_button.click(fn=pipe.run, inputs=inputs, outputs=result)
+
+demo.queue().launch()

app_training.py

@@ -1,135 +0,0 @@
-#!/usr/bin/env python
-
-from __future__ import annotations
-
-import os
-
-import gradio as gr
-
-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) -> 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('''
-                        - Upload a video and write a `Training Prompt` that describes the video.
-                        ''')
-
-            with gr.Column():
-                with gr.Box():
-                    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)
-
-                    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('''
-                        - 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')
-            with gr.Column():
-                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 = gr.Radio(
-                        label='Upload to',
-                        choices=[_.value for _ in UploadTarget],
-                        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')
-
-        with gr.Box():
-            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,
-                remove_gpu_after_training, input_token
-            ],
-            outputs=output_message)
-    return demo
-
-
-if __name__ == '__main__':
-    hf_token = os.getenv('HF_TOKEN')
-    trainer = Trainer(hf_token)
-    demo = create_training_demo(trainer)
-    demo.queue(max_size=1).launch(share=False)

app_upload.py

@@ -1,106 +0,0 @@
-#!/usr/bin/env python
-
-from __future__ import annotations
-
-import pathlib
-
-import gradio as gr
-import slugify
-
-from constants import MODEL_LIBRARY_ORG_NAME, UploadTarget
-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(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')
-            model_dir = gr.Dropdown(
-                label='Model names',
-                choices=model_dirs,
-                value=model_dirs[0] if model_dirs else None)
-        with gr.Box():
-            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')
-            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')
-            output_message = gr.Markdown()
-
-        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__':
-    import os
-
-    hf_token = os.getenv('HF_TOKEN')
-    demo = create_upload_demo(hf_token)
-    demo.queue(max_size=1).launch(share=False)

trainer.py

@@ -1,166 +0,0 @@
-from __future__ import annotations
-
-import datetime
-import os
-import pathlib
-import shlex
-import shutil
-import subprocess
-import sys
-
-import gradio as gr
-import slugify
-import torch
-from huggingface_hub import HfApi
-from omegaconf import OmegaConf
-
-from app_upload import ModelUploader
-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'
-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, hf_token: str | None = None):
-        self.hf_token = hf_token
-        self.model_uploader = ModelUploader(hf_token)
-
-        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_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)
-        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,
-        training_prompt: str,
-        output_model_name: str,
-        overwrite_existing_model: bool,
-        validation_prompt: str,
-        base_model: str,
-        resolution_s: str,
-        n_steps: int,
-        learning_rate: float,
-        gradient_accumulation: int,
-        seed: int,
-        fp16: bool,
-        use_8bit_adam: bool,
-        checkpointing_steps: int,
-        validation_epochs: int,
-        upload_to_hub: bool,
-        use_private_repo: bool,
-        delete_existing_repo: bool,
-        upload_to: str,
-        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 gr.Error('CUDA is not available.')
-        if training_video is None:
-            raise gr.Error('You need to upload a video.')
-        if not training_prompt:
-            raise gr.Error('The training prompt is missing.')
-        if not validation_prompt:
-            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}'
-        output_model_name = slugify.slugify(output_model_name)
-
-        repo_dir = pathlib.Path(__file__).parent
-        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)
-
-        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
-        config.train_data.prompt = training_prompt
-        config.train_data.n_sample_frames = 8
-        config.train_data.width = resolution
-        config.train_data.height = resolution
-        config.train_data.sample_start_idx = 0
-        config.train_data.sample_frame_rate = 1
-        config.validation_data.prompts = [validation_prompt]
-        config.validation_data.video_length = 8
-        config.validation_data.width = resolution
-        config.validation_data.height = resolution
-        config.validation_data.num_inference_steps = 50
-        config.validation_data.guidance_scale = 7.5
-        config.learning_rate = learning_rate
-        config.gradient_accumulation_steps = gradient_accumulation
-        config.train_batch_size = 1
-        config.max_train_steps = n_steps
-        config.checkpointing_steps = checkpointing_steps
-        config.validation_steps = validation_epochs
-        config.seed = seed
-        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:
-            OmegaConf.save(config, f)
-
-        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')
-
-        message = 'Training completed!'
-        print(message)
-
-        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_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,44 +0,0 @@
-from __future__ import annotations
-
-from huggingface_hub import HfApi
-
-
-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:
-                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

@@ -1,65 +0,0 @@
-from __future__ import annotations
-
-import pathlib
-
-
-def find_exp_dirs() -> list[str]:
-    repo_dir = pathlib.Path(__file__).parent
-    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()
-    ]
-    return [path.relative_to(repo_dir).as_posix() for path in exp_dirs]
-
-
-def save_model_card(
-    save_dir: pathlib.Path,
-    base_model: str,
-    training_prompt: str,
-    test_prompt: str = '',
-    test_image_dir: str = '',
-) -> None:
-    image_str = ''
-    if test_prompt and test_image_dir:
-        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
-Test prompt: {test_prompt}
-
-![{image_path.stem}]({rel_path})'''
-
-    model_card = f'''---
-license: creativeml-openrail-m
-base_model: {base_model}
-training_prompt: {training_prompt}
-tags:
-- stable-diffusion
-- stable-diffusion-diffusers
-- text-to-image
-- diffusers
-- text-to-video
-- tune-a-video
-inference: false
----
-
-# Tune-A-Video - {save_dir.name}
-
-## Model description
-- Base model: [{base_model}](https://huggingface.co/{base_model})
-- Training prompt: {training_prompt}
-
-{image_str}
-
-## 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:
-        f.write(model_card)