Add files

.gitattributes

@@ -1,3 +1,4 @@
+*.whl filter=lfs diff=lfs merge=lfs -text
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,6 +1,5 @@
-training_data/
+checkpoints/
 experiments/
-wandb/
 
 
 # Byte-compiled / optimized / DLL files

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "Tune-A-Video"]
+	path = Tune-A-Video
+	url = https://github.com/showlab/Tune-A-Video

.pre-commit-config.yaml

@@ -1,4 +1,4 @@
-exclude: train_dreambooth_lora.py
+exclude: patch
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v4.2.0
@@ -21,7 +21,7 @@ repos:
   - id: docformatter
     args: ['--in-place']
 - repo: https://github.com/pycqa/isort
-  rev: 5.10.1
+  rev: 5.12.0
   hooks:
     - id: isort
 - repo: https://github.com/pre-commit/mirrors-mypy

Dockerfile

@@ -0,0 +1,59 @@
+FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get upgrade -y && \
+    apt-get install -y --no-install-recommends \
+    git \
+    git-lfs \
+    wget \
+    curl \
+    # ffmpeg \
+    ffmpeg \
+    x264 \
+    # python build dependencies \
+    build-essential \
+    libssl-dev \
+    zlib1g-dev \
+    libbz2-dev \
+    libreadline-dev \
+    libsqlite3-dev \
+    libncursesw5-dev \
+    xz-utils \
+    tk-dev \
+    libxml2-dev \
+    libxmlsec1-dev \
+    libffi-dev \
+    liblzma-dev && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:${PATH}
+WORKDIR ${HOME}/app
+
+RUN curl https://pyenv.run | bash
+ENV PATH=${HOME}/.pyenv/shims:${HOME}/.pyenv/bin:${PATH}
+ENV PYTHON_VERSION=3.10.9
+RUN pyenv install ${PYTHON_VERSION} && \
+    pyenv global ${PYTHON_VERSION} && \
+    pyenv rehash && \
+    pip install --no-cache-dir -U pip setuptools wheel
+
+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 wheel/xformers-0.0.16+bc08bbc.d20230130-cp310-cp310-linux_x86_64.whl /tmp/xformers-0.0.16+bc08bbc.d20230130-cp310-cp310-linux_x86_64.whl
+RUN pip install --no-cache-dir -U /tmp/xformers-0.0.16+bc08bbc.d20230130-cp310-cp310-linux_x86_64.whl
+
+COPY --chown=1000 . ${HOME}/app
+RUN cd Tune-A-Video && patch -p1 < ../patch
+ENV PYTHONPATH=${HOME}/app \
+    PYTHONUNBUFFERED=1 \
+    GRADIO_ALLOW_FLAGGING=never \
+    GRADIO_NUM_PORTS=1 \
+    GRADIO_SERVER_NAME=0.0.0.0 \
+    GRADIO_THEME=huggingface \
+    SYSTEM=spaces
+CMD ["python", "app.py"]

README.md

@@ -1,12 +1,9 @@
 ---
-title: LoRA DreamBooth Training UI
+title: Tune-A-Video Training UI
 emoji: ⚡
 colorFrom: red
 colorTo: purple
-sdk: gradio
-sdk_version: 3.16.2
-python_version: 3.10.9
-app_file: app.py
+sdk: docker
 pinned: false
 license: mit
 duplicated_from: lora-library/LoRA-DreamBooth-Training-UI

Tune-A-Video

@@ -0,0 +1 @@
+Subproject commit b2c8c3eeac0df5c5d9eccc4dd2153e17b83c638c

app.py

@@ -13,11 +13,11 @@ from app_upload import create_upload_demo
 from inference import InferencePipeline
 from trainer import Trainer
 
-TITLE = '# LoRA DreamBooth Training UI'
+TITLE = '# Tune-A-Video Training UI'
 
-ORIGINAL_SPACE_ID = 'lora-library/LoRA-DreamBooth-Training-UI'
+ORIGINAL_SPACE_ID = 'hysts/Tune-A-Video-Training-UI'
 SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
-SHARED_UI_WARNING = f'''# Attention - This Space doesn't work in this shared UI. You can duplicate and use it with a paid private T4 GPU.
+SHARED_UI_WARNING = f'''# Attention - This Space doesn't work in this shared UI. You can duplicate and use it with a paid private A100 GPU.
 
 <center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img 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>
 '''
@@ -29,7 +29,7 @@ else:
 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.
-"T4 small" is sufficient to run this demo.
+"A100 large" is required to run this demo.
 </center>
 '''
 

app_inference.py

@@ -7,18 +7,13 @@ import enum
 import gradio as gr
 from huggingface_hub import HfApi
 
+from constants import MODEL_LIBRARY_ORG_NAME, UploadTarget
 from inference import InferencePipeline
 from utils import find_exp_dirs
 
-SAMPLE_MODEL_IDS = [
-    'patrickvonplaten/lora_dreambooth_dog_example',
-    'sayakpaul/sd-model-finetuned-lora-t4',
-]
-
 
 class ModelSource(enum.Enum):
-    SAMPLE = 'Sample'
-    HUB_LIB = 'Hub (lora-library)'
+    HUB_LIB = UploadTarget.MODEL_LIBRARY.value
     LOCAL = 'Local'
 
 
@@ -26,47 +21,41 @@ class InferenceUtil:
     def __init__(self, hf_token: str | None):
         self.hf_token = hf_token
 
-    @staticmethod
-    def load_sample_lora_model_list():
-        return gr.update(choices=SAMPLE_MODEL_IDS, value=SAMPLE_MODEL_IDS[0])
-
-    def load_hub_lora_model_list(self) -> dict:
+    def load_hub_model_list(self) -> dict:
         api = HfApi(token=self.hf_token)
         choices = [
-            info.modelId for info in api.list_models(author='lora-library')
+            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_lora_model_list() -> dict:
+    def load_local_model_list() -> dict:
         choices = find_exp_dirs()
         return gr.update(choices=choices,
                          value=choices[0] if choices else None)
 
-    def reload_lora_model_list(self, model_source: str) -> dict:
-        if model_source == ModelSource.SAMPLE.value:
-            return self.load_sample_lora_model_list()
-        elif model_source == ModelSource.HUB_LIB.value:
-            return self.load_hub_lora_model_list()
+    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_lora_model_list()
+            return self.load_local_model_list()
         else:
             raise ValueError
 
-    def load_model_info(self, lora_model_id: str) -> tuple[str, str]:
+    def load_model_info(self, model_id: str) -> tuple[str, str]:
         try:
-            card = InferencePipeline.get_model_card(lora_model_id,
-                                                    self.hf_token)
+            card = InferencePipeline.get_model_card(model_id, self.hf_token)
         except Exception:
             return '', ''
         base_model = getattr(card.data, 'base_model', '')
-        instance_prompt = getattr(card.data, 'instance_prompt', '')
-        return base_model, instance_prompt
+        training_prompt = getattr(card.data, 'training_prompt', '')
+        return base_model, training_prompt
 
-    def reload_lora_model_list_and_update_model_info(
+    def reload_model_list_and_update_model_info(
             self, model_source: str) -> tuple[dict, str, str]:
-        model_list_update = self.reload_lora_model_list(model_source)
+        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 '')
         return model_list_update, *model_info
@@ -83,30 +72,34 @@ def create_inference_demo(pipe: InferencePipeline,
                     model_source = gr.Radio(
                         label='Model Source',
                         choices=[_.value for _ in ModelSource],
-                        value=ModelSource.SAMPLE.value)
+                        value=ModelSource.HUB_LIB.value)
                     reload_button = gr.Button('Reload Model List')
-                    lora_model_id = gr.Dropdown(label='LoRA Model ID',
-                                                choices=SAMPLE_MODEL_IDS,
-                                                value=SAMPLE_MODEL_IDS[0])
+                    model_id = gr.Dropdown(label='Model ID',
+                                           choices=None,
+                                           value=None)
                     with gr.Accordion(
                             label=
-                            'Model info (Base model and instance prompt used for training)',
+                            '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)
-                            instance_prompt_used_for_training = gr.Text(
-                                label='Instance prompt', interactive=False)
+                            prompt_used_for_training = gr.Text(
+                                label='Training prompt', interactive=False)
                 prompt = gr.Textbox(
                     label='Prompt',
                     max_lines=1,
-                    placeholder='Example: "A picture of a sks dog in a bucket"'
-                )
-                alpha = gr.Slider(label='LoRA alpha',
-                                  minimum=0,
-                                  maximum=2,
-                                  step=0.05,
-                                  value=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,
@@ -117,7 +110,7 @@ def create_inference_demo(pipe: InferencePipeline,
                                           minimum=0,
                                           maximum=100,
                                           step=1,
-                                          value=25)
+                                          value=50)
                     guidance_scale = gr.Slider(label='CFG Scale',
                                                minimum=0,
                                                maximum=50,
@@ -130,34 +123,33 @@ def create_inference_demo(pipe: InferencePipeline,
                 - After training, you can press "Reload Model List" button to load your trained model names.
                 ''')
             with gr.Column():
-                result = gr.Image(label='Result')
-
-        model_source.change(
-            fn=app.reload_lora_model_list_and_update_model_info,
-            inputs=model_source,
-            outputs=[
-                lora_model_id,
-                base_model_used_for_training,
-                instance_prompt_used_for_training,
-            ])
-        reload_button.click(
-            fn=app.reload_lora_model_list_and_update_model_info,
-            inputs=model_source,
-            outputs=[
-                lora_model_id,
-                base_model_used_for_training,
-                instance_prompt_used_for_training,
-            ])
-        lora_model_id.change(fn=app.load_model_info,
-                             inputs=lora_model_id,
-                             outputs=[
-                                 base_model_used_for_training,
-                                 instance_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 = [
-            lora_model_id,
+            model_id,
             prompt,
-            alpha,
+            video_length,
+            fps,
             seed,
             num_steps,
             guidance_scale,

app_training.py

@@ -6,7 +6,7 @@ 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
 
@@ -18,12 +18,13 @@ def create_training_demo(trainer: Trainer,
             with gr.Column():
                 with gr.Box():
                     gr.Markdown('Training Data')
-                    instance_images = gr.Files(label='Instance images')
-                    instance_prompt = gr.Textbox(label='Instance prompt',
-                                                 max_lines=1)
+                    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 images of the style you are planning on training on.
-                        - For an instance prompt, use a unique, made up word to avoid collisions.
+                        - Upload a video and write a prompt describing the video.
                         ''')
                 with gr.Box():
                     gr.Markdown('Output Model')
@@ -46,25 +47,26 @@ def create_training_demo(trainer: Trainer,
                     upload_to = gr.Radio(
                         label='Upload to',
                         choices=[_.value for _ in UploadTarget],
-                        value=UploadTarget.LORA_LIBRARY.value)
-                    gr.Markdown('''
-                    - By default, trained models will be uploaded to [LoRA Library](https://huggingface.co/lora-library) (see [this example model](https://huggingface.co/lora-library/lora-dreambooth-sample-dog)).
-                    - You can also choose "Personal Profile", in which case, the model will be uploaded to https://huggingface.co/{your_username}/{model_name}.
+                        value=UploadTarget.MODEL_LIBRARY.value)
+                    gr.Markdown(f'''
+                    - By default, trained models will be uploaded to [Tune-A-Video Library](https://huggingface.co/{MODEL_LIBRARY_ORG_NAME}) (see [this example model](https://huggingface.co/{MODEL_LIBRARY_ORG_NAME}/{SAMPLE_MODEL_REPO})).
+                    - You can also choose "Personal Profile", in which case, the model will be uploaded to https://huggingface.co/{{your_username}}/{{model_name}}.
                     ''')
 
             with gr.Box():
                 gr.Markdown('Training Parameters')
                 with gr.Row():
-                    base_model = gr.Text(
-                        label='Base Model',
-                        value='stabilityai/stable-diffusion-2-1-base',
-                        max_lines=1)
+                    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')
+                                             label='Resolution',
+                                             visible=False)
                 num_training_steps = gr.Number(
-                    label='Number of Training Steps', value=1000, precision=0)
-                learning_rate = gr.Number(label='Learning Rate', value=0.0001)
+                    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,
@@ -75,25 +77,20 @@ def create_training_demo(trainer: Trainer,
                                  step=1,
                                  value=0)
                 fp16 = gr.Checkbox(label='FP16', value=True)
-                use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=True)
+                use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=False)
                 checkpointing_steps = gr.Number(label='Checkpointing Steps',
-                                                value=100,
+                                                value=1000,
                                                 precision=0)
-                use_wandb = gr.Checkbox(label='Use W&B',
-                                        value=False,
-                                        interactive=bool(
-                                            os.getenv('WANDB_API_KEY')))
                 validation_epochs = gr.Number(label='Validation Epochs',
                                               value=100,
                                               precision=0)
                 gr.Markdown('''
                     - The base model must be a model that is compatible with [diffusers](https://github.com/huggingface/diffusers) library.
                     - It takes a few minutes to download the base model first.
-                    - It will take about 8 minutes to train for 1000 steps with a T4 GPU.
+                    - It will take about 4 minutes to train for 300 steps with an A100 GPU.
+                    - It takes a few minutes to upload your trained model.
                     - You may want to try a small number of steps first, like 1, to see if everything works fine in your environment.
                     - You can check the training status by pressing the "Open logs" button if you are running this on your Space.
-                    - You need to set the environment variable `WANDB_API_KEY` if you'd like to use [W&B](https://wandb.ai/site). See [W&B documentation](https://docs.wandb.ai/guides/track/advanced/environment-variables).
-                    - **Note:** Due to [this issue](https://github.com/huggingface/accelerate/issues/944), currently, training will not terminate properly if you use W&B.
                     ''')
 
         remove_gpu_after_training = gr.Checkbox(
@@ -111,8 +108,8 @@ def create_training_demo(trainer: Trainer,
             run_button.click(fn=pipe.clear)
         run_button.click(fn=trainer.run,
                          inputs=[
-                             instance_images,
-                             instance_prompt,
+                             training_video,
+                             training_prompt,
                              output_model_name,
                              delete_existing_model,
                              validation_prompt,
@@ -125,7 +122,6 @@ def create_training_demo(trainer: Trainer,
                              fp16,
                              use_8bit_adam,
                              checkpointing_steps,
-                             use_wandb,
                              validation_epochs,
                              upload_to_hub,
                              use_private_repo,

app_upload.py

@@ -7,13 +7,13 @@ import pathlib
 import gradio as gr
 import slugify
 
-from constants import UploadTarget
+from constants import MODEL_LIBRARY_ORG_NAME, UploadTarget
 from uploader import Uploader
 from utils import find_exp_dirs
 
 
-class LoRAModelUploader(Uploader):
-    def upload_lora_model(
+class ModelUploader(Uploader):
+    def upload_model(
         self,
         folder_path: str,
         repo_name: str,
@@ -29,8 +29,8 @@ class LoRAModelUploader(Uploader):
 
         if upload_to == UploadTarget.PERSONAL_PROFILE.value:
             organization = ''
-        elif upload_to == UploadTarget.LORA_LIBRARY.value:
-            organization = 'lora-library'
+        elif upload_to == UploadTarget.MODEL_LIBRARY.value:
+            organization = MODEL_LIBRARY_ORG_NAME
         else:
             raise ValueError
 
@@ -41,14 +41,14 @@ class LoRAModelUploader(Uploader):
                            delete_existing_repo=delete_existing_repo)
 
 
-def load_local_lora_model_list() -> dict:
-    choices = find_exp_dirs(ignore_repo=True)
+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 = LoRAModelUploader(hf_token)
-    model_dirs = find_exp_dirs(ignore_repo=True)
+    uploader = ModelUploader(hf_token)
+    model_dirs = find_exp_dirs()
 
     with gr.Blocks() as demo:
         with gr.Box():
@@ -66,20 +66,20 @@ def create_upload_demo(hf_token: str | None) -> gr.Blocks:
                     label='Delete existing repo of the same name', value=False)
             upload_to = gr.Radio(label='Upload to',
                                  choices=[_.value for _ in UploadTarget],
-                                 value=UploadTarget.LORA_LIBRARY.value)
+                                 value=UploadTarget.MODEL_LIBRARY.value)
             model_name = gr.Textbox(label='Model Name')
         upload_button = gr.Button('Upload')
-        gr.Markdown('''
-            - You can upload your trained model to your personal profile (i.e. https://huggingface.co/{your_username}/{model_name}) or to the public [LoRA Concepts Library](https://huggingface.co/lora-library) (i.e. https://huggingface.co/lora-library/{model_name}).
+        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_lora_model_list,
+        reload_button.click(fn=load_local_model_list,
                             inputs=None,
                             outputs=model_dir)
-        upload_button.click(fn=uploader.upload_lora_model,
+        upload_button.click(fn=uploader.upload_model,
                             inputs=[
                                 model_dir,
                                 model_name,

constants.py

@@ -3,4 +3,8 @@ import enum
 
 class UploadTarget(enum.Enum):
     PERSONAL_PROFILE = 'Personal Profile'
-    LORA_LIBRARY = 'LoRA Library'
+    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'

inference.py

@@ -2,13 +2,21 @@ from __future__ import annotations
 
 import gc
 import pathlib
+import sys
+import tempfile
 
 import gradio as gr
+import imageio
 import PIL.Image
 import torch
-from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+from einops import rearrange
 from huggingface_hub import ModelCard
 
+sys.path.append('Tune-A-Video')
+
+from tuneavideo.models.unet import UNet3DConditionModel
+from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
+
 
 class InferencePipeline:
     def __init__(self, hf_token: str | None = None):
@@ -16,20 +24,18 @@ class InferencePipeline:
         self.pipe = None
         self.device = torch.device(
             'cuda:0' if torch.cuda.is_available() else 'cpu')
-        self.lora_model_id = None
-        self.base_model_id = None
+        self.model_id = None
 
     def clear(self) -> None:
-        self.lora_model_id = None
-        self.base_model_id = None
+        self.model_id = None
         del self.pipe
         self.pipe = None
         torch.cuda.empty_cache()
         gc.collect()
 
     @staticmethod
-    def check_if_model_is_local(lora_model_id: str) -> bool:
-        return pathlib.Path(lora_model_id).exists()
+    def check_if_model_is_local(model_id: str) -> bool:
+        return pathlib.Path(model_id).exists()
 
     @staticmethod
     def get_model_card(model_id: str,
@@ -41,39 +47,30 @@ class InferencePipeline:
         return ModelCard.load(card_path, token=hf_token)
 
     @staticmethod
-    def get_base_model_info(lora_model_id: str,
-                            hf_token: str | None = None) -> str:
-        card = InferencePipeline.get_model_card(lora_model_id, hf_token)
+    def get_base_model_info(model_id: str, hf_token: str | None = None) -> str:
+        card = InferencePipeline.get_model_card(model_id, hf_token)
         return card.data.base_model
 
-    def load_pipe(self, lora_model_id: str) -> None:
-        if lora_model_id == self.lora_model_id:
+    def load_pipe(self, model_id: str) -> None:
+        if model_id == self.model_id:
             return
-        base_model_id = self.get_base_model_info(lora_model_id, self.hf_token)
-        if base_model_id != self.base_model_id:
-            if self.device.type == 'cpu':
-                pipe = DiffusionPipeline.from_pretrained(
-                    base_model_id, use_auth_token=self.hf_token)
-            else:
-                pipe = DiffusionPipeline.from_pretrained(
-                    base_model_id,
-                    torch_dtype=torch.float16,
-                    use_auth_token=self.hf_token)
-                pipe = pipe.to(self.device)
-            pipe.scheduler = DPMSolverMultistepScheduler.from_config(
-                pipe.scheduler.config)
-            self.pipe = pipe
-        self.pipe.unet.load_attn_procs(  # type: ignore
-            lora_model_id, use_auth_token=self.hf_token)
-
-        self.lora_model_id = lora_model_id  # type: ignore
-        self.base_model_id = base_model_id  # type: ignore
+        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)
+        pipe = TuneAVideoPipeline.from_pretrained(base_model_id,
+                                                  unet=unet,
+                                                  torch_dtype=torch.float16)
+        pipe = pipe.to(self.device)
+        self.pipe = pipe
+        self.model_id = model_id  # type: ignore
 
     def run(
         self,
-        lora_model_id: str,
+        model_id: str,
         prompt: str,
-        lora_scale: float,
+        video_length: int,
+        fps: int,
         seed: int,
         n_steps: int,
         guidance_scale: float,
@@ -81,14 +78,26 @@ class InferencePipeline:
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
 
-        self.load_pipe(lora_model_id)
+        self.load_pipe(model_id)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         out = self.pipe(
             prompt,
+            video_length=video_length,
+            width=512,
+            height=512,
             num_inference_steps=n_steps,
             guidance_scale=guidance_scale,
             generator=generator,
-            cross_attention_kwargs={'scale': lora_scale},
         )  # type: ignore
-        return out.images[0]
+
+        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)
+        writer = imageio.get_writer(out_file.name, fps=fps)
+        for frame in frames:
+            writer.append_data(frame)
+        writer.close()
+
+        return out_file.name

packages.txt

@@ -0,0 +1 @@
+ffmpeg

patch

@@ -0,0 +1,15 @@
+diff --git a/train_tuneavideo.py b/train_tuneavideo.py
+index 66d51b2..86b2a5d 100644
+--- a/train_tuneavideo.py
++++ b/train_tuneavideo.py
+@@ -94,8 +94,8 @@ def main(
+ 
+     # Handle the output folder creation
+     if accelerator.is_main_process:
+-        now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+-        output_dir = os.path.join(output_dir, now)
++        #now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
++        #output_dir = os.path.join(output_dir, now)
+         os.makedirs(output_dir, exist_ok=True)
+         OmegaConf.save(config, os.path.join(output_dir, 'config.yaml'))
+ 

requirements.txt

@@ -1,14 +1,18 @@
 accelerate==0.15.0
-bitsandbytes==0.36.0.post2
-datasets==2.8.0
-git+https://github.com/huggingface/diffusers@31be42209ddfdb69d9640a777b32e9b5c6259bf0#egg=diffusers
+bitsandbytes==0.35.4
+decord==0.6.0
+diffusers[torch]==0.11.1
+einops==0.6.0
 ftfy==6.1.1
 gradio==3.16.2
 huggingface-hub==0.12.0
+imageio==2.25.0
+imageio-ffmpeg==0.4.8
+omegaconf==2.3.0
 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
-wandb==0.13.9
+triton==2.0.0.dev20221202

train_dreambooth_lora.py

@@ -1,1026 +0,0 @@
-#!/usr/bin/env python
-# coding=utf-8
-#
-# This file is adapted from https://github.com/huggingface/diffusers/blob/febaf863026bd014b7a14349336544fc109d0f57/examples/dreambooth/train_dreambooth_lora.py
-# The original license is as below:
-#
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-
-import argparse
-import hashlib
-import logging
-import math
-import os
-import warnings
-from pathlib import Path
-from typing import Optional
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch.utils.data import Dataset
-
-import datasets
-import diffusers
-import transformers
-from accelerate import Accelerator
-from accelerate.logging import get_logger
-from accelerate.utils import set_seed
-from diffusers import (
-    AutoencoderKL,
-    DDPMScheduler,
-    DiffusionPipeline,
-    DPMSolverMultistepScheduler,
-    UNet2DConditionModel,
-)
-from diffusers.loaders import AttnProcsLayers
-from diffusers.models.cross_attention import LoRACrossAttnProcessor
-from diffusers.optimization import get_scheduler
-from diffusers.utils import check_min_version, is_wandb_available
-from diffusers.utils.import_utils import is_xformers_available
-from huggingface_hub import HfFolder, Repository, create_repo, whoami
-from PIL import Image
-from torchvision import transforms
-from tqdm.auto import tqdm
-from transformers import AutoTokenizer, PretrainedConfig
-
-
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.12.0.dev0")
-
-logger = get_logger(__name__)
-
-
-def save_model_card(repo_name, images=None, base_model=str, prompt=str, repo_folder=None):
-    img_str = ""
-    for i, image in enumerate(images):
-        image.save(os.path.join(repo_folder, f"image_{i}.png"))
-        img_str += f"![img_{i}](./image_{i}.png)\n"
-
-    yaml = f"""
----
-license: creativeml-openrail-m
-base_model: {base_model}
-tags:
-- stable-diffusion
-- stable-diffusion-diffusers
-- text-to-image
-- diffusers
-- lora
-inference: true
----
-    """
-    model_card = f"""
-# LoRA DreamBooth - {repo_name}
-
-These are LoRA adaption weights for {repo_name}. The weights were trained on {prompt} using [DreamBooth](https://dreambooth.github.io/). You can find some example images in the following. \n
-{img_str}
-"""
-    with open(os.path.join(repo_folder, "README.md"), "w") as f:
-        f.write(yaml + model_card)
-
-
-def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str, revision: str):
-    text_encoder_config = PretrainedConfig.from_pretrained(
-        pretrained_model_name_or_path,
-        subfolder="text_encoder",
-        revision=revision,
-    )
-    model_class = text_encoder_config.architectures[0]
-
-    if model_class == "CLIPTextModel":
-        from transformers import CLIPTextModel
-
-        return CLIPTextModel
-    elif model_class == "RobertaSeriesModelWithTransformation":
-        from diffusers.pipelines.alt_diffusion.modeling_roberta_series import RobertaSeriesModelWithTransformation
-
-        return RobertaSeriesModelWithTransformation
-    else:
-        raise ValueError(f"{model_class} is not supported.")
-
-
-def parse_args(input_args=None):
-    parser = argparse.ArgumentParser(description="Simple example of a training script.")
-    parser.add_argument(
-        "--pretrained_model_name_or_path",
-        type=str,
-        default=None,
-        required=True,
-        help="Path to pretrained model or model identifier from huggingface.co/models.",
-    )
-    parser.add_argument(
-        "--revision",
-        type=str,
-        default=None,
-        required=False,
-        help="Revision of pretrained model identifier from huggingface.co/models.",
-    )
-    parser.add_argument(
-        "--tokenizer_name",
-        type=str,
-        default=None,
-        help="Pretrained tokenizer name or path if not the same as model_name",
-    )
-    parser.add_argument(
-        "--instance_data_dir",
-        type=str,
-        default=None,
-        required=True,
-        help="A folder containing the training data of instance images.",
-    )
-    parser.add_argument(
-        "--class_data_dir",
-        type=str,
-        default=None,
-        required=False,
-        help="A folder containing the training data of class images.",
-    )
-    parser.add_argument(
-        "--instance_prompt",
-        type=str,
-        default=None,
-        required=True,
-        help="The prompt with identifier specifying the instance",
-    )
-    parser.add_argument(
-        "--class_prompt",
-        type=str,
-        default=None,
-        help="The prompt to specify images in the same class as provided instance images.",
-    )
-    parser.add_argument(
-        "--validation_prompt",
-        type=str,
-        default=None,
-        help="A prompt that is used during validation to verify that the model is learning.",
-    )
-    parser.add_argument(
-        "--num_validation_images",
-        type=int,
-        default=4,
-        help="Number of images that should be generated during validation with `validation_prompt`.",
-    )
-    parser.add_argument(
-        "--validation_epochs",
-        type=int,
-        default=50,
-        help=(
-            "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt"
-            " `args.validation_prompt` multiple times: `args.num_validation_images`."
-        ),
-    )
-    parser.add_argument(
-        "--with_prior_preservation",
-        default=False,
-        action="store_true",
-        help="Flag to add prior preservation loss.",
-    )
-    parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.")
-    parser.add_argument(
-        "--num_class_images",
-        type=int,
-        default=100,
-        help=(
-            "Minimal class images for prior preservation loss. If there are not enough images already present in"
-            " class_data_dir, additional images will be sampled with class_prompt."
-        ),
-    )
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default="lora-dreambooth-model",
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
-    parser.add_argument(
-        "--resolution",
-        type=int,
-        default=512,
-        help=(
-            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
-            " resolution"
-        ),
-    )
-    parser.add_argument(
-        "--center_crop",
-        default=False,
-        action="store_true",
-        help=(
-            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
-            " cropped. The images will be resized to the resolution first before cropping."
-        ),
-    )
-    parser.add_argument(
-        "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
-    )
-    parser.add_argument(
-        "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images."
-    )
-    parser.add_argument("--num_train_epochs", type=int, default=1)
-    parser.add_argument(
-        "--max_train_steps",
-        type=int,
-        default=None,
-        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
-    )
-    parser.add_argument(
-        "--checkpointing_steps",
-        type=int,
-        default=500,
-        help=(
-            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
-            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
-            " training using `--resume_from_checkpoint`."
-        ),
-    )
-    parser.add_argument(
-        "--resume_from_checkpoint",
-        type=str,
-        default=None,
-        help=(
-            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
-            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
-        ),
-    )
-    parser.add_argument(
-        "--gradient_accumulation_steps",
-        type=int,
-        default=1,
-        help="Number of updates steps to accumulate before performing a backward/update pass.",
-    )
-    parser.add_argument(
-        "--gradient_checkpointing",
-        action="store_true",
-        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=5e-4,
-        help="Initial learning rate (after the potential warmup period) to use.",
-    )
-    parser.add_argument(
-        "--scale_lr",
-        action="store_true",
-        default=False,
-        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
-    )
-    parser.add_argument(
-        "--lr_scheduler",
-        type=str,
-        default="constant",
-        help=(
-            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-            ' "constant", "constant_with_warmup"]'
-        ),
-    )
-    parser.add_argument(
-        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
-    )
-    parser.add_argument(
-        "--lr_num_cycles",
-        type=int,
-        default=1,
-        help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
-    )
-    parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.")
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=0,
-        help=(
-            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
-        ),
-    )
-    parser.add_argument(
-        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
-    )
-    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
-    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
-    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
-    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
-    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
-    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
-    parser.add_argument(
-        "--hub_model_id",
-        type=str,
-        default=None,
-        help="The name of the repository to keep in sync with the local `output_dir`.",
-    )
-    parser.add_argument(
-        "--logging_dir",
-        type=str,
-        default="logs",
-        help=(
-            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
-        ),
-    )
-    parser.add_argument(
-        "--allow_tf32",
-        action="store_true",
-        help=(
-            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
-            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
-        ),
-    )
-    parser.add_argument(
-        "--report_to",
-        type=str,
-        default="tensorboard",
-        help=(
-            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
-            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
-        ),
-    )
-    parser.add_argument(
-        "--mixed_precision",
-        type=str,
-        default=None,
-        choices=["no", "fp16", "bf16"],
-        help=(
-            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
-            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
-            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
-        ),
-    )
-    parser.add_argument(
-        "--prior_generation_precision",
-        type=str,
-        default=None,
-        choices=["no", "fp32", "fp16", "bf16"],
-        help=(
-            "Choose prior generation precision between fp32, fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
-            " 1.10.and an Nvidia Ampere GPU.  Default to  fp16 if a GPU is available else fp32."
-        ),
-    )
-    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
-    parser.add_argument(
-        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
-    )
-
-    if input_args is not None:
-        args = parser.parse_args(input_args)
-    else:
-        args = parser.parse_args()
-
-    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
-    if env_local_rank != -1 and env_local_rank != args.local_rank:
-        args.local_rank = env_local_rank
-
-    if args.with_prior_preservation:
-        if args.class_data_dir is None:
-            raise ValueError("You must specify a data directory for class images.")
-        if args.class_prompt is None:
-            raise ValueError("You must specify prompt for class images.")
-    else:
-        # logger is not available yet
-        if args.class_data_dir is not None:
-            warnings.warn("You need not use --class_data_dir without --with_prior_preservation.")
-        if args.class_prompt is not None:
-            warnings.warn("You need not use --class_prompt without --with_prior_preservation.")
-
-    return args
-
-
-class DreamBoothDataset(Dataset):
-    """
-    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
-    It pre-processes the images and the tokenizes prompts.
-    """
-
-    def __init__(
-        self,
-        instance_data_root,
-        instance_prompt,
-        tokenizer,
-        class_data_root=None,
-        class_prompt=None,
-        size=512,
-        center_crop=False,
-    ):
-        self.size = size
-        self.center_crop = center_crop
-        self.tokenizer = tokenizer
-
-        self.instance_data_root = Path(instance_data_root)
-        if not self.instance_data_root.exists():
-            raise ValueError("Instance images root doesn't exists.")
-
-        self.instance_images_path = list(Path(instance_data_root).iterdir())
-        self.num_instance_images = len(self.instance_images_path)
-        self.instance_prompt = instance_prompt
-        self._length = self.num_instance_images
-
-        if class_data_root is not None:
-            self.class_data_root = Path(class_data_root)
-            self.class_data_root.mkdir(parents=True, exist_ok=True)
-            self.class_images_path = list(self.class_data_root.iterdir())
-            self.num_class_images = len(self.class_images_path)
-            self._length = max(self.num_class_images, self.num_instance_images)
-            self.class_prompt = class_prompt
-        else:
-            self.class_data_root = None
-
-        self.image_transforms = transforms.Compose(
-            [
-                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
-                transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
-                transforms.ToTensor(),
-                transforms.Normalize([0.5], [0.5]),
-            ]
-        )
-
-    def __len__(self):
-        return self._length
-
-    def __getitem__(self, index):
-        example = {}
-        instance_image = Image.open(self.instance_images_path[index % self.num_instance_images])
-        if not instance_image.mode == "RGB":
-            instance_image = instance_image.convert("RGB")
-        example["instance_images"] = self.image_transforms(instance_image)
-        example["instance_prompt_ids"] = self.tokenizer(
-            self.instance_prompt,
-            truncation=True,
-            padding="max_length",
-            max_length=self.tokenizer.model_max_length,
-            return_tensors="pt",
-        ).input_ids
-
-        if self.class_data_root:
-            class_image = Image.open(self.class_images_path[index % self.num_class_images])
-            if not class_image.mode == "RGB":
-                class_image = class_image.convert("RGB")
-            example["class_images"] = self.image_transforms(class_image)
-            example["class_prompt_ids"] = self.tokenizer(
-                self.class_prompt,
-                truncation=True,
-                padding="max_length",
-                max_length=self.tokenizer.model_max_length,
-                return_tensors="pt",
-            ).input_ids
-
-        return example
-
-
-def collate_fn(examples, with_prior_preservation=False):
-    input_ids = [example["instance_prompt_ids"] for example in examples]
-    pixel_values = [example["instance_images"] for example in examples]
-
-    # Concat class and instance examples for prior preservation.
-    # We do this to avoid doing two forward passes.
-    if with_prior_preservation:
-        input_ids += [example["class_prompt_ids"] for example in examples]
-        pixel_values += [example["class_images"] for example in examples]
-
-    pixel_values = torch.stack(pixel_values)
-    pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
-
-    input_ids = torch.cat(input_ids, dim=0)
-
-    batch = {
-        "input_ids": input_ids,
-        "pixel_values": pixel_values,
-    }
-    return batch
-
-
-class PromptDataset(Dataset):
-    "A simple dataset to prepare the prompts to generate class images on multiple GPUs."
-
-    def __init__(self, prompt, num_samples):
-        self.prompt = prompt
-        self.num_samples = num_samples
-
-    def __len__(self):
-        return self.num_samples
-
-    def __getitem__(self, index):
-        example = {}
-        example["prompt"] = self.prompt
-        example["index"] = index
-        return example
-
-
-def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
-    if token is None:
-        token = HfFolder.get_token()
-    if organization is None:
-        username = whoami(token)["name"]
-        return f"{username}/{model_id}"
-    else:
-        return f"{organization}/{model_id}"
-
-
-def main(args):
-    logging_dir = Path(args.output_dir, args.logging_dir)
-
-    accelerator = Accelerator(
-        gradient_accumulation_steps=args.gradient_accumulation_steps,
-        mixed_precision=args.mixed_precision,
-        log_with=args.report_to,
-        logging_dir=logging_dir,
-    )
-
-    if args.report_to == "wandb":
-        if not is_wandb_available():
-            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
-        import wandb
-
-    # Currently, it's not possible to do gradient accumulation when training two models with accelerate.accumulate
-    # This will be enabled soon in accelerate. For now, we don't allow gradient accumulation when training two models.
-    # TODO (patil-suraj): Remove this check when gradient accumulation with two models is enabled in accelerate.
-    # Make one log on every process with the configuration for debugging.
-    logging.basicConfig(
-        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
-        datefmt="%m/%d/%Y %H:%M:%S",
-        level=logging.INFO,
-    )
-    logger.info(accelerator.state, main_process_only=False)
-    if accelerator.is_local_main_process:
-        datasets.utils.logging.set_verbosity_warning()
-        transformers.utils.logging.set_verbosity_warning()
-        diffusers.utils.logging.set_verbosity_info()
-    else:
-        datasets.utils.logging.set_verbosity_error()
-        transformers.utils.logging.set_verbosity_error()
-        diffusers.utils.logging.set_verbosity_error()
-
-    # If passed along, set the training seed now.
-    if args.seed is not None:
-        set_seed(args.seed)
-
-    # Generate class images if prior preservation is enabled.
-    if args.with_prior_preservation:
-        class_images_dir = Path(args.class_data_dir)
-        if not class_images_dir.exists():
-            class_images_dir.mkdir(parents=True)
-        cur_class_images = len(list(class_images_dir.iterdir()))
-
-        if cur_class_images < args.num_class_images:
-            torch_dtype = torch.float16 if accelerator.device.type == "cuda" else torch.float32
-            if args.prior_generation_precision == "fp32":
-                torch_dtype = torch.float32
-            elif args.prior_generation_precision == "fp16":
-                torch_dtype = torch.float16
-            elif args.prior_generation_precision == "bf16":
-                torch_dtype = torch.bfloat16
-            pipeline = DiffusionPipeline.from_pretrained(
-                args.pretrained_model_name_or_path,
-                torch_dtype=torch_dtype,
-                safety_checker=None,
-                revision=args.revision,
-            )
-            pipeline.set_progress_bar_config(disable=True)
-
-            num_new_images = args.num_class_images - cur_class_images
-            logger.info(f"Number of class images to sample: {num_new_images}.")
-
-            sample_dataset = PromptDataset(args.class_prompt, num_new_images)
-            sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)
-
-            sample_dataloader = accelerator.prepare(sample_dataloader)
-            pipeline.to(accelerator.device)
-
-            for example in tqdm(
-                sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
-            ):
-                images = pipeline(example["prompt"]).images
-
-                for i, image in enumerate(images):
-                    hash_image = hashlib.sha1(image.tobytes()).hexdigest()
-                    image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg"
-                    image.save(image_filename)
-
-            del pipeline
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-    # Handle the repository creation
-    if accelerator.is_main_process:
-        if args.push_to_hub:
-            if args.hub_model_id is None:
-                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
-            else:
-                repo_name = args.hub_model_id
-
-            create_repo(repo_name, exist_ok=True, token=args.hub_token)
-            repo = Repository(args.output_dir, clone_from=repo_name, token=args.hub_token)
-
-            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
-                if "step_*" not in gitignore:
-                    gitignore.write("step_*\n")
-                if "epoch_*" not in gitignore:
-                    gitignore.write("epoch_*\n")
-        elif args.output_dir is not None:
-            os.makedirs(args.output_dir, exist_ok=True)
-
-    # Load the tokenizer
-    if args.tokenizer_name:
-        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
-    elif args.pretrained_model_name_or_path:
-        tokenizer = AutoTokenizer.from_pretrained(
-            args.pretrained_model_name_or_path,
-            subfolder="tokenizer",
-            revision=args.revision,
-            use_fast=False,
-        )
-
-    # import correct text encoder class
-    text_encoder_cls = import_model_class_from_model_name_or_path(args.pretrained_model_name_or_path, args.revision)
-
-    # Load scheduler and models
-    noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
-    text_encoder = text_encoder_cls.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
-    )
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
-    unet = UNet2DConditionModel.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
-    )
-
-    # We only train the additional adapter LoRA layers
-    vae.requires_grad_(False)
-    text_encoder.requires_grad_(False)
-    unet.requires_grad_(False)
-
-    # For mixed precision training we cast the text_encoder and vae weights to half-precision
-    # as these models are only used for inference, keeping weights in full precision is not required.
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-
-    # Move unet, vae and text_encoder to device and cast to weight_dtype
-    unet.to(accelerator.device, dtype=weight_dtype)
-    vae.to(accelerator.device, dtype=weight_dtype)
-    text_encoder.to(accelerator.device, dtype=weight_dtype)
-
-    if args.enable_xformers_memory_efficient_attention:
-        if is_xformers_available():
-            unet.enable_xformers_memory_efficient_attention()
-        else:
-            raise ValueError("xformers is not available. Make sure it is installed correctly")
-
-    # now we will add new LoRA weights to the attention layers
-    # It's important to realize here how many attention weights will be added and of which sizes
-    # The sizes of the attention layers consist only of two different variables:
-    # 1) - the "hidden_size", which is increased according to `unet.config.block_out_channels`.
-    # 2) - the "cross attention size", which is set to `unet.config.cross_attention_dim`.
-
-    # Let's first see how many attention processors we will have to set.
-    # For Stable Diffusion, it should be equal to:
-    # - down blocks (2x attention layers) * (2x transformer layers) * (3x down blocks) = 12
-    # - mid blocks (2x attention layers) * (1x transformer layers) * (1x mid blocks) = 2
-    # - up blocks (2x attention layers) * (3x transformer layers) * (3x down blocks) = 18
-    # => 32 layers
-
-    # Set correct lora layers
-    lora_attn_procs = {}
-    for name in unet.attn_processors.keys():
-        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
-        if name.startswith("mid_block"):
-            hidden_size = unet.config.block_out_channels[-1]
-        elif name.startswith("up_blocks"):
-            block_id = int(name[len("up_blocks.")])
-            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
-        elif name.startswith("down_blocks"):
-            block_id = int(name[len("down_blocks.")])
-            hidden_size = unet.config.block_out_channels[block_id]
-
-        lora_attn_procs[name] = LoRACrossAttnProcessor(
-            hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
-        )
-
-    unet.set_attn_processor(lora_attn_procs)
-    lora_layers = AttnProcsLayers(unet.attn_processors)
-
-    accelerator.register_for_checkpointing(lora_layers)
-
-    if args.scale_lr:
-        args.learning_rate = (
-            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
-        )
-
-    # Enable TF32 for faster training on Ampere GPUs,
-    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
-    if args.allow_tf32:
-        torch.backends.cuda.matmul.allow_tf32 = True
-
-    if args.scale_lr:
-        args.learning_rate = (
-            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
-        )
-
-    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
-    if args.use_8bit_adam:
-        try:
-            import bitsandbytes as bnb
-        except ImportError:
-            raise ImportError(
-                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
-            )
-
-        optimizer_class = bnb.optim.AdamW8bit
-    else:
-        optimizer_class = torch.optim.AdamW
-
-    # Optimizer creation
-    optimizer = optimizer_class(
-        lora_layers.parameters(),
-        lr=args.learning_rate,
-        betas=(args.adam_beta1, args.adam_beta2),
-        weight_decay=args.adam_weight_decay,
-        eps=args.adam_epsilon,
-    )
-
-    # Dataset and DataLoaders creation:
-    train_dataset = DreamBoothDataset(
-        instance_data_root=args.instance_data_dir,
-        instance_prompt=args.instance_prompt,
-        class_data_root=args.class_data_dir if args.with_prior_preservation else None,
-        class_prompt=args.class_prompt,
-        tokenizer=tokenizer,
-        size=args.resolution,
-        center_crop=args.center_crop,
-    )
-
-    train_dataloader = torch.utils.data.DataLoader(
-        train_dataset,
-        batch_size=args.train_batch_size,
-        shuffle=True,
-        collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
-        num_workers=args.dataloader_num_workers,
-    )
-
-    # Scheduler and math around the number of training steps.
-    overrode_max_train_steps = False
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
-    if args.max_train_steps is None:
-        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
-        overrode_max_train_steps = True
-
-    lr_scheduler = get_scheduler(
-        args.lr_scheduler,
-        optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
-        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-        num_cycles=args.lr_num_cycles,
-        power=args.lr_power,
-    )
-
-    # Prepare everything with our `accelerator`.
-    lora_layers, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        lora_layers, optimizer, train_dataloader, lr_scheduler
-    )
-
-    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
-    if overrode_max_train_steps:
-        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
-    # Afterwards we recalculate our number of training epochs
-    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
-
-    # We need to initialize the trackers we use, and also store our configuration.
-    # The trackers initializes automatically on the main process.
-    if accelerator.is_main_process:
-        accelerator.init_trackers("dreambooth-lora", config=vars(args))
-
-    # Train!
-    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-
-    logger.info("***** Running training *****")
-    logger.info(f"  Num examples = {len(train_dataset)}")
-    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
-    logger.info(f"  Num Epochs = {args.num_train_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
-    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
-    logger.info(f"  Total optimization steps = {args.max_train_steps}")
-    global_step = 0
-    first_epoch = 0
-
-    # Potentially load in the weights and states from a previous save
-    if args.resume_from_checkpoint:
-        if args.resume_from_checkpoint != "latest":
-            path = os.path.basename(args.resume_from_checkpoint)
-        else:
-            # Get the mos recent checkpoint
-            dirs = os.listdir(args.output_dir)
-            dirs = [d for d in dirs if d.startswith("checkpoint")]
-            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
-            path = dirs[-1] if len(dirs) > 0 else None
-
-        if path is None:
-            accelerator.print(
-                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
-            )
-            args.resume_from_checkpoint = None
-        else:
-            accelerator.print(f"Resuming from checkpoint {path}")
-            accelerator.load_state(os.path.join(args.output_dir, path))
-            global_step = int(path.split("-")[1])
-
-            resume_global_step = global_step * args.gradient_accumulation_steps
-            first_epoch = global_step // num_update_steps_per_epoch
-            resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)
-
-    # Only show the progress bar once on each machine.
-    progress_bar = tqdm(range(global_step, args.max_train_steps), disable=not accelerator.is_local_main_process)
-    progress_bar.set_description("Steps")
-
-    for epoch in range(first_epoch, args.num_train_epochs):
-        unet.train()
-        for step, batch in enumerate(train_dataloader):
-            # Skip steps until we reach the resumed step
-            if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
-                if step % args.gradient_accumulation_steps == 0:
-                    progress_bar.update(1)
-                continue
-
-            with accelerator.accumulate(unet):
-                # Convert images to latent space
-                latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
-                latents = latents * 0.18215
-
-                # Sample noise that we'll add to the latents
-                noise = torch.randn_like(latents)
-                bsz = latents.shape[0]
-                # Sample a random timestep for each image
-                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
-                timesteps = timesteps.long()
-
-                # Add noise to the latents according to the noise magnitude at each timestep
-                # (this is the forward diffusion process)
-                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
-
-                # Get the text embedding for conditioning
-                encoder_hidden_states = text_encoder(batch["input_ids"])[0]
-
-                # Predict the noise residual
-                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
-
-                # Get the target for loss depending on the prediction type
-                if noise_scheduler.config.prediction_type == "epsilon":
-                    target = noise
-                elif noise_scheduler.config.prediction_type == "v_prediction":
-                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
-                else:
-                    raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
-
-                if args.with_prior_preservation:
-                    # Chunk the noise and model_pred into two parts and compute the loss on each part separately.
-                    model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
-                    target, target_prior = torch.chunk(target, 2, dim=0)
-
-                    # Compute instance loss
-                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
-
-                    # Compute prior loss
-                    prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="mean")
-
-                    # Add the prior loss to the instance loss.
-                    loss = loss + args.prior_loss_weight * prior_loss
-                else:
-                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
-
-                accelerator.backward(loss)
-                if accelerator.sync_gradients:
-                    params_to_clip = lora_layers.parameters()
-                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
-                optimizer.step()
-                lr_scheduler.step()
-                optimizer.zero_grad()
-
-            # Checks if the accelerator has performed an optimization step behind the scenes
-            if accelerator.sync_gradients:
-                progress_bar.update(1)
-                global_step += 1
-
-                if global_step % args.checkpointing_steps == 0:
-                    if accelerator.is_main_process:
-                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
-                        accelerator.save_state(save_path)
-                        logger.info(f"Saved state to {save_path}")
-
-            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
-            progress_bar.set_postfix(**logs)
-            accelerator.log(logs, step=global_step)
-
-            if global_step >= args.max_train_steps:
-                break
-
-        if args.validation_prompt is not None and epoch % args.validation_epochs == 0:
-            logger.info(
-                f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
-                f" {args.validation_prompt}."
-            )
-            # create pipeline
-            pipeline = DiffusionPipeline.from_pretrained(
-                args.pretrained_model_name_or_path,
-                unet=accelerator.unwrap_model(unet),
-                text_encoder=accelerator.unwrap_model(text_encoder),
-                revision=args.revision,
-                torch_dtype=weight_dtype,
-            )
-            pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
-            pipeline = pipeline.to(accelerator.device)
-            pipeline.set_progress_bar_config(disable=True)
-
-            # run inference
-            generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
-            prompt = args.num_validation_images * [args.validation_prompt]
-            images = pipeline(prompt, num_inference_steps=25, generator=generator).images
-
-            for tracker in accelerator.trackers:
-                if tracker.name == "tensorboard":
-                    np_images = np.stack([np.asarray(img) for img in images])
-                    tracker.writer.add_images("validation", np_images, epoch, dataformats="NHWC")
-                if tracker.name == "wandb":
-                    tracker.log(
-                        {
-                            "validation": [
-                                wandb.Image(image, caption=f"{i}: {args.validation_prompt}")
-                                for i, image in enumerate(images)
-                            ]
-                        }
-                    )
-
-            del pipeline
-            torch.cuda.empty_cache()
-
-    # Save the lora layers
-    accelerator.wait_for_everyone()
-    if accelerator.is_main_process:
-        unet = unet.to(torch.float32)
-        unet.save_attn_procs(args.output_dir)
-
-        # Final inference
-        # Load previous pipeline
-        pipeline = DiffusionPipeline.from_pretrained(
-            args.pretrained_model_name_or_path, revision=args.revision, torch_dtype=weight_dtype
-        )
-        pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
-        pipeline = pipeline.to(accelerator.device)
-
-        # load attention processors
-        pipeline.unet.load_attn_procs(args.output_dir)
-
-        # run inference
-        if args.validation_prompt and args.num_validation_images > 0:
-            generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
-            prompt = args.num_validation_images * [args.validation_prompt]
-            images = pipeline(prompt, num_inference_steps=25, generator=generator).images
-
-            test_image_dir = Path(args.output_dir) / 'test_images'
-            test_image_dir.mkdir()
-            for i, image in enumerate(images):
-                out_path = test_image_dir / f'image_{i}.png'
-                image.save(out_path)
-
-            for tracker in accelerator.trackers:
-                if tracker.name == "tensorboard":
-                    np_images = np.stack([np.asarray(img) for img in images])
-                    tracker.writer.add_images("test", np_images, epoch, dataformats="NHWC")
-                if tracker.name == "wandb":
-                    tracker.log(
-                        {
-                            "test": [
-                                wandb.Image(image, caption=f"{i}: {args.validation_prompt}")
-                                for i, image in enumerate(images)
-                            ]
-                        }
-                    )
-
-        if args.push_to_hub:
-            save_model_card(
-                repo_name,
-                images=images,
-                base_model=args.pretrained_model_name_or_path,
-                prompt=args.instance_prompt,
-                repo_folder=args.output_dir,
-            )
-            repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True)
-
-    accelerator.end_training()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

trainer.py

@@ -6,61 +6,52 @@ import pathlib
 import shlex
 import shutil
 import subprocess
+import sys
 
 import gradio as gr
-import PIL.Image
 import slugify
 import torch
 from huggingface_hub import HfApi
+from omegaconf import OmegaConf
 
-from app_upload import LoRAModelUploader
+from app_upload import ModelUploader
 from utils import save_model_card
 
-URL_TO_JOIN_LORA_LIBRARY_ORG = 'https://huggingface.co/organizations/lora-library/share/hjetHAcKjnPHXhHfbeEcqnBqmhgilFfpOL'
+sys.path.append('Tune-A-Video')
 
-
-def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
-    w, h = image.size
-    if w == h:
-        return image
-    elif w > h:
-        new_image = PIL.Image.new(image.mode, (w, w), (0, 0, 0))
-        new_image.paste(image, (0, (w - h) // 2))
-        return new_image
-    else:
-        new_image = PIL.Image.new(image.mode, (h, h), (0, 0, 0))
-        new_image.paste(image, ((h - w) // 2, 0))
-        return new_image
+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.api = HfApi(token=hf_token)
-        self.model_uploader = LoRAModelUploader(hf_token)
-
-    def prepare_dataset(self, instance_images: list, resolution: int,
-                        instance_data_dir: pathlib.Path) -> None:
-        shutil.rmtree(instance_data_dir, ignore_errors=True)
-        instance_data_dir.mkdir(parents=True)
-        for i, temp_path in enumerate(instance_images):
-            image = PIL.Image.open(temp_path.name)
-            image = pad_image(image)
-            image = image.resize((resolution, resolution))
-            image = image.convert('RGB')
-            out_path = instance_data_dir / f'{i:03d}.jpg'
-            image.save(out_path, format='JPEG', quality=100)
-
-    def join_lora_library_org(self) -> None:
+        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) -> None:
         subprocess.run(
             shlex.split(
-                f'curl -X POST -H "Authorization: Bearer {self.hf_token}" -H "Content-Type: application/json" {URL_TO_JOIN_LORA_LIBRARY_ORG}'
+                f'curl -X POST -H "Authorization: Bearer {self.hf_token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}'
             ))
 
     def run(
         self,
-        instance_images: list | None,
-        instance_prompt: str,
+        training_video: str,
+        training_prompt: str,
         output_model_name: str,
         overwrite_existing_model: bool,
         validation_prompt: str,
@@ -73,7 +64,6 @@ class Trainer:
         fp16: bool,
         use_8bit_adam: bool,
         checkpointing_steps: int,
-        use_wandb: bool,
         validation_epochs: int,
         upload_to_hub: bool,
         use_private_repo: bool,
@@ -83,10 +73,10 @@ class Trainer:
     ) -> str:
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
-        if instance_images is None:
-            raise gr.Error('You need to upload images.')
-        if not instance_prompt:
-            raise gr.Error('The instance prompt is missing.')
+        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.')
 
@@ -94,7 +84,7 @@ class Trainer:
 
         if not output_model_name:
             timestamp = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
-            output_model_name = f'lora-dreambooth-{timestamp}'
+            output_model_name = f'tune-a-video-{timestamp}'
         output_model_name = slugify.slugify(output_model_name)
 
         repo_dir = pathlib.Path(__file__).parent
@@ -103,52 +93,52 @@ class Trainer:
             shutil.rmtree(output_dir, ignore_errors=True)
         output_dir.mkdir(parents=True)
 
-        instance_data_dir = repo_dir / 'training_data' / output_model_name
-        self.prepare_dataset(instance_images, resolution, instance_data_dir)
-
         if upload_to_hub:
-            self.join_lora_library_org()
-
-        command = f'''
-        accelerate launch train_dreambooth_lora.py \
-          --pretrained_model_name_or_path={base_model}  \
-          --instance_data_dir={instance_data_dir} \
-          --output_dir={output_dir} \
-          --instance_prompt="{instance_prompt}" \
-          --resolution={resolution} \
-          --train_batch_size=1 \
-          --gradient_accumulation_steps={gradient_accumulation} \
-          --learning_rate={learning_rate} \
-          --lr_scheduler=constant \
-          --lr_warmup_steps=0 \
-          --max_train_steps={n_steps} \
-          --checkpointing_steps={checkpointing_steps} \
-          --validation_prompt="{validation_prompt}" \
-          --validation_epochs={validation_epochs} \
-          --seed={seed}
-        '''
-        if fp16:
-            command += ' --mixed_precision fp16'
-        if use_8bit_adam:
-            command += ' --use_8bit_adam'
-        if use_wandb:
-            command += ' --report_to wandb'
-
-        with open(output_dir / 'train.sh', 'w') as f:
-            command_s = ' '.join(command.split())
-            f.write(command_s)
+            self.join_model_library_org()
+
+        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,
-                        instance_prompt=instance_prompt,
+                        training_prompt=training_prompt,
                         test_prompt=validation_prompt,
-                        test_image_dir='test_images')
+                        test_image_dir='samples')
 
         message = 'Training completed!'
         print(message)
 
         if upload_to_hub:
-            upload_message = self.model_uploader.upload_lora_model(
+            upload_message = self.model_uploader.upload_model(
                 folder_path=output_dir.as_posix(),
                 repo_name=output_model_name,
                 upload_to=upload_to,

utils.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import pathlib
 
 
-def find_exp_dirs(ignore_repo: bool = False) -> list[str]:
+def find_exp_dirs() -> list[str]:
     repo_dir = pathlib.Path(__file__).parent
     exp_root_dir = repo_dir / 'experiments'
     if not exp_root_dir.exists():
@@ -11,46 +11,45 @@ def find_exp_dirs(ignore_repo: bool = False) -> list[str]:
     exp_dirs = sorted(exp_root_dir.glob('*'))
     exp_dirs = [
         exp_dir for exp_dir in exp_dirs
-        if (exp_dir / 'pytorch_lora_weights.bin').exists()
+        if (exp_dir / 'model_index.json').exists()
     ]
-    if ignore_repo:
-        exp_dirs = [
-            exp_dir for exp_dir in exp_dirs if not (exp_dir / '.git').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,
-    instance_prompt: 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('*'))
+        image_paths = sorted((save_dir / test_image_dir).glob('*.gif'))
         if image_paths:
-            image_str = f'Test prompt: {test_prompt}\n'
-            for image_path in image_paths:
-                rel_path = image_path.relative_to(save_dir)
-                image_str += f'![{image_path.stem}]({rel_path})\n'
+            image_path = image_paths[-1]
+            rel_path = image_path.relative_to(save_dir)
+            image_str = f'Test prompt: {test_prompt}\n' + f'![{image_path.stem}]({rel_path})\n'
 
     model_card = f'''---
 license: creativeml-openrail-m
 base_model: {base_model}
-instance_prompt: {instance_prompt}
+training_prompt: {training_prompt}
 tags:
 - stable-diffusion
 - stable-diffusion-diffusers
 - text-to-image
 - diffusers
-- lora
-inference: true
+- text-to-video
+- tune-a-video
+inference: false
 ---
-# LoRA DreamBooth - {save_dir.name}
 
-These are LoRA adaption weights for [{base_model}](https://huggingface.co/{base_model}). The weights were trained on the instance prompt "{instance_prompt}" using [DreamBooth](https://dreambooth.github.io/). You can find some example images in the following.
+# Tune-A-Video - {save_dir.name}
+
+Base model: [{base_model}](https://huggingface.co/{base_model}).
+
+Training prompt: {training_prompt}
 
 {image_str}
 '''

wheel/xformers-0.0.16+bc08bbc.d20230130-cp310-cp310-linux_x86_64.whl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:016219e017ce06b351ef0f98fc074ee60be06ee1d700cfe0a45c9b59e25bb938
+size 134437916