diff --git a/configs/man-skiing.yaml b/configs/man-skiing.yaml
index 3089b11..a2ff273 100644
--- a/configs/man-skiing.yaml
+++ b/configs/man-skiing.yaml
@@ -14,7 +14,7 @@ validation_data:
prompts:
- "mickey mouse is skiing on the snow"
- "spider man is skiing on the beach, cartoon style"
- - "wonder woman, wearing a cowboy hat, is skiing”"
+ - "wonder woman, wearing a cowboy hat, is skiing"
- "a man, wearing pink clothes, is skiing at sunset"
video_length: 24
width: 512
diff --git a/notebooks/Tune-A-Video.ipynb b/notebooks/Tune-A-Video.ipynb
index 276fd05..ee77a84 100644
--- a/notebooks/Tune-A-Video.ipynb
+++ b/notebooks/Tune-A-Video.ipynb
@@ -1 +1,385 @@
-{"cells":[{"cell_type":"markdown","metadata":{"id":"fZ_xQvU70UQc"},"source":["# Tune-A-Video\n","\n","**[Tune-A-Video: One-Shot Tuning of Image Diffusion Models for Text-to-Video Generation](https://arxiv.org/abs/2212.11565)** \n","[Jay Zhangjie Wu](https://zhangjiewu.github.io/), \n","[Yixiao Ge](https://geyixiao.com/), \n","[Xintao Wang](https://xinntao.github.io/), \n","[Stan Weixian Lei](), \n","[Yuchao Gu](https://ycgu.site/), \n","[Wynne Hsu](https://www.comp.nus.edu.sg/~whsu/), \n","[Ying Shan](https://scholar.google.com/citations?user=4oXBp9UAAAAJ&hl=en), \n","[Xiaohu Qie](https://scholar.google.com/citations?user=mk-F69UAAAAJ&hl=en), \n","[Mike Zheng Shou](https://sites.google.com/view/showlab) \n","\n","[](https://tuneavideo.github.io/)\n","[](https://arxiv.org/abs/2212.11565)\n","[](https://huggingface.co/spaces/Tune-A-Video-library/Tune-A-Video-Training-UI)\n","[](https://github.com/showlab/Tune-A-Video)"]},{"cell_type":"markdown","metadata":{"id":"wnTMyW41cC1E"},"source":["## Setup"]},{"cell_type":"code","execution_count":1,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":1617,"status":"ok","timestamp":1675847464041,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"XU7NuMAA2drw","outputId":"6652be79-2978-490c-80f1-a5a0afe8a165"},"outputs":[{"name":"stdout","output_type":"stream","text":["Tesla T4, 15360 MiB, 15109 MiB\n"]}],"source":["#@markdown Check type of GPU and VRAM available.\n","!nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader"]},{"cell_type":"code","execution_count":2,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":48415,"status":"ok","timestamp":1675847514322,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"D1PRgre3Gt5U","outputId":"85079eb4-f9f9-4d4c-bbb5-c1fcc829c0ce"},"outputs":[{"name":"stdout","output_type":"stream","text":["Cloning into '/content/Tune-A-Video'...\n","remote: Enumerating objects: 42, done.\u001b[K\n","remote: Counting objects: 100% (42/42), done.\u001b[K\n","remote: Compressing objects: 100% (32/32), done.\u001b[K\n","remote: Total 42 (delta 9), reused 32 (delta 6), pack-reused 0\u001b[K\n","Unpacking objects: 100% (42/42), 796.30 KiB | 11.06 MiB/s, done.\n","/content/Tune-A-Video\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 MB\u001b[0m \u001b[31m11.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m145.0/145.0 KB\u001b[0m \u001b[31m14.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25h Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"," Building wheel for lit (setup.py) ... \u001b[?25l\u001b[?25hdone\n"]}],"source":["#@title Install requirements\n","\n","!git clone https://github.com/showlab/Tune-A-Video.git /content/Tune-A-Video\n","%cd /content/Tune-A-Video \n","# %pip install -r requirements.txt\n","%pip install -q -U --pre triton\n","%pip install -q diffusers[torch]==0.11.1 transformers==4.26.0 bitsandbytes==0.35.4 \\\n","decord accelerate omegaconf einops ftfy gradio imageio-ffmpeg xformers"]},{"cell_type":"code","execution_count":3,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":82042,"status":"ok","timestamp":1675847618104,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"m6I6kZNG3Inb","outputId":"378860d5-2fa0-4aa5-d5fb-0e960254a93a"},"outputs":[{"name":"stdout","output_type":"stream","text":["Updated git hooks.\n","Git LFS initialized.\n","Cloning into 'checkpoints/CompVis/stable-diffusion-v1-4'...\n","remote: Enumerating objects: 719, done.\u001b[K\n","remote: Total 719 (delta 0), reused 0 (delta 0), pack-reused 719\u001b[K\n","Receiving objects: 100% (719/719), 679.70 KiB | 1.42 MiB/s, done.\n","Resolving deltas: 100% (119/119), done.\n","Filtering content: 100% (4/4), 5.10 GiB | 66.30 MiB/s, done.\n","[*] MODEL_NAME=./checkpoints/CompVis/stable-diffusion-v1-4\n"]}],"source":["#@title Download pretrained model\n","\n","#@markdown Name/Path of the initial model.\n","MODEL_NAME = \"CompVis/stable-diffusion-v1-4\" #@param {type:\"string\"}\n","\n","#@markdown If model should be download from a remote repo. Untick it if the model is loaded from a local path.\n","download_pretrained_model = True #@param {type:\"boolean\"}\n","if download_pretrained_model:\n"," !git lfs install\n"," !git clone https://huggingface.co/$MODEL_NAME checkpoints/$MODEL_NAME\n"," MODEL_NAME = f\"./checkpoints/{MODEL_NAME}\"\n","print(f\"[*] MODEL_NAME={MODEL_NAME}\")"]},{"cell_type":"markdown","metadata":{"id":"qn5ILIyDJIcX"},"source":["## Usage\n"]},{"cell_type":"markdown","metadata":{"id":"REmFAHfz9Y_X"},"source":["### Training\n"]},{"cell_type":"code","execution_count":4,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":556,"status":"ok","timestamp":1675848002295,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"Rxg0y5MBudmd","outputId":"06cebc9e-27fe-4073-f370-0900d555a628"},"outputs":[{"name":"stdout","output_type":"stream","text":["[*] Weights will be saved at outputs/man-surfing\n"]}],"source":["#@markdown If model weights should be saved directly in google drive (takes around 4-5 GB).\n","save_to_gdrive = False #@param {type:\"boolean\"}\n","if save_to_gdrive:\n"," from google.colab import drive\n"," drive.mount('/content/drive')\n","\n","#@markdown Enter the directory name to save model at.\n","\n","OUTPUT_DIR = \"outputs/man-surfing\" #@param {type:\"string\"}\n","if save_to_gdrive:\n"," OUTPUT_DIR = \"/content/drive/MyDrive/\" + OUTPUT_DIR\n","\n","print(f\"[*] Weights will be saved at {OUTPUT_DIR}\")\n","\n","!mkdir -p $OUTPUT_DIR\n"]},{"cell_type":"code","execution_count":5,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/","height":39},"executionInfo":{"elapsed":3955,"status":"ok","timestamp":1675848008563,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"32gYIDDR1aCp","outputId":"74be2d00-844e-4a8b-8773-1aecd9cb427f"},"outputs":[{"data":{"text/html":["\n"," \n"," \n"," "],"text/plain":[""]},"metadata":{},"output_type":"display_data"}],"source":["#@markdown Upload your video by running this cell.\n","\n","#@markdown OR\n","\n","#@markdown You can use the file manager on the left panel to upload (drag and drop) to `data` folder.\n","\n","import os\n","from google.colab import files\n","import shutil\n","\n","uploaded = files.upload()\n","for filename in uploaded.keys():\n"," dst_path = os.path.join(\"data\", filename)\n"," shutil.move(filename, dst_path)"]},{"cell_type":"code","execution_count":6,"metadata":{"executionInfo":{"elapsed":1,"status":"ok","timestamp":1675848018997,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"wGGFFpNcR2d_"},"outputs":[],"source":["#@markdown Train config\n","\n","from omegaconf import OmegaConf\n","\n","CONFIG_NAME = \"configs/man-surfing.yaml\" #@param {type:\"string\"}\n","\n","train_video_path = \"data/man-surfing.mp4\" #@param {type:\"string\"}\n","train_prompt = \"a man is surfing\" #@param {type:\"string\"}\n","video_length = 8 #@param {type:\"number\"}\n","width = 512 #@param {type:\"number\"}\n","height = 512 #@param {type:\"number\"}\n","learning_rate = 3e-5 #@param {type:\"number\"}\n","train_steps = 300 #@param {type:\"number\"}\n","\n","config = {\n"," \"pretrained_model_path\": MODEL_NAME,\n"," \"output_dir\": OUTPUT_DIR,\n"," \"train_data\": {\n"," \"video_path\": train_video_path,\n"," \"prompt\": train_prompt,\n"," \"n_sample_frames\": video_length,\n"," \"width\": width,\n"," \"height\": height,\n"," \"sample_start_idx\": 0,\n"," \"sample_frame_rate\": 1,\n"," },\n"," \"validation_data\": {\n"," \"prompts\": [\n"," \"a panda is surfing\",\n"," \"a boy, wearing a birthday hat, is surfing\",\n"," \"a raccoon is surfing, cartoon style\",\n"," \"Iron Man is surfing in the desert\",\n"," ],\n"," \"video_length\": video_length,\n"," \"width\": width,\n"," \"height\": height,\n"," \"num_inference_steps\": 20,\n"," \"guidance_scale\": 7.5,\n"," },\n"," \"learning_rate\": learning_rate,\n"," \"train_batch_size\": 1,\n"," \"max_train_steps\": train_steps,\n"," \"checkpointing_steps\": 1000,\n"," \"validation_steps\": 1000,\n"," \"trainable_modules\": [\n"," \"attn1.to_q\",\n"," \"attn2.to_q\",\n"," \"attn_temp\",\n"," ],\n"," \"seed\": 33,\n"," \"mixed_precision\": \"fp16\",\n"," \"use_8bit_adam\": False,\n"," \"gradient_checkpointing\": True,\n"," \"enable_xformers_memory_efficient_attention\": True,\n","}\n","\n","OmegaConf.save(config, CONFIG_NAME)"]},{"cell_type":"code","execution_count":7,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":1109123,"status":"ok","timestamp":1675849135846,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"jjcSXTp-u-Eg","outputId":"dbe96b90-9d60-407e-a277-7ccaaa218a5a"},"outputs":[{"name":"stdout","output_type":"stream","text":["WARNING:accelerate.commands.launch:The following values were not passed to `accelerate launch` and had defaults used instead:\n","\t`--num_processes` was set to a value of `1`\n","\t`--num_machines` was set to a value of `1`\n","\t`--mixed_precision` was set to a value of `'no'`\n","\t`--dynamo_backend` was set to a value of `'no'`\n","To avoid this warning pass in values for each of the problematic parameters or run `accelerate config`.\n","INFO:__main__:Distributed environment: NO\n","Num processes: 1\n","Process index: 0\n","Local process index: 0\n","Device: cuda\n","Mixed precision type: fp16\n","\n","{'prediction_type', 'variance_type'} was not found in config. Values will be initialized to default values.\n","{'norm_num_groups'} was not found in config. Values will be initialized to default values.\n","{'use_linear_projection', 'upcast_attention', 'class_embed_type', 'mid_block_type', 'resnet_time_scale_shift', 'dual_cross_attention', 'only_cross_attention', 'num_class_embeds'} was not found in config. Values will be initialized to default values.\n","{'prediction_type'} was not found in config. Values will be initialized to default values.\n","INFO:__main__:***** Running training *****\n","INFO:__main__: Num examples = 1\n","INFO:__main__: Num Epochs = 300\n","INFO:__main__: Instantaneous batch size per device = 1\n","INFO:__main__: Total train batch size (w. parallel, distributed & accumulation) = 1\n","INFO:__main__: Gradient Accumulation steps = 1\n","INFO:__main__: Total optimization steps = 300\n","Steps: 100% 300/300 [17:06<00:00, 3.43s/it, lr=3e-5, step_loss=0.0428]{'prediction_type'} was not found in config. Values will be initialized to default values.\n","Configuration saved in outputs/man-surfing/2023-02-08T09-20-46/model_index.json\n","Configuration saved in outputs/man-surfing/2023-02-08T09-20-46/vae/config.json\n","Model weights saved in outputs/man-surfing/2023-02-08T09-20-46/vae/diffusion_pytorch_model.bin\n","Configuration saved in outputs/man-surfing/2023-02-08T09-20-46/unet/config.json\n","Model weights saved in outputs/man-surfing/2023-02-08T09-20-46/unet/diffusion_pytorch_model.bin\n","Configuration saved in outputs/man-surfing/2023-02-08T09-20-46/scheduler/scheduler_config.json\n","Steps: 100% 300/300 [17:22<00:00, 3.48s/it, lr=3e-5, step_loss=0.0428]\n"]}],"source":["!accelerate launch train_tuneavideo.py --config=$CONFIG_NAME"]},{"cell_type":"markdown","metadata":{"id":"ToNG4fd_dTbF"},"source":["### Inference"]},{"cell_type":"code","execution_count":13,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":3,"status":"ok","timestamp":1675849329160,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"89Az5NUxOWdy","outputId":"80d6f2ad-f710-4fbf-d115-38667a620b24"},"outputs":[{"name":"stdout","output_type":"stream","text":["[*] WEIGHTS_DIR=outputs/man-surfing/2023-02-08T09-20-46\n"]}],"source":["#@markdown Specify the weights directory to use (leave blank for latest)\n","WEIGHTS_DIR = \"\" #@param {type:\"string\"}\n","if WEIGHTS_DIR == \"\":\n"," from natsort import natsorted\n"," from glob import glob\n"," import os\n"," WEIGHTS_DIR = natsorted(glob(OUTPUT_DIR + os.sep + \"*\"))[-1]\n","print(f\"[*] WEIGHTS_DIR={WEIGHTS_DIR}\")"]},{"cell_type":"code","execution_count":14,"metadata":{"executionInfo":{"elapsed":20083,"status":"ok","timestamp":1675849351518,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"91bsSFv2Punm"},"outputs":[],"source":["import torch\n","from torch import autocast\n","from diffusers import DDIMScheduler\n","from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline\n","from tuneavideo.models.unet import UNet3DConditionModel\n","from tuneavideo.util import save_videos_grid\n","\n","\n","unet = UNet3DConditionModel.from_pretrained(WEIGHTS_DIR, subfolder='unet', torch_dtype=torch.float16).to('cuda')\n","scheduler = DDIMScheduler.from_pretrained(MODEL_NAME, subfolder='scheduler')\n","pipe = TuneAVideoPipeline.from_pretrained(MODEL_NAME, unet=unet, scheduler=scheduler, torch_dtype=torch.float16).to(\"cuda\")\n","pipe.enable_xformers_memory_efficient_attention()\n","g_cuda = None"]},{"cell_type":"code","execution_count":18,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":4,"status":"ok","timestamp":1675849749322,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"oIzkltjpVO_f","outputId":"44604d9a-db00-48e4-8ea2-6c9378532e2e"},"outputs":[{"data":{"text/plain":[""]},"execution_count":18,"metadata":{},"output_type":"execute_result"}],"source":["#@markdown Can set random seed here for reproducibility.\n","g_cuda = torch.Generator(device='cuda')\n","seed = 1234 #@param {type:\"number\"}\n","g_cuda.manual_seed(seed)"]},{"cell_type":"code","execution_count":16,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/","height":49,"referenced_widgets":["49097988b4a4457baeae21183c14bdd3","a49725f4344a49e0a1c974ff795c76b4","fb35a24428564299a99d51004cccdfa0","dfef05b6c77e480e8a3184169190d936","a08b256dc4f244e98bd0a23a8eacb34f","bb8aa031ddb24d238f6f747c2bffc32d","3ae8a0b6018c485ebd33af4bd976b4ed","4439a02ffe694feb9c2b3caf27d3cffb","b8b58e20668142038672a28ec3f429e3","1a24a695cf60498baded5a9194c8b521","5fa928aea7804dac903c454c6fc1e668"]},"executionInfo":{"elapsed":81227,"status":"ok","timestamp":1675849439987,"user":{"displayName":"Tune-A-Video Show Lab","userId":"17565151047989293791"},"user_tz":-480},"id":"K6xoHWSsbcS3","outputId":"94e7ff06-2997-45a6-d948-c7ddeacdb50d","scrolled":false},"outputs":[{"data":{"application/vnd.jupyter.widget-view+json":{"model_id":"49097988b4a4457baeae21183c14bdd3","version_major":2,"version_minor":0},"text/plain":[" 0%| | 0/50 [00:00 {\n"," if (!google.colab.kernel.accessAllowed && !cache) {\n"," return;\n"," }\n"," element.appendChild(document.createTextNode(''));\n"," const url = await google.colab.kernel.proxyPort(port, {cache});\n","\n"," const external_link = document.createElement('div');\n"," external_link.innerHTML = `\n","